def vis_iou(self, img, scale, aspect):
output = self.true_pred[scale]
bbox = self.targets[img]['bbox'].cuda() * self.img_size
res = boxes.box_iou(helper.get_abs_coord(output[img, :, :4]),
helper.get_abs_coord(bbox)).max(axis=1)[0]
in_h = (output.shape[1] / self.num_anchors)**0.5
in_h = int(in_h)
vis_res = res.view(self.num_anchors, in_h, in_h)
ax = sns.heatmap(vis_res[aspect].cpu().numpy())
def show_gt(self, k, color=(0, 0, 0), linelen=1):
img2show = self.image[k]
mean = torch.tensor([[[0.485, 0.456, 0.406]]]).T
std = torch.tensor([[[0.229, 0.224, 0.225]]]).T
img2show = img2show * std + mean
img2show = img2show * 255
img2show = img2show.transpose(0, 1)
img2show = img2show.transpose(1, 2)
img2show = img2show.numpy().astype(np.uint8)
im = img2show.copy()
bbox = self.targets[k]['bbox'].cuda() * self.img_size
cords = helper.get_abs_coord(bbox)
labels = self.targets[k]['category_id'].cpu().numpy()
k = 0
for cord in cords:
pt1, pt2 = (cord[0], cord[1]), (cord[2], cord[3])
pt1 = int(pt1[0]), int(pt1[1])
pt2 = int(pt2[0]), int(pt2[1])
im = cv2.rectangle(im.copy(), pt1, pt2, color, linelen)
lb = self.get_category_name(labels[k])
im = cv2.putText(im.copy(), lb, pt1, 0, 1e-3 * self.img_size,
color, linelen // 2)
k = k + 1
io.imshow(im)
def vis_performance(self, img, scale, aspect):
output = self.true_pred[scale]
bbox = self.targets[img]['bbox'].cuda() * self.img_size
res = boxes.box_iou(helper.get_abs_coord(output[img, :, :4]),
helper.get_abs_coord(bbox)).max(axis=1)[0]
in_h = (output.shape[1] / self.num_anchors)**0.5
in_h = int(in_h)
iou = res.view(self.num_anchors, in_h, in_h)[aspect]
confidence = self.out[scale][img, aspect, :, :, 4]
classes = self.out[scale][img, aspect, :, :, 5:].max(axis=-1)
classes_values = classes[0]
classes_labels = classes[1]
gt_labels = self.targets[img]['category_id']
mask = torch.stack([classes_labels == lab for lab in gt_labels])
mask = mask.sum(axis=0)
mask[mask > 0] = 1
mask[mask == 0] = -1
performance_heatmap = confidence * iou * mask
ax = sns.heatmap(performance_heatmap.cpu().numpy())
def draw_bbs(self,
k,
confidence=0.1,
iou_threshold=0.5,
color=(0, 0, 0),
linelen=1):
predictions = torch.cat(self.true_pred, axis=1)[k]
predictions[:, :4] = helper.get_abs_coord(predictions[:, :4])
score = predictions[:, 4] * (predictions[:, 5:].max(axis=1)[0])
pred_mask = score > confidence
pred_conf = predictions[pred_mask]
indices = boxes.nms(pred_conf[:, :4], pred_conf[:, 4], iou_threshold)
pred_final = pred_conf[indices, :]
img2show = self.image[k]
mean = torch.tensor([[[0.485, 0.456, 0.406]]]).T
std = torch.tensor([[[0.229, 0.224, 0.225]]]).T
img2show = img2show * std + mean
img2show = img2show * 255
img2show = img2show.transpose(0, 1)
img2show = img2show.transpose(1, 2)
img2show = img2show.numpy().astype(np.uint8)
im = img2show.copy()
cords = pred_final.clone().cpu().numpy()
labels = (pred_final[:, 5:].max(axis=1)[1]).cpu().numpy()
k = 0
for cord in cords:
pt1, pt2 = (cord[0], cord[1]), (cord[2], cord[3])
pt1 = int(pt1[0]), int(pt1[1])
pt2 = int(pt2[0]), int(pt2[1])
im = cv2.rectangle(im.copy(), pt1, pt2, color, linelen)
lb = self.get_category_name(labels[k])
im = cv2.putText(im.copy(), lb, pt1, 0, 1e-3 * self.img_size,
color, linelen // 2)
k = k + 1
io.imshow(im)
targets2.append(dd)
# targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
targets = targets2
forw_dur = time.time()
out = model(images)
forw_dur = time.time() - forw_dur
forw_sum = forw_sum + forw_dur
outcome = []
for k, y in enumerate(out):
outcome.append(losses[k](y))
predictions = torch.cat(outcome, axis=1)
predictions[:, :, :4] = helper.get_abs_coord(
predictions[:, :, :4])
score = predictions[:, :, 4] * (predictions[:, :,
5:].max(axis=2)[0])
pred_mask = score > confidence
pred_conf = [(predictions[e][m])
for e, m in enumerate(pred_mask)]
indices = [
boxes.nms(pred_conf[i][:, :4], pred_conf[i][:, 4],
iou_threshold) for i in range(len(pred_conf))
]
pred_final = [
pred_conf[i][indices[i], :] for i in range(len(pred_conf))
]
pred_final = list(filter(lambda t: t.shape[0] != 0,
def test_one_epoch(dataloader, model, yolo_loss, cfg):
confidence = cfg.yolo.inf_confidence
iou_threshold = cfg.yolo.inf_iou_threshold
inp_dim = cfg.dataset.inp_dim
yolo_loss.set_img_size(inp_dim)
model.eval()
results = []
dset_name = dataloader.dset_name
torch.backends.cudnn.benchmark = True
with torch.no_grad():
for batch_idx, (images, targets) in enumerate(dataloader):
# measure data loading time
images = images.to('cuda', non_blocking=True)
targets = [{
k: v.to('cuda', non_blocking=True)
for k, v in t.items()
} for t in targets]
out = model(images)
predictions = yolo_loss(out)
predictions[:, :, :4] = helper.get_abs_coord(predictions[:, :, :4])
score = predictions[:, :, 4] * (predictions[:, :,
5:].max(axis=2)[0])
pred_mask = score > confidence
pred_conf = [(predictions[e][m]) for e, m in enumerate(pred_mask)]
indices = [
boxes.nms(pred_conf[i][:, :4], pred_conf[i][:, 4],
iou_threshold) for i in range(len(pred_conf))
]
pred_final = [
pred_conf[i][indices[i], :] for i in range(len(pred_conf))
]
pred_final = list(filter(lambda t: t.shape[0] != 0, pred_final))
for i, atrbs in enumerate(pred_final):
xmin = atrbs[:, 0] / inp_dim * targets[i]['img_size'][1]
ymin = atrbs[:, 1] / inp_dim * targets[i]['img_size'][0]
xmax = atrbs[:, 2] / inp_dim * targets[i]['img_size'][1]
ymax = atrbs[:, 3] / inp_dim * targets[i]['img_size'][0]
w = xmax - xmin
h = ymax - ymin
scores = (atrbs[:, 4] * atrbs[:, 5:].max(axis=1)[0]).tolist()
labels = (atrbs[:, 5:].max(axis=1)[1])
if dset_name == 'coco':
labels = helper.torch80_to_91(labels).tolist()
else:
labels = (labels + 1).tolist()
bboxes = torch.stack((xmin, ymin, w, h), axis=1)
areas = (bboxes[:, 2] * bboxes[:, 3]).tolist()
bboxes = bboxes.tolist()
temp = [{
'bbox': b,
'area': a,
'category_id': l,
'score': s,
'image_id': targets[i]['image_id'].item()
} for b, a, l, s in zip(bboxes, areas, labels, scores)]
results = list(itertools.chain(results, temp))
return results