def get_points_desc(fe, sp, device, img, thresh):
"""
Returns points in (x,y) format alongside with descriptors
"""
if fe is not None:
pts_2, desc_2, heatmap_2 = fe.run(img.astype('float32').squeeze() / 255.)
else:
with torch.no_grad():
pts_2, desc_2 = sp.to(device).points_desc(torch.from_numpy(img).to(device), threshold=thresh)
pts_2 = pts_2.T
desc_2 = desc_2[0].T.cpu().detach().numpy()
pts_2 = numpy.concatenate([util.swap_rows(pts_2[:2]), pts_2[2, :][numpy.newaxis, :]])
return pts_2, desc_2
def draw_m(img1, img2, matches, points, points2):
from super_debug import draw_matches
img_pair = numpy.hstack([img1, img2]).astype(numpy.uint8)
img_pair = numpy.stack([img_pair, img_pair, img_pair], axis=2).squeeze()
draw_matches(0, img_pair, matches, util.swap_rows(points.cpu().numpy().T),
util.swap_rows(points2.cpu().numpy().T))
def main(sp, thresh, nn_thresh=0.8, draw=True):
matcher = PointTracker(max_length=2, nn_thresh=nn_thresh)
fe = None
dist_thresh = 0.8
root_fire = '/mnt/fileserver/shared/datasets/fundus/FIRE/'
gt_dir = os.path.join(root_fire, 'Ground Truth')
images_dir = os.path.join(root_fire, 'Images')
mask_dir = os.path.join(root_fire, 'Masks')
gt_files = [os.path.join(gt_dir, x) for x in os.listdir(gt_dir)]
nCasesTotal = 0
mask = imageio.imread(os.path.join(mask_dir, 'mask.png'))
mask = imageio.imread(os.path.join(mask_dir, 'feature_mask.png'))
scale = 0.15
mask = cv2.resize(mask, dsize=None, fx=scale, fy=scale)
if draw:
cv2.imshow('mask', mask)
cov_all = []
acc_all = []
rep_all = []
for gt_file in gt_files:
name = gt_file.split('control_points_')[1].split('_')[0]
img1_path = os.path.join(images_dir, name + '_1.jpg')
img2_path = os.path.join(images_dir, name + '_2.jpg')
img1 = imageio.imread(img1_path)
img2 = imageio.imread(img2_path)
img1 = cv2.resize(img1, dsize=None, fx=scale, fy=scale)
img2 = cv2.resize(img2, dsize=None, fx=scale, fy=scale)
img_h, img_w = img1.shape[:2]
arr = numpy.loadtxt(gt_file)
arr[:, (0, 2)] = arr[:, (0, 2)] * scale
arr[:, (1, 3)] = arr[:, (1, 3)] * scale
points1 = arr[:, :2]
points2 = arr[:, 2:]
H, H_inv = calculate_h(points1, points2)
# (x, y) to (row, col)
points1 = swap_rows(arr[:, :2].T).T
points2 = swap_rows(arr[:, 2:].T).T
#in_bounds, points1_projected = project_points2(torch.from_numpy(H), None, torch.from_numpy(points1), img_h, img_w)
#drawing.show_points(img1 / 256, points1.astype(numpy.int16), 'points1_img1', 0.5)
#drawing.show_points(img2 / 256, points2.astype(numpy.int16), 'points2_img2', 0.5)
#drawing.show_points(img2 / 256, points1_projected.cpu().long(), 'points1_img2', 0.5)
desc_1, pts_1, desc_2, pts_2 = extract(device, fe, img1, img2, mask, sp, thresh)
### debug
#img11 = drawing.draw_points(pts_1, img1.copy(), False)
#img12 = drawing.draw_points(pts_2, img2.copy(), False)
#in_bounds, points1_projected = project_points2(torch.from_numpy(H),
# None,
# torch.from_numpy(pts_1[:2].T),
# img_h,
# img_w)
#img22 = drawing.draw_points(points1_projected.numpy().T, img2.copy(), False)
#cv2.imshow('points11_img1', img11)
#cv2.imshow('points22_img2', img12)
#cv2.imshow('points11_img2', img22)
#cv2.waitKey(0)
nCasesTotal += 1
matches = matcher.nn_match_two_way(desc_1, desc_2, nn_thresh=dist_thresh)
nMatches = matches.shape[1]
pt1 = pts_1[:2, matches[0, :].astype('int32')]
pt2 = pts_2[:2, matches[1, :].astype('int32')]
in_bounds, pt2_p = project_points2(torch.from_numpy(H), None, torch.from_numpy(pt1.T),
img_h, img_w)
nGoodMatches = 0
pt2_p = pt2_p.T
dx = pt2_p[0, :] - pt2[0, :]
dy = pt2_p[1, :] - pt2[1, :]
err = numpy.sqrt(dx * dx + dy * dy)
rep = 0
for i in range(nMatches):
if err[i] <= GOOD_MATCH_THRESHOLD:
nGoodMatches += 1
matches[2, i] = 0
else:
matches[2, i] = err[i]
replication = replication_ratio(pt2_p.cpu().numpy(), pts_2, GOOD_MATCH_THRESHOLD)
accuracy = float(nGoodMatches) / float(nMatches) if nMatches else 0.0
rep_all.append(replication)
acc_all.append(accuracy)
print('accuracy:', accuracy)
print('replication:', replication)
if draw:
img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2RGB)
from super_debug import draw_desc_interpolate
draw_desc_interpolate(img1, img2, swap_rows(pts_1[:2]).T, swap_rows(pts_2[:2]).T, desc_1.T, desc_2.T, 0)
img_output = draw_matches(matches, pts_1, pts_2, img1, img2)
# img_output[:, img_output.shape[1] // 2:] = drawing.draw_points(pt2_p.numpy(),
# img_output[:, img_output.shape[1] // 2:], iscolor=False)
cv2.imshow('matches', img_output)
coverage1, coverage_mask, intersect = coverage(img1[:,:,-1], mask, matches, pt1)
if draw:
cv2.imshow('coverage', coverage_mask)
cov_all.append(coverage1)
if draw:
cv2.imshow('itersect', intersect)
print('converage:', coverage1)
if draw:
key = cv2.waitKey(0)
if key == ord('q'):
break
sys.stdout.flush()
acc = numpy.mean(acc_all)
cov = numpy.mean(cov_all)
rep = numpy.mean(rep_all)
print('average accuracy:', acc)
print('average coverage:', cov)
print('average replication:', rep)
print('harmonic mean:', harmonic(acc, cov, rep))
def points_desc(self, x, threshold=0.5):
"""
Method computes points and descriptors from torch.Tensor of images
:param x: torch.Tensor
:param threshold: float
:return:
"""
prob, desc = self.forward_expand(x.float())
if prob.shape[1] % 2 == 0:
prob = prob[:, 1]
else:
prob = prob[:, 0]
if self.nms:
prob_n = prob
if len(prob_n.shape) == 3:
prob_n = prob.unsqueeze(0)
if isinstance(self.nms, nonmaximum.HeatmapNMS):
prob = self.nms(prob_n).squeeze(0)
desc_result = []
rows, cols = x.shape[-2], x.shape[-1]
coords_result = None
for i, heatmap in enumerate(prob):
coords = (heatmap > threshold).nonzero()
coords = coords.cpu()
prob_i = prob[i].cpu()
heat = prob[i][coords[:, 0], coords[:, 1]]
heat = heat.cpu()
if isinstance(self.nms, nonmaximum.CoordsNMS):
coords_swapped = util.swap_rows(coords.transpose(1, 0))
coords_swapped = coords_swapped.cpu()
xy_heat = numpy.vstack([coords_swapped, heat.detach()])
coords_new = self.nms(rows, cols, xy_heat)
if not isinstance(coords_new, torch.Tensor):
coords_new = torch.from_numpy(coords_new).long()
coords = util.swap_rows(coords_new[0:2]).transpose(1, 0)
with_idx = numpy.hstack([
coords, prob_i[coords[:, 0],
coords[:, 1]].detach().unsqueeze(1)
])
else:
with_idx = numpy.hstack([
coords, prob_i[coords[:, 0],
coords[:, 1]].detach().unsqueeze(1)
])
if coords_result is None:
coords_result = with_idx
else:
coords_result = numpy.vstack([coords_result,
with_idx]).astype(numpy.long)
desc_result.append(
util.descriptor_interpolate(desc[0],
rows,
cols,
coords,
align_corners=self.align_corners))
return coords_result, desc_result