def test(model, dataloader, epoch, opt, test_logger, visualizer=None):
labels = []
sample_matrics = []
for i, (images, targets) in enumerate(dataloader):
if targets.size(0) == 0:
continue
batches_done = len(dataloader) * epoch + i
if not opt.no_cuda:
images_cpu = images.clone()
model = model.to(opt.device)
images = Variable(images.to(opt.device))
if targets is not None:
targets = Variable(targets.to(opt.device), requires_grad=False)
labels += targets[:, 1].tolist()
targets[:, 2:] = xywh2xyxy(targets[:, 2:])
targets[:, 2:] *= opt.image_size
detections = model.forward(images)
detections = non_max_suppression(detections, opt.conf_thres,
opt.nms_thres)
sample_matrics += get_batch_statistics(detections,
targets.cpu(),
iou_threshold=0.5)
if visualizer is not None and not opt.no_vis_preds:
visualizer.plot_predictions(images.cpu(), detections,
env='main') # plot prediction
if visualizer is not None and not opt.no_vis_gt:
visualizer.plot_ground_truth(images.cpu(),
targets.cpu(),
env='main') # plot ground truth
true_positives, pred_scores, pred_labels = [
np.concatenate(x, 0) for x in list(zip(*sample_matrics))
]
precision, recall, AP, f1, ap_class = ap_per_class(true_positives,
pred_scores,
pred_labels, labels)
# logging
metric_table_data = [['Metrics', 'Value'], ['precision',
precision.mean()],
['recall', recall.mean()], ['f1', f1.mean()],
['mAP', AP.mean()]]
metric_table = AsciiTable(metric_table_data,
title='[Epoch {:d}/{:d}'.format(
epoch, opt.num_epochs))
print('{}\n\n\n'.format(metric_table.table))
class_names = load_classe_names(opt.classname_path)
for i, c in enumerate(ap_class):
metric_table_data += [['AP-{}'.format(class_names[c]), AP[i]]]
metric_table.table_data = metric_table_data
test_logger.write('{}\n\n\n'.format(metric_table.table))
if visualizer is not None:
vis.plot_metrics(metric_table_data, batches_done, env='main')
def val(model, optimizer, dataloader, epoch, vis, opt):
for i, (images, targets) in enumerate(dataloader):
labels += targets[:, 1].tolist()
targets[:, 2:] = xywh2xyxy(targets[:, 2:])
targets[:, 2:] *= opt.img_size
loss, yolo_outputs = model.forward(images, targets)
outputs = non_max_suppression(outputs, opt.conf_thres, opt.nms_thres)
error52 = model.yolo_layer52.metrics
error26 = model.yolo_layer26.metrics
error13 = model.yolo_layer13.metrics
vis.print_current_losses([error52, error26, error13], epoch, i,
len(dataloader))
vis.plot_current_visuals(images_cpu, yolo_outputs)
sample_matrics += get_batch_statistics(outputs,
targets,
iou_threshold=0.5)
true_positives, pred_scores, pred_labels = [
np.concatenate(x, 0) for x in list(zip(*sample_matrics))
]
precision, recall, AP, f1, ap_class = ap_per_class(true_positives,
pred_scores,
pred_labels, labels)
print("Average Precision for Val:")
for i, c in enumerate(ap_class):
print('Class %s - AP: %.4f' % (class_names[c], AP[i]))
print('mAP: %.4f' % (AP.mean()))
def test(model,dataloader,epoch,opt):
labels = []
sample_matrics = []
total_time = 0
pic_num = 0
first_run = True
if opt.gpu:
model.to(opt.device)
model.eval()
# warm-up
if opt.gpu:
input_shape = (3,416,416)
dummy_input = Variable(torch.randn(1,*input_shape))
model.forward(dummy_input.to(opt.device))
for i,(images,targets) in enumerate(dataloader):
pic_num += opt.batch_size
labels += targets[:,1].tolist()
targets[:,2:] = xywh2xyxy(targets[:,2:])
targets[:,2:] *= opt.image_size
if opt.gpu:
images = Variable(images.to(opt.device))
t_start = time.time()
detections = model.forward(images)
t_end = time.time()
detections = non_max_suppression(detections,opt.conf_thresh,opt.nms_thresh)
print("forward time:"+str(t_end - t_start))
total_time += t_end - t_start
sample_matrics += get_batch_statistics(detections,targets,iou_threshold=0.5)
print("Average time:"+str(total_time/pic_num) + "s")
true_positives,pred_scores,pred_labels = [np.concatenate(x,0) for x in list(zip(*sample_matrics))]
precision,recall,AP,f1,ap_class = ap_per_class(true_positives,pred_scores,pred_labels,labels)
metric_table_data = [
['Metrics','Value'],['precision',precision.mean()],['recall',recall.mean()],
['f1',f1.mean()],['mAP',AP.mean()]]
metric_table = AsciiTable(
metric_table_data,
title='[Epoch {:d}/{:d}'.format(epoch,opt.num_epochs))
class_names = load_classe_names(opt.classname_path)
for i,c in enumerate(ap_class):
metric_table_data += [['AP-{}'.format(class_names[c]),AP[i]]]
metric_table.table_data = metric_table_data
print('{}\n'.format(metric_table.table))
def val(model, optimizer, dataloader, epoch, opt, val_logger, visualizer=None):
labels = []
sample_matrics = []
for i, (images, targets) in enumerate(dataloader):
labels += targets[:, 1].tolist()
targets[:, 2:] = xywh2xyxy(targets[:, 2:])
targets[:, 2:] *= opt.image_size
outputs = model.forward(images)
outputs = non_max_suppression(outputs, opt.conf_thres, opt.nms_thres)
sample_matrics += get_batch_statistics(outputs,
targets,
iou_threshold=0.5)
if visualizer is not None:
vis.plot_current_visuals(images, outputs)
true_positives, pred_scores, pred_labels = [
np.concatenate(x, 0) for x in list(zip(*sample_matrics))
]
precision, recall, AP, f1, ap_class = ap_per_class(true_positives,
pred_scores,
pred_labels, labels)
# logging
metric_table_data = [['Metrics', 'Value'], ['precision',
precision.mean()],
['recall', recall.mean()], ['f1', f1.mean()],
['mAP', AP.mean()]]
metric_table = AsciiTable(metric_table_data,
title='[Epoch {:d}/{:d}'.format(
epoch, opt.num_epochs))
print('{}\n\n\n'.format(metric_table.table))
class_names = load_classe_names(opt.classname_path)
for i, c in enumerate(ap_class):
metric_table_data += [['AP-{}'.format(class_names[c]), AP[i]]]
metric_table.table_data = metric_table_data
val_logger.write('{}\n\n\n'.format(metric_table.table))
def detect(model, input, opt):
w,h = input.size
if not opt.no_pad2square:
if w == h:
image = input
else:
dim_diff = abs(w - h)
padding_1,padding_2 = dim_diff // 2,dim_diff - dim_diff // 2
padding = (0,padding_1,0,padding_2) if w > h else (padding_1,0,padding_2,0)
image = F.pad(input,padding,fill=0,padding_mode='constant')
else:
image = input
image = torchvision.transforms.ToTensor()(image).float().unsqueeze(0)
image = resize(image,opt.image_size).float().unsqueeze(0)
if opt.gpu:
model.to(opt.device)
image = image.to(opt.device)
t_start = time.time()
detections = model.forward(image)
t_end = time.time()
detections = non_max_suppression(detections,opt.conf_thresh,opt.nms_thresh)
print(detections)
print("inference time:" + str(t_end - t_start))
class_names = load_classe_names(opt.classname_path)
image = image.squeeze(0).transpose(0,1).transpose(1,2)
detection = detections[0]
fig,ax = plt.pyplot.subplots(1)
plt.pyplot.axis('off')
plt.pyplot.gca().xaxis.set_major_locator(plt.ticker.NullLocator())
plt.pyplot.gca().yaxis.set_major_locator(plt.ticker.NullLocator())
plt.pyplot.tight_layout(pad=0)
ax.imshow(image.cpu().numpy())
unique_labels = detection[:,-1].unique()
num_cls_preds = len(unique_labels)
color_map = plt.pyplot.get_cmap('tab20b')
colors = [color_map(i) for i in np.linspace(0,1,opt.num_classes)]
bbox_colors = random.sample(colors, num_cls_preds)
for xmin, ymin, xmax, ymax, conf, cls_conf, cls_pred in detection:
box_w = xmax - xmin
box_h = ymax - ymin
color = bbox_colors[int(np.where(unique_labels == int(cls_pred))[0])]
bbox = plt.patches.Rectangle((xmin,ymin),box_w,box_h,linewidth=2,edgecolor=color,facecolor="none")
ax.add_patch(bbox)
plt.pyplot.text(
xmin,
ymin,
s='{} {:.4f}'.format(class_names[int(cls_pred)],cls_conf),
color='white',
verticalalignment='top',
bbox={'color': color,'pad': 0},
)
plt.pyplot.show()
def train(model, optimizer, dataloader, epoch, opt, logger, visualizer=None):
for i, (images, targets) in enumerate(dataloader):
# targets: [idx, class_id, x, y, h, w] in yolo format
# idx is used to associate the bounding boxes with its image
# skip images without bounding boxes (mainly because coco has unlabelled images)
if targets.size(0) == 0:
continue
batches_done = len(dataloader) * epoch + i
if not opt.no_cuda:
model = model.to(opt.device)
images = Variable(images.to(opt.device))
if targets is not None:
targets = Variable(targets.to(opt.device), requires_grad=False)
loss, detections = model.forward(images, targets)
detections = non_max_suppression(detections.cpu(), opt.conf_thres,
opt.nms_thres)
loss.backward()
if batches_done % opt.gradient_accumulations == 0 or i == len(
dataloader) - 1:
optimizer.step()
optimizer.zero_grad()
# logging
metric_keys = model.yolo_layer52.metrics.keys()
yolo_metrics = [
model.yolo_layer52.metrics, model.yolo_layer26.metrics,
model.yolo_layer13.metrics
]
metric_table_data = [[
'Metrics', 'YOLO Layer 0', 'YOLO Layer 1', 'YOLO Layer 2'
]]
formats = {m: '%.6f' for m in metric_keys}
for metric in metric_keys:
row_metrics = [
formats[metric] % ym.get(metric, 0) for ym in yolo_metrics
]
metric_table_data += [[metric, *row_metrics]]
metric_table_data += [[
'total loss', '{:.6f}'.format(loss.item()), '', ''
]]
# beautify log message
metric_table = AsciiTable(
metric_table_data,
title='[Epoch {:d}/{:d}, Batch {:d}/{:d}]'.format(
epoch, opt.num_epochs, i, len(dataloader)))
metric_table.inner_footing_row_border = True
logger.print_and_write('{}\n\n\n'.format(metric_table.table))
if visualizer is not None and not opt.no_vis_preds:
visualizer.plot_predictions(images.cpu(), detections,
env='main') # plot prediction
if visualizer is not None and not opt.no_vis_gt:
visualizer.plot_ground_truth(images.cpu(),
targets.cpu(),
env='main') # plot ground truth
metrics_to_vis = []
# uncomment code below to plot the metrics of each YOLO layer
# for j, ym in enumerate(yolo_metrics):
# for key, metric in ym.items():
# if key != 'grid_size':
# metrics_to_vis += [('{}_yolo_layer_{}'.format(key, j), metric)]
metrics_to_vis += [('total_loss', loss.item())]
if visualizer is not None:
visualizer.plot_metrics(metrics_to_vis, batches_done, env='main')
# save checkpoints
if epoch % opt.checkpoint_interval == 0:
save_file_path = os.path.join(opt.checkpoint_path,
'epoch_{}.pth'.format(epoch))
states = {
'epoch': epoch + 1,
'model': opt.model,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
}
torch.save(states, save_file_path)
def val(model, optimizer, dataloader, epoch, opt, val_logger, best_mAP=0):
labels = []
sample_matrics = []
if opt.gpu:
model = model.to(opt.device)
model.eval()
for i, (images, targets) in enumerate(dataloader):
labels += targets[:, 1].tolist()
targets[:, 2:] = xywh2xyxy(targets[:, 2:])
targets[:, 2:] *= opt.image_size
batches_done = len(dataloader) * epoch + i
if opt.gpu:
images = Variable(images.to(opt.device))
detections = model.forward(images)
detections = non_max_suppression(detections, opt.conf_thresh,
opt.nms_thresh)
sample_matrics += get_batch_statistics(detections,
targets.cpu(),
iou_threshold=0.5)
true_positives, pred_scores, pred_labels = [
np.concatenate(x, 0) for x in list(zip(*sample_matrics))
]
precision, recall, AP, f1, ap_class = ap_per_class(true_positives,
pred_scores,
pred_labels, labels)
# logging
metric_table_data = [['Metrics', 'Value'], ['precision',
precision.mean()],
['recall', recall.mean()], ['f1', f1.mean()],
['mAP', AP.mean()]]
metric_table = AsciiTable(metric_table_data,
title='[Epoch {:d}/{:d}'.format(
epoch, opt.num_epochs))
class_names = load_classe_names(opt.classname_path)
for i, c in enumerate(ap_class):
metric_table_data += [['AP-{}'.format(class_names[c]), AP[i]]]
metric_table.table_data = metric_table_data
val_logger.print_and_write('{}\n'.format(metric_table.table))
if best_mAP < AP.mean():
save_file_path = os.path.join(opt.checkpoint_path, 'best.pth')
states = {
'epoch': epoch + 1,
'model': opt.model,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
'best_mAP': best_mAP,
}
torch.save(states, save_file_path)
best_mAP = AP.mean()
print("current best mAP:" + str(best_mAP))
return best_mAP