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