demo.py

import os
import sys
import torch
import numpy as np
import math
import cv2

from models.liftfeat_wrapper import LiftFeat,MODEL_PATH

import argparse

parser=argparse.ArgumentParser(description='HPatch dataset evaluation script')
parser.add_argument('--name',type=str,default='LiftFeat',help='experiment name')
parser.add_argument('--img1',type=str,default='./assert/ref.jpg',help='reference image path')
parser.add_argument('--img2',type=str,default='./assert/query.jpg',help='query image path')
parser.add_argument('--size',type=str,default=None,help='Resize images to w,h, None means disable resize')
parser.add_argument('--use_opencv_match',action='store_true',help='Enable OpenCV match function')
parser.add_argument('--gpu',type=str,default='0',help='GPU ID')
args=parser.parse_args()

os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu


def warp_corners_and_draw_matches(ref_points, dst_points, img1, img2):
    h, w = img1.shape[:2]
    # Calculate the Homography matrix
    H, mask = cv2.findHomography(ref_points, dst_points, cv2.USAC_MAGSAC, h*0.01+w*0.01, maxIters=1_000, confidence=0.999)
    mask = mask.flatten()

    # Get corners of the first image (image1)
    corners_img1 = np.array([[0, 0], [w-1, 0], [w-1, h-1], [0, h-1]], dtype=np.float32).reshape(-1, 1, 2)

    # Warp corners to the second image (image2) space
    warped_corners = cv2.perspectiveTransform(corners_img1, H)

    # Draw the warped corners in image2
    img2_with_corners = img2.copy()

    # Prepare keypoints and matches for drawMatches function
    keypoints1 = [cv2.KeyPoint(float(p[0]), float(p[1]), 5) for p in ref_points]
    keypoints2 = [cv2.KeyPoint(float(p[0]), float(p[1]), 5) for p in dst_points]
    matches = [cv2.DMatch(i,i,0) for i in range(len(mask)) if mask[i]]

    # Draw inlier matches
    img_matches = cv2.drawMatches(img1, keypoints1, img2_with_corners, keypoints2, matches, None,
                                  matchColor=(0, 255, 0), flags=2)

    return img_matches


def opencv_knn_match(descs1,descs2,kpts1,kpts2):
    bf = cv2.BFMatcher()
    
    matches = bf.knnMatch(descs1,descs2,k=2)
    
    good_matches = []
    for m, n in matches:
        if m.distance < 0.9 * n.distance:
            good_matches.append(m)
    
    mkpts1 = [];mkpts2 = []
    
    for m in good_matches:
        mkpt1=kpts1[m.queryIdx];mkpt2=kpts2[m.trainIdx]
        mkpts1.append(mkpt1);mkpts2.append(mkpt2)
        
    mkpts1 = np.array(mkpts1)
    mkpts2 = np.array(mkpts2)
    
    return mkpts1,mkpts2


if __name__=="__main__":
    if args.size:
        print(f'resize images to {args.size}')
        w=int(args.size.split(',')[0])
        h=int(args.size.split(',')[1])
        dst_size=(w,h)
    else:
        print(f'disable resize')
        
    if args.use_opencv_match:
        print(f'Use OpenCV knnMatch')
    else:
        print(f'Use original match function')
    
    liftfeat=LiftFeat(weight=MODEL_PATH,detect_threshold=0.05)
    
    img1=cv2.imread(args.img1)
    img2=cv2.imread(args.img2)
    
    if args.size:
        img1=cv2.resize(img1,dst_size)
        img2=cv2.resize(img2,dst_size)
    
    if args.use_opencv_match:
        data1 = liftfeat.extract(img1)
        data2 = liftfeat.extract(img2)
        kpts1,descs1=data1['keypoints'].cpu().numpy(),data1['descriptors'].cpu().numpy()
        kpts2,descs2=data2['keypoints'].cpu().numpy(),data2['descriptors'].cpu().numpy()
        
        mkpts1,mkpts2 = opencv_knn_match(descs1,descs2,kpts1,kpts2)
    else:
        mkpts1,mkpts2=liftfeat.match_liftfeat(img1,img2)
    
    
    canvas=warp_corners_and_draw_matches(mkpts1,mkpts2,img1,img2)
    
    import matplotlib.pyplot as plt
    plt.figure(figsize=[12,12])
    plt.imshow(canvas[...,::-1])
    
    plt.savefig(os.path.join(os.path.dirname(__file__),'match.jpg'), dpi=300, bbox_inches='tight')
    
    plt.show()