def eval_model(model, parsed_batch, DES_THRSH, COO_THRSH):
"""
only support one bach size cause function ditance_matrix_vector
"""
im1_data, im1_info, homo12, im2_data, im2_info, homo21, im1_raw, im2_raw = (
parsed_batch
)
# predict key points in each image and extract descripotor from each patch
scale1, kp1, des1,_,_ = model.inference(im1_data, im1_info, im1_raw)
scale2, kp2, des2,_,_ = model.inference(im2_data, im2_info, im2_raw)
kp1w = ptCltoCr(kp1, homo12, scale2, right_imorint=None, clamp=False)[0]
_, _, maxh, maxw = im2_data.size()
visible = kp1w[:, 2].lt(maxw) * kp1w[:, 1].lt(maxh)
TPNN, PNN = nearest_neighbor_match_score(des1, des2, kp1w, kp2, visible, COO_THRSH)
TPNNT, PNNT = nearest_neighbor_threshold_match_score(
des1, des2, kp1w, kp2, visible, DES_THRSH, COO_THRSH
)
TPNNDR, PNNDR = nearest_neighbor_distance_ratio_match_score(
des1, des2, kp1w, kp2, visible, COO_THRSH
)
return TPNN, PNN, TPNNT, PNNT, TPNNDR, PNNDR
def pair(
left_imtopk,
left_topkvalue,
left_imscale,
left_imorint,
left_iminfo,
left_imraw,
homolr,
right_imscale,
right_imorint,
right_iminfo,
right_imraw,
PSIZE,
):
"""
generate patch pair based on left topk_mask
"""
left_imC = left_imtopk.nonzero() # (B*topk, 4)
left_imS = left_imscale.masked_select(
left_imtopk) # (B*topk) have grad to detect network
if left_imorint is not None:
left_cos, left_sim = left_imorint.squeeze().chunk(chunks=2, dim=-1)
left_cos = left_cos.masked_select(left_imtopk) # (B*topk)
left_sim = left_sim.masked_select(left_imtopk) # (B*topk)
left_imO = torch.cat((left_cos.unsqueeze(-1), left_sim.unsqueeze(-1)),
dim=-1)
else:
left_imO = None
left_imP = clip_patch(left_imC,
left_imS,
left_imO,
left_iminfo,
left_imraw,
PSIZE=PSIZE) # (numkp, 1, 32, 32)
right_imC, right_imS, right_imO = ptCltoCr(left_imC, homolr, right_imscale,
right_imorint)
right_imP = clip_patch(right_imC,
right_imS,
right_imO,
right_iminfo,
right_imraw,
PSIZE=PSIZE) # (numkp, 1, 32, 32)
left_impair = torch.cat((left_imP, right_imP), 1) # (numkp, 2, 32, 32)
# import pdb
# pdb.set_trace()
return left_impair, left_imC, right_imC
def eval_model(det, des, parsed_batch, DES_THRSH, COO_THRSH, TOPK, PSIZE):
"""
only support one bach size cause function ditance_matrix_vector
"""
im1_data, im1_info, homo12, im2_data, im2_info, homo21, im1_raw, im2_raw = (
parsed_batch)
# predict key points in each image and extract descripotor from each patch
def inference(im_data, im_info, im_raw):
score_maps, orint_maps = det(im_data)
im_rawsc, im_scale, im_orint = handle_det_out(score_maps, orint_maps,
det.scale_list,
det.score_com_strength,
det.scale_com_strength)
im_score = det.process(im_rawsc)[0]
im_topk = topk_map(im_score, TOPK)
kpts = im_topk.nonzero() # (B*topk, 4)
cos, sim = im_orint.squeeze().chunk(chunks=2, dim=-1)
cos = cos.masked_select(im_topk) # (B*topk)
sim = sim.masked_select(im_topk) # (B*topk)
im_orint = torch.cat((cos.unsqueeze(-1), sim.unsqueeze(-1)), dim=-1)
im_patches = clip_patch(
kpts,
im_scale.masked_select(im_topk),
im_orint,
im_info,
im_raw,
PSIZE=PSIZE,
) # (numkp, 1, 32, 32)
im_patches = des.input_norm(im_patches)
im_des = des(im_patches)
return im_scale, kpts, im_des
scale1, kp1, des1 = inference(im1_data, im1_info, im1_raw)
scale2, kp2, des2 = inference(im2_data, im2_info, im2_raw)
kp1w = ptCltoCr(kp1, homo12, scale2, right_imorint=None, clamp=False)[0]
_, _, maxh, maxw = im2_data.size()
visible = kp1w[:, 2].lt(maxw) * kp1w[:, 1].lt(maxh)
TPNN, PNN = nearest_neighbor_match_score(des1, des2, kp1w, kp2, visible,
COO_THRSH)
TPNNT, PNNT = nearest_neighbor_threshold_match_score(
des1, des2, kp1w, kp2, visible, DES_THRSH, COO_THRSH)
TPNNDR, PNNDR = nearest_neighbor_distance_ratio_match_score(
des1, des2, kp1w, kp2, visible, COO_THRSH)
return TPNN, PNN, TPNNT, PNNT, TPNNDR, PNNDR
def eval_model(model, parsed_batch, DES_THRSH, COO_THRSH):
"""
only support one bach size cause function ditance_matrix_vector
"""
im1_data, im1_info, homo12, im2_data, im2_info, homo21, im1_raw, im2_raw = (
parsed_batch)
# predict key points in each image and extract descripotor from each patch
scale1, kp1, des1 = model.inference(im1_data, im1_info, im1_raw)
scale2, kp2, des2 = model.inference(im2_data, im2_info, im2_raw)
kp1w = ptCltoCr(kp1, homo12, scale2, right_imorint=None, clamp=False)[0]
predict_label, nn_kp2 = nearest_neighbor_distance_ratio_match(
des1, des2, kp2, 0.8)
idx = predict_label.nonzero().view(-1)
kp1 = kp1w
mkp1 = kp1.index_select(dim=0,
index=idx.long()) # predict match keypoints in I1
mkp2 = nn_kp2.index_select(
dim=0, index=idx.long()) # predict match keypoints in I2
def to_cv2_kp(kp):
# kp is like [batch_idx, y, x, channel]
return cv2.KeyPoint(kp[2], kp[1], 0)
def to_cv2_dmatch(m):
return cv2.DMatch(m, m, m, m)
# DMatch.distance
# DMatch.trainIdx
# DMatch.queryIdx
# DMatch.imgIdx
def reverse_img(img):
"""
reverse image from tensor to cv2 format
:param img: tensor
:return: RBG image
"""
img = img.permute(0, 2, 3, 1)[0].cpu().detach().numpy()
img = (img * 255).astype(np.uint8) # change to opencv format
img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB) # gray to rgb
return img
img1 = im1_data
img2 = im2_data
img1, img2 = reverse_img(img1), reverse_img(img2)
keypoints1 = list(map(to_cv2_kp, mkp1))
keypoints2 = list(map(to_cv2_kp, mkp2))
DMatch = list(map(to_cv2_dmatch, np.arange(0, len(keypoints1))))
# matches1to2 Matches from the first image to the second one, which means that
# keypoints1[i] has a corresponding point in keypoints2[matches[i]] .
outImg = cv2.drawMatches(img1, keypoints1, img2, keypoints2, DMatch, None)
# cv2.imwrite("outImg.png", outImg)
name = time.time()
name = str(name)
cv2.imwrite(name + 'match.png', outImg)
img_kps1 = np.copy(img1)
img_kps1 = cv2.drawKeypoints(img1, keypoints1, img_kps1, flags=None)
cv2.imwrite("img_kps3.png", img_kps1)
img_kps2 = np.copy(img2)
img_kps2 = cv2.drawKeypoints(img2, keypoints2, img_kps2, flags=None)
cv2.imwrite("img_kps4.png", img_kps2)
_, _, maxh, maxw = im2_data.size()
visible = kp1w[:, 2].lt(maxw) * kp1w[:, 1].lt(maxh)
TPNN, PNN = nearest_neighbor_match_score(des1, des2, kp1w, kp2, visible,
COO_THRSH)
TPNNT, PNNT = nearest_neighbor_threshold_match_score(
des1, des2, kp1w, kp2, visible, DES_THRSH, COO_THRSH)
TPNNDR, PNNDR = nearest_neighbor_distance_ratio_match_score(
des1, des2, kp1w, kp2, visible, COO_THRSH)
print('TPNN, PNN, TPNNT, PNNT, TPNNDR, PNNDR are as followings:')
print(TPNN, PNN, TPNNT, PNNT, TPNNDR, PNNDR)
return TPNN, PNN, TPNNT, PNNT, TPNNDR, PNNDR
