def _vis_eval(self, epoch, i, len_loader, img, heatmap, adj_mtx,
junctions_pred, junctions_gt, adj_mtx_gt):
junctions_gt = np.int32(junctions_gt)
lines_gt, scores_gt = graph2line(junctions_gt,
adj_mtx_gt,
threshold=self.vis_junc_th)
vis_line_gt = vutils.make_grid(draw_lines(img, lines_gt, scores_gt))
img_with_junc = draw_jucntions(img, junctions_pred)
img_with_junc = torch.stack(img_with_junc,
dim=0).numpy()[:, ::-1, :, :]
lines_pred, score_pred = graph2line(junctions_pred, adj_mtx)
vis_line_pred = vutils.make_grid(
draw_lines(img_with_junc, lines_pred, score_pred))
junc_score = []
line_score = []
for m, juncs in zip(heatmap, junctions_gt):
juncs = juncs[juncs.sum(axis=1) > 0]
junc_score += m[juncs[:, 1], juncs[:, 0]].tolist()
for s in score_pred:
line_score += s.tolist()
junc_pooling = vutils.make_grid(draw_jucntions(heatmap,
junctions_pred))
self.writer.add_image(self.exp_name + "/" + "eval/junction_pooling",
junc_pooling, epoch * len_loader + i)
self.writer.add_image(self.exp_name + "/" + "eval/lines_gt",
vis_line_gt, epoch * len_loader + i)
self.writer.add_image(self.exp_name + "/" + "eval/lines_pred",
vis_line_pred, epoch * len_loader + i)
self.writer.add_scalar(self.exp_name + "/" + "eval/mean_junc_score",
np.mean(junc_score), epoch * len_loader + i)
self.writer.add_scalar(self.exp_name + "/" + "eval/mean_line_score",
np.mean(line_score), epoch * len_loader + i)
def _vis_train(self, epoch, i, len_loader, img, heatmap, adj_mtx,
junctions_gt, adj_mtx_gt):
junctions_gt = np.int32(junctions_gt)
lines_gt, scores_gt = graph2line(junctions_gt, adj_mtx_gt)
vis_line_gt = vutils.make_grid(draw_lines(img, lines_gt, scores_gt))
lines_pred, score_pred = graph2line(junctions_gt,
adj_mtx,
threshold=self.vis_line_th)
vis_line_pred = vutils.make_grid(
draw_lines(img, lines_pred, score_pred))
junc_score = []
line_score = []
for m, juncs in zip(heatmap, junctions_gt):
juncs = juncs[juncs.sum(axis=1) > 0]
junc_score += m[juncs[:, 1], juncs[:, 0]].tolist()
for s in score_pred:
line_score += s.tolist()
self.writer.add_image(self.exp_name + "/" + "train/lines_gt",
vis_line_gt, epoch * len_loader + i)
self.writer.add_image(self.exp_name + "/" + "train/lines_pred",
vis_line_pred, epoch * len_loader + i)
self.writer.add_scalar(self.exp_name + "/" + "train/mean_junc_score",
np.mean(junc_score), epoch * len_loader + i)
self.writer.add_scalar(self.exp_name + "/" + "train/mean_line_score",
np.mean(line_score), epoch * len_loader + i)
def test(self, path_to_image):
# main loop
torch.set_grad_enabled(False)
print(f"test for image: {path_to_image}", flush=True)
if self.is_cuda:
model = self.model.cuda().eval()
else:
model = self.model.eval()
img = cv2.imread(path_to_image)
img = cv2.resize(img, (self.img_size, self.img_size))
img = torch.from_numpy(img[:, :, ::-1]).float().permute(2, 0,
1).unsqueeze(0)
if self.is_cuda:
img = img.cuda()
# measure elapsed time
junc_pred, heatmap_pred, adj_mtx_pred = model(img)
# visualize eval
img = img.cpu().numpy()
junctions_pred = junc_pred.cpu().numpy()
adj_mtx = adj_mtx_pred.cpu().numpy()
img_with_junc = draw_jucntions(img, junctions_pred)
img_with_junc = img_with_junc[0].numpy()[:, ::-1, :, :]
lines_pred, score_pred = graph2line(junctions_pred, adj_mtx)
vis_line_pred = draw_lines(img_with_junc, lines_pred, score_pred)[0]
cv2.imshow("result", vis_line_pred)
return self