def test(opt):
if torch.cuda.is_available():
if opt.pre_trained_model_type == "model":
model = torch.load(opt.pre_trained_model_path)
else:
model = Yolo(80)
model.load_state_dict(torch.load(opt.pre_trained_model_path))
else:
if opt.pre_trained_model_type == "model":
model = torch.load(opt.pre_trained_model_path,
map_location=lambda storage, loc: storage)
else:
model = Yolo(80)
model.load_state_dict(
torch.load(opt.pre_trained_model_path,
map_location=lambda storage, loc: storage))
model.eval()
colors = pickle.load(open("src/pallete", "rb"))
for image_path in glob.iglob(opt.input + os.sep + '*.png'):
if "prediction" in image_path:
continue
image = cv2.imread(image_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
height, width = image.shape[:2]
image = cv2.resize(image, (opt.image_size, opt.image_size))
image = np.transpose(np.array(image, dtype=np.float32), (2, 0, 1))
image = image[None, :, :, :]
width_ratio = float(opt.image_size) / width
height_ratio = float(opt.image_size) / height
data = Variable(torch.FloatTensor(image))
if torch.cuda.is_available():
data = data.cuda()
with torch.no_grad():
logits = model(data)
predictions = post_processing(logits, opt.image_size, CLASSES,
model.anchors, opt.conf_threshold,
opt.nms_threshold)
if len(predictions) != 0:
predictions = predictions[0]
output_image = cv2.imread(image_path)
for pred in predictions:
xmin = int(max(pred[0] / width_ratio, 0))
ymin = int(max(pred[1] / height_ratio, 0))
xmax = int(min((pred[0] + pred[2]) / width_ratio, width))
ymax = int(min((pred[1] + pred[3]) / height_ratio, height))
color = colors[CLASSES.index(pred[5])]
cv2.rectangle(output_image, (xmin, ymin), (xmax, ymax), color,
2)
cv2.rectangle(output_image, (0, 0), (width - 1, height - 1),
(150, 150, 150), 1)
text_size = cv2.getTextSize(pred[5] + ' : %.2f' % pred[4],
cv2.FONT_HERSHEY_PLAIN, 1, 1)[0]
# cv2.rectangle(output_image, (xmin, ymin), (xmin + text_size[0] + 3, ymin + text_size[1] + 4), color, -1)
cv2.putText(output_image, pred[5] + ' : %.2f' % pred[4],
(xmin, ymin - text_size[1] - 0 * 4),
cv2.FONT_HERSHEY_PLAIN, 1, (255, 255, 255), 1)
print("Object: {}, Bounding box: ({},{}) ({},{})".format(
pred[5], xmin, xmax, ymin, ymax))
cv2.imwrite(image_path[:-4] + "_prediction.jpg", output_image)
def test(opt):
if torch.cuda.is_available():
if opt.pre_trained_model_type == "model":
model = torch.load(opt.pre_trained_model_path)
else:
model = Yolo(20)
model.load_state_dict(torch.load(opt.pre_trained_model_path))
else:
if opt.pre_trained_model_type == "model":
model = torch.load(opt.pre_trained_model_path, map_location=lambda storage, loc: storage)
else:
model = Yolo(20)
model.load_state_dict(torch.load(opt.pre_trained_model_path, map_location=lambda storage, loc: storage))
model.eval()
colors = pickle.load(open("src/pallete", "rb"))
cap = cv2.VideoCapture(opt.input)
out = cv2.VideoWriter(opt.output, cv2.VideoWriter_fourcc(*"MJPG"), int(cap.get(cv2.CAP_PROP_FPS)),
(int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))))
while cap.isOpened():
flag, image = cap.read()
output_image = np.copy(image)
if flag:
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
else:
break
height, width = image.shape[:2]
image = cv2.resize(image, (opt.image_size, opt.image_size))
image = np.transpose(np.array(image, dtype=np.float32), (2, 0, 1))
image = image[None, :, :, :]
width_ratio = float(opt.image_size) / width
height_ratio = float(opt.image_size) / height
data = Variable(torch.FloatTensor(image))
if torch.cuda.is_available():
data = data.cuda()
with torch.no_grad():
logits = model(data)
predictions = post_processing(logits, opt.image_size, CLASSES, model.anchors, opt.conf_threshold,
opt.nms_threshold)
if len(predictions) != 0:
predictions = predictions[0]
for pred in predictions:
xmin = int(max(pred[0] / width_ratio, 0))
ymin = int(max(pred[1] / height_ratio, 0))
xmax = int(min((pred[0] + pred[2]) / width_ratio, width))
ymax = int(min((pred[1] + pred[3]) / height_ratio, height))
color = colors[CLASSES.index(pred[5])]
cv2.rectangle(output_image, (xmin, ymin), (xmax, ymax), color, 2)
text_size = cv2.getTextSize(pred[5] + ' : %.2f' % pred[4], cv2.FONT_HERSHEY_PLAIN, 1, 1)[0]
cv2.rectangle(output_image, (xmin, ymin), (xmin + text_size[0] + 3, ymin + text_size[1] + 4), color, -1)
cv2.putText(
output_image, pred[5] + ' : %.2f' % pred[4],
(xmin, ymin + text_size[1] + 4), cv2.FONT_HERSHEY_PLAIN, 1,
(255, 255, 255), 1)
out.write(output_image)
cap.release()
out.release()
def __init__(self, opt = get_args()):
if torch.cuda.is_available():
if opt.pre_trained_model_type == "model":
model = torch.load(opt.pre_trained_model_path)
else:
model = Yolo(80)
model.load_state_dict(torch.load(opt.pre_trained_model_path))
else:
if opt.pre_trained_model_type == "model":
model = torch.load(opt.pre_trained_model_path, map_location=lambda storage, loc: storage)
else:
model = Yolo(80)
model.load_state_dict(torch.load(opt.pre_trained_model_path, map_location=lambda storage, loc: storage))
model.eval()
self.model = model
self.opt = opt
def test(opt):
input_list_path = os.path.join(
opt.data_path, "VOC{}".format(opt.year),
"ImageSets/Main/{}.txt".format(opt.test_set))
image_ids = [id.strip() for id in open(input_list_path)]
output_folder = os.path.join(opt.output,
"VOC{}_{}".format(opt.year, opt.test_set))
colors = pickle.load(open("src/pallete", "rb"))
if os.path.isdir(output_folder):
shutil.rmtree(output_folder)
os.makedirs(output_folder)
if torch.cuda.is_available():
if opt.pre_trained_model_type == "model":
model = torch.load(opt.pre_trained_model_path)
else:
model = Yolo(20)
model.load_state_dict(torch.load(opt.pre_trained_model_path))
else:
if opt.pre_trained_model_type == "model":
model = torch.load(opt.pre_trained_model_path,
map_location=lambda storage, loc: storage)
else:
model = Yolo(20)
model.load_state_dict(
torch.load(opt.pre_trained_model_path,
map_location=lambda storage, loc: storage))
model.eval()
for id in image_ids:
image_path = os.path.join(opt.data_path, "VOC{}".format(opt.year),
"JPEGImages", "{}.jpg".format(id))
image = cv2.imread(image_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
height, width = image.shape[:2]
image = cv2.resize(image, (opt.image_size, opt.image_size))
image = np.transpose(np.array(image, dtype=np.float32), (2, 0, 1))
image = image[None, :, :, :]
width_ratio = float(opt.image_size) / width
height_ratio = float(opt.image_size) / height
data = Variable(torch.FloatTensor(image))
if torch.cuda.is_available():
data = data.cuda()
with torch.no_grad():
logits = model(data)
predictions = post_processing(logits, opt.image_size, CLASSES,
model.anchors, opt.conf_threshold,
opt.nms_threshold)
if len(predictions) == 0:
continue
else:
predictions = predictions[0]
output_image = cv2.imread(image_path)
for pred in predictions:
xmin = int(max(pred[0] / width_ratio, 0))
ymin = int(max(pred[1] / height_ratio, 0))
xmax = int(min((pred[0] + pred[2]) / width_ratio, width))
ymax = int(min((pred[1] + pred[3]) / height_ratio, height))
color = colors[CLASSES.index(pred[5])]
cv2.rectangle(output_image, (xmin, ymin), (xmax, ymax), color, 2)
text_size = cv2.getTextSize(pred[5] + ' : %.2f' % pred[4],
cv2.FONT_HERSHEY_PLAIN, 1, 1)[0]
cv2.rectangle(output_image, (xmin, ymin),
(xmin + text_size[0] + 3, ymin + text_size[1] + 4),
color, -1)
cv2.putText(output_image, pred[5] + ' : %.2f' % pred[4],
(xmin, ymin + text_size[1] + 4),
cv2.FONT_HERSHEY_PLAIN, 1, (255, 255, 255), 1)
print("Object: {}, Bounding box: ({},{}) ({},{})".format(
pred[5], xmin, xmax, ymin, ymax))
cv2.imwrite("{}/{}_prediction.jpg".format(output_folder, id),
output_image)
def train(opt):
if torch.cuda.is_available():
torch.cuda.manual_seed(123)
else:
torch.manual_seed(123)
learning_rate_schedule = {"0": 1e-5, "5": 1e-4, "80": 1e-5, "110": 1e-6}
twoShapesAnchors = [(1.0, 1.0), (3.0, 3.0), (6, 6), (9.0, 9.0),
(10.0, 5.0), (5.0, 10.0)]
twoShapesAnchors = [(1.3221, 1.73145), (3.19275, 4.00944),
(5.05587, 8.09892), (9.47112, 4.84053),
(11.2364, 10.0071)]
training_params = {
"batch_size": opt.batch_size,
"shuffle": True,
"drop_last": True,
"collate_fn": custom_collate_fn
}
test_params = {
"batch_size": opt.batch_size,
"shuffle": False,
"drop_last": False,
"collate_fn": custom_collate_fn
}
training_set = TwoShapesDataset(root_path=opt.data_path,
mode=opt.train_set,
trainingSet="annotations",
image_size=opt.image_size)
training_generator = DataLoader(training_set, **training_params)
test_set = TwoShapesDataset(root_path=opt.data_path,
mode=opt.test_set,
trainingSet="annotations",
image_size=opt.image_size,
is_training=False)
test_generator = DataLoader(test_set, **test_params)
if torch.cuda.is_available():
if opt.pre_trained_model_type == "model":
model = torch.load(opt.pre_trained_model_path)
elif opt.pre_trained_model_type == "params":
model = Yolo(training_set.num_classes, twoShapesAnchors)
model.load_state_dict(torch.load(opt.pre_trained_model_path))
else:
print("Just loading the model definintion")
model = Yolo(training_set.num_classes, twoShapesAnchors)
else:
if opt.pre_trained_model_type == "model":
model = torch.load(opt.pre_trained_model_path,
map_location=lambda storage, loc: storage)
elif opt.pre_trained_model_type == "params":
model = Yolo(training_set.num_classes, twoShapesAnchors)
model.load_state_dict(
torch.load(opt.pre_trained_model_path,
map_location=lambda storage, loc: storage))
# The following line will re-initialize weight for the last layer, which is useful
# when you want to retrain the model based on my trained weights. if you uncomment it,
# you will see the loss is already very small at the beginning.
# nn.init.normal_(list(model.modules())[:].weight, 0, 0.01)
# model.apply(weights_init)
log_path = os.path.join(opt.log_path, "{}".format("test1"))
if os.path.isdir(log_path):
shutil.rmtree(log_path)
os.makedirs(log_path)
writer = SummaryWriter(log_path)
if torch.cuda.is_available():
writer.add_graph(
model.cpu(),
torch.rand(opt.batch_size, 3, opt.image_size, opt.image_size))
model.cuda()
else:
writer.add_graph(
model, torch.rand(opt.batch_size, 3, opt.image_size,
opt.image_size))
criterion = YoloLoss(training_set.num_classes, model.anchors,
opt.reduction)
optimizer = torch.optim.SGD(model.parameters(),
lr=1e-5,
momentum=opt.momentum,
weight_decay=opt.decay)
best_loss = 1e10
best_epoch = 0
model.train()
num_iter_per_epoch = len(training_generator)
for epoch in range(opt.num_epoches):
if str(epoch) in learning_rate_schedule.keys():
for param_group in optimizer.param_groups:
param_group['lr'] = learning_rate_schedule[str(epoch)]
for iter, batch in enumerate(training_generator):
image, label = batch
if torch.cuda.is_available():
image = Variable(image.cuda(), requires_grad=True)
else:
image = Variable(image, requires_grad=True)
optimizer.zero_grad()
logits = model(image)
loss, loss_coord, loss_conf, loss_cls = criterion(logits, label)
loss.backward()
optimizer.step()
print(
"Epoch: {}/{}, Iteration: {}/{}, Lr: {}, Loss:{:.2f} (Coord:{:.2f} Conf:{:.2f} Cls:{:.2f})"
.format(epoch + 1, opt.num_epoches, iter + 1,
num_iter_per_epoch, optimizer.param_groups[0]['lr'],
loss, loss_coord, loss_conf, loss_cls))
writer.add_scalar('Train/Total_loss', loss,
epoch * num_iter_per_epoch + iter)
writer.add_scalar('Train/Coordination_loss', loss_coord,
epoch * num_iter_per_epoch + iter)
writer.add_scalar('Train/Confidence_loss', loss_conf,
epoch * num_iter_per_epoch + iter)
writer.add_scalar('Train/Class_loss', loss_cls,
epoch * num_iter_per_epoch + iter)
if epoch % opt.test_interval == 0:
model.eval()
loss_ls = []
loss_coord_ls = []
loss_conf_ls = []
loss_cls_ls = []
for te_iter, te_batch in enumerate(test_generator):
te_image, te_label = te_batch
num_sample = len(te_label)
if torch.cuda.is_available():
te_image = te_image.cuda()
with torch.no_grad():
te_logits = model(te_image)
batch_loss, batch_loss_coord, batch_loss_conf, batch_loss_cls = criterion(
te_logits, te_label)
loss_ls.append(batch_loss * num_sample)
loss_coord_ls.append(batch_loss_coord * num_sample)
loss_conf_ls.append(batch_loss_conf * num_sample)
loss_cls_ls.append(batch_loss_cls * num_sample)
te_loss = sum(loss_ls) / test_set.__len__()
te_coord_loss = sum(loss_coord_ls) / test_set.__len__()
te_conf_loss = sum(loss_conf_ls) / test_set.__len__()
te_cls_loss = sum(loss_cls_ls) / test_set.__len__()
print(
"Epoch: {}/{}, Lr: {}, Loss:{:.2f} (Coord:{:.2f} Conf:{:.2f} Cls:{:.2f})"
.format(epoch + 1, opt.num_epoches,
optimizer.param_groups[0]['lr'], te_loss,
te_coord_loss, te_conf_loss, te_cls_loss))
writer.add_scalar('Test/Total_loss', te_loss, epoch)
writer.add_scalar('Test/Coordination_loss', te_coord_loss, epoch)
writer.add_scalar('Test/Confidence_loss', te_conf_loss, epoch)
writer.add_scalar('Test/Class_loss', te_cls_loss, epoch)
model.train()
if te_loss + opt.es_min_delta < best_loss:
best_loss = te_loss
best_epoch = epoch
# torch.save(model, opt.saved_path + os.sep + "trained_yolo_coco")
torch.save(
model.state_dict(),
opt.saved_path + os.sep + "only_params_trained_yolo_coco")
torch.save(model, opt.saved_path + os.sep + opt.save_file)
# Early stopping
if epoch - best_epoch > opt.es_patience > 0:
print(
"Stop training at epoch {}. The lowest loss achieved is {}"
.format(epoch, te_loss))
break
writer.export_scalars_to_json(log_path + os.sep + "all_logs.json")
writer.close()
def test(opt):
model = Yolo(1).cuda()
model.load_state_dict(torch.load(opt.pre_trained_model_path))
model.eval()
colors = pickle.load(open("src/pallete", "rb"))
result_csv = []
result_csv.append(['patientId','x','y','width','height','Target'])
# for id in range(4998):
for id in range(20000,21764):
msg = "solving [%06d/%06d]" % (id+1, 4998)
print(msg, end='', flush=True)
back = '\b'*len(msg)
print(back, end='', flush=True)
image_path = 'data/norm_train_image/train%05d.png' % id
image = cv2.imread(image_path)
# image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
height, width = image.shape[:2]
image = cv2.resize(image, (512, 512))
image = np.transpose(np.array(image, dtype=np.float32), (2, 0, 1))
# image = image[0]
# image = np.expand_dims(image, axis=0)
image = np.expand_dims(image, axis=0)
# print(image.shape)
width_ratio = float(opt.image_size) / width
height_ratio = float(opt.image_size) / height
data = Variable(torch.FloatTensor(image))
if (width_ratio != 1 or height_ratio != 1):
print('\nNOT ratioa 1!!!\n')
if torch.cuda.is_available():
data = data.cuda()
with torch.no_grad():
logits = model(data)
predictions = post_processing(logits, opt.image_size, CLASSES, model.anchors, opt.conf_threshold,
opt.nms_threshold)
if len(predictions) != 0:
predictions = predictions[0]
output_image = cv2.imread(image_path)
for pred in predictions:
xmin = int(max(pred[0] / width_ratio, 0))
ymin = int(max(pred[1] / height_ratio, 0))
xmax = int(min((pred[0]+pred[2]) / width_ratio, width))
ymax = int(min(pred[1]+(pred[3]) / height_ratio, height))
w = int(min((pred[2]) / width_ratio, width))
h = int(min((pred[3]) / height_ratio, height))
if(xmin+w > 511):
w = 511-xmin
if(ymin+h > 511):
h = 511-ymin
output = ['train%05d.png' % id,xmin*2,ymin*2,w*2,h*2,1]
result_csv.append(output)
color = colors[CLASSES.index(pred[5])]
cv2.rectangle(output_image, (xmin, ymin), (xmax, ymax), color, 2)
text_size = cv2.getTextSize(pred[5] + ' : %.2f' % pred[4], cv2.FONT_HERSHEY_PLAIN, 1, 1)[0]
cv2.rectangle(output_image, (xmin, ymin), (xmin + text_size[0] + 3, ymin + text_size[1] + 4), color, -1)
cv2.putText(
output_image, pred[5] + ' : %.2f' % pred[4],
(xmin, ymin + text_size[1] + 4), cv2.FONT_HERSHEY_PLAIN, 1,
(255, 255, 255), 1)
print("Image: {} Object: {}, Bounding box: ({},{}) ({},{})".format(image_path, pred[5], xmin, xmax, ymin, ymax))
cv2.imwrite('predictions/'+str(id) + "_prediction.jpg", output_image)
# test_images/
else:
output = ['test%04d.png' % id,'','','','',0]
result_csv.append(output)
def train(opt):
if not os.path.isdir(opt.saved_path):
os.makedirs(opt.saved_path)
os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpu_devices
if torch.cuda.is_available():
torch.cuda.manual_seed(123)
else:
torch.manual_seed(123)
learning_rate_schedule = {"0": 1e-5, "5": 1e-4, "80": 1e-5, "110": 1e-6}
training_params = {
"batch_size": opt.batch_size,
"shuffle": True,
"drop_last": True,
"collate_fn": custom_collate_fn
}
test_params = {
"batch_size": opt.batch_size,
"shuffle": False,
"drop_last": False,
"collate_fn": custom_collate_fn
}
training_set = BOTDataset_scene(opt.data_path,
opt.train_set,
opt.image_size,
scene=2)
# mmm = training_set[0]
# print(1)
training_generator = DataLoader(training_set, **training_params)
test_set = BOTDataset_scene(opt.data_path,
opt.test_set,
opt.image_size,
scene=2,
is_training=False)
test_generator = DataLoader(test_set, **test_params)
if torch.cuda.is_available():
if opt.pre_trained_model_type == "model":
model = torch.load(opt.pre_trained_model_path)
else:
# model = Yolo(training_set.num_classes)
# model.load_state_dict(torch.load(opt.pre_trained_model_path))
print('总类别数据：', training_set.num_classes)
model = Yolo_Person2(training_set.num_classes)
pretrained_dict = torch.load(opt.pre_trained_model_path)
model_dict = model.state_dict()
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in model_dict
}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
print('Success Init BOT single class model')
else:
if opt.pre_trained_model_type == "model":
model = torch.load(opt.pre_trained_model_path,
map_location=lambda storage, loc: storage)
else:
model = Yolo(training_set.num_classes)
model.load_state_dict(
torch.load(opt.pre_trained_model_path,
map_location=lambda storage, loc: storage))
# The following line will re-initialize weight for the last layer, which is useful
# when you want to retrain the model based on my trained weights. if you uncomment it,
# you will see the loss is already very small at the beginning.
nn.init.normal_(list(model.modules())[-1].weight, 0, 0.01)
log_path = os.path.join(opt.log_path, "{}".format(opt.year))
if os.path.isdir(log_path):
shutil.rmtree(log_path)
os.makedirs(log_path)
writer = SummaryWriter(log_path)
if torch.cuda.is_available():
writer.add_graph(
model.cpu(),
torch.rand(opt.batch_size, 3, opt.image_size, opt.image_size))
model.cuda()
else:
writer.add_graph(
model, torch.rand(opt.batch_size, 3, opt.image_size,
opt.image_size))
# print('training_set.num_classes',training_set.num_classes)
criterion = YoloLoss(training_set.num_classes, model.anchors,
opt.reduction)
# criterion = YoloLoss(training_set.num_classes, model.anchors, opt.reduction)
optimizer = torch.optim.SGD(model.parameters(),
lr=1e-5,
momentum=opt.momentum,
weight_decay=opt.decay)
best_loss = 1e10
best_epoch = 0
model.train()
num_iter_per_epoch = len(training_generator)
for epoch in range(opt.num_epoches):
if str(epoch) in learning_rate_schedule.keys():
for param_group in optimizer.param_groups:
param_group['lr'] = learning_rate_schedule[str(epoch)]
for iter, batch in enumerate(training_generator):
image, label = batch
if torch.cuda.is_available():
image = Variable(image.cuda(), requires_grad=True)
else:
image = Variable(image, requires_grad=True)
optimizer.zero_grad()
logits = model(image)
# print('logits.shape',logits.shape)
# print('label.shape',len(label))
# print('label',label)
loss, loss_coord, loss_conf, loss_cls = criterion(logits, label)
loss.backward()
optimizer.step()
print(
"Epoch: {}/{}, Iteration: {}/{}, Lr: {}, Loss:{:.2f} (Coord:{:.2f} Conf:{:.2f} Cls:{:.2f})"
.format(epoch + 1, opt.num_epoches, iter + 1,
num_iter_per_epoch, optimizer.param_groups[0]['lr'],
loss, loss_coord, loss_conf, loss_cls))
writer.add_scalar('Train/Total_loss', loss,
epoch * num_iter_per_epoch + iter)
writer.add_scalar('Train/Coordination_loss', loss_coord,
epoch * num_iter_per_epoch + iter)
writer.add_scalar('Train/Confidence_loss', loss_conf,
epoch * num_iter_per_epoch + iter)
writer.add_scalar('Train/Class_loss', loss_cls,
epoch * num_iter_per_epoch + iter)
if epoch % opt.test_interval == 0:
model.eval()
loss_ls = []
loss_coord_ls = []
loss_conf_ls = []
loss_cls_ls = []
for te_iter, te_batch in enumerate(test_generator):
te_image, te_label = te_batch
num_sample = len(te_label)
if torch.cuda.is_available():
te_image = te_image.cuda()
with torch.no_grad():
te_logits = model(te_image)
batch_loss, batch_loss_coord, batch_loss_conf, batch_loss_cls = criterion(
te_logits, te_label)
loss_ls.append(batch_loss * num_sample)
loss_coord_ls.append(batch_loss_coord * num_sample)
loss_conf_ls.append(batch_loss_conf * num_sample)
loss_cls_ls.append(batch_loss_cls * num_sample)
te_loss = sum(loss_ls) / test_set.__len__()
te_coord_loss = sum(loss_coord_ls) / test_set.__len__()
te_conf_loss = sum(loss_conf_ls) / test_set.__len__()
te_cls_loss = sum(loss_cls_ls) / test_set.__len__()
print(
"Epoch: {}/{}, Lr: {}, Loss:{:.2f} (Coord:{:.2f} Conf:{:.2f} Cls:{:.2f})"
.format(epoch + 1, opt.num_epoches,
optimizer.param_groups[0]['lr'], te_loss,
te_coord_loss, te_conf_loss, te_cls_loss))
writer.add_scalar('Test/Total_loss', te_loss, epoch)
writer.add_scalar('Test/Coordination_loss', te_coord_loss, epoch)
writer.add_scalar('Test/Confidence_loss', te_conf_loss, epoch)
writer.add_scalar('Test/Class_loss', te_cls_loss, epoch)
model.train()
if te_loss + opt.es_min_delta < best_loss:
best_loss = te_loss
best_epoch = epoch
# torch.save(model, opt.saved_path + os.sep + "trained_yolo_voc")
torch.save(
model.state_dict(),
opt.saved_path + os.sep + "only_params_trained_yolo_bot")
torch.save(
model,
opt.saved_path + os.sep + "whole_model_trained_yolo_bot")
# Early stopping
if epoch - best_epoch > opt.es_patience > 0:
print(
"Stop training at epoch {}. The lowest loss achieved is {}"
.format(epoch, te_loss))
break
writer.export_scalars_to_json(log_path + os.sep + "all_logs.json")
writer.close()
def test(opt):
global colors
if torch.cuda.is_available():
if opt.pre_trained_model_type == "model":
model1 = torch.load(opt.pre_trained_model_path)
else:
model1 = Yolo(1)
model1.load_state_dict(torch.load(opt.pre_trained_model_path))
colors = pickle.load(open("src/pallete", "rb"))
model1.eval()
img_list = sorted(glob.glob(os.path.join(opt.data_path_test, '*.jpg')))
print(img_list)
save_dir = './anotherMissOh_Test_Result/'
for idx, idx_img in enumerate(img_list):
image = Image.open(idx_img)
np_img = np.asarray(image)
image = cv2.cvtColor(np.float32(np_img), cv2.COLOR_RGB2BGR)
height, width = image.shape[:2]
image = cv2.resize(image, (opt.image_size, opt.image_size))
image = np.transpose(np.array(image, dtype=np.float32), (2, 0, 1))
image = image[None, :, :, :]
width_ratio = float(opt.image_size) / width
height_ratio = float(opt.image_size) / height
data = Variable(torch.FloatTensor(image))
if torch.cuda.is_available():
data = data.cuda()
with torch.no_grad():
logits = model1(data)
predictions = post_processing(logits, opt.image_size,
MissOh_CLASSES, model1.anchors,
opt.conf_threshold,
opt.nms_threshold)
if len(predictions) != 0:
predictions = predictions[0]
output_image = cv2.cvtColor(np.float32(np_img), cv2.COLOR_RGB2BGR)
for pred in predictions:
xmin = int(max(pred[0] / width_ratio, 0))
ymin = int(max(pred[1] / height_ratio, 0))
xmax = int(min((pred[0] + pred[2]) / width_ratio, width))
ymax = int(min((pred[1] + pred[3]) / height_ratio, height))
color = colors[MissOh_CLASSES.index(pred[5])]
cv2.rectangle(output_image, (xmin, ymin), (xmax, ymax), color,
2)
text_size = cv2.getTextSize(pred[5] + ' : %.2f' % pred[4],
cv2.FONT_HERSHEY_PLAIN, 1, 1)[0]
cv2.rectangle(
output_image, (xmin, ymin),
(xmin + text_size[0] + 3, ymin + text_size[1] + 4), color,
-1)
cv2.putText(output_image, pred[5] + ' : %.2f' % pred[4],
(xmin, ymin + text_size[1] + 4),
cv2.FONT_HERSHEY_PLAIN, 1, (255, 255, 255), 1)
cv2.imwrite(save_dir + "entire_{}.png".format(idx), output_image)
image_list.append(image_id)
num_obj = (len(image_file) - 1) // 5
for i in range(num_obj):
x1 = int(image_file[1 + 5 * i])
y1 = int(image_file[2 + 5 * i])
x2 = int(image_file[3 + 5 * i])
y2 = int(image_file[4 + 5 * i])
c = int(image_file[5 + 5 * i])
class_name = VOC_CLASSES[c]
target[(image_id, class_name)].append([x1, y1, x2, y2])
# start test
print('---start test---')
model = Yolo(20)
model.load_state_dict(torch.load('./trained_models/only_params_trained_yolo_voc'))
model.eval()
model.cuda()
with torch.no_grad():
count = 0
for image_path in tqdm(image_list):
image = cv2.imread(os.path.join('/BS/database11/allimgs/', image_path))
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
height, width = image.shape[:2]
image = cv2.resize(image, (image_size, image_size))
image = np.transpose(np.array(image, dtype=np.float32), (2, 0, 1))
image = image[None, :, :, :]
width_ratio = float(image_size) / width
height_ratio = float(image_size) / height
data = Variable(torch.FloatTensor(image))