Upload infer.py

infer.py

@@ -0,0 +1,164 @@
+import argparse
+import numpy as np
+import cv2
+import axengine as axe
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", type=str, required=True, help="Path to the axmodel")
+    parser.add_argument("--img1", type=str, required=True, help="Path to the first image")
+    parser.add_argument("--img2", type=str, required=True, help="Path to the second image")
+    parser.add_argument("--output", type=str, default="matches.jpg", help="The output image directory")
+    parser.add_argument("--threshold", type=float, default=0.005, help="The keypoint threshold")
+    parser.add_argument("--max_points", type=int, default=100, help="The max num for keypoints")
+    return parser.parse_args()
+
+def preprocess_image(path: str, h: int, w: int):
+    img = cv2.imread(path)
+    raw_h, raw_w = img.shape[:2]
+
+    if (raw_h, raw_w) != (h, w):
+        img = cv2.resize(img, (w, h))
+        scale_h = raw_h / h
+        scale_w = raw_w / w
+    else:
+        scale_h = 1.0
+        scale_w = 1.0
+
+    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    img_tensor = img_gray.astype(np.float32) / 255.0
+    img_tensor = img_tensor[None, None, :, :]  # -> (1, 1, H, W)
+
+    return img_tensor, img, (scale_h, scale_w)
+
+def get_keypoints(score_map, threshold):
+    row, col = np.where(score_map > threshold)  # y, x
+    if len(row) == 0:
+        return np.zeros((0, 2), dtype=np.float32), np.zeros((0,), dtype=np.float32)
+    
+    scores = score_map[row, col]
+    keypoints = np.stack([col, row], axis=1).astype(np.float32)
+    return keypoints, scores
+
+def get_descriptors(kp, desc_map):
+    if len(kp) == 0:
+        return np.zeros((0, 256), dtype=np.float32)
+    
+    c, h, w = desc_map.shape
+    x = kp[:, 0] / 8.0
+    y = kp[:, 1] / 8.0
+
+    x0 = np.floor(x).astype(np.int32)
+    x1 = x0 + 1
+    y0 = np.floor(y).astype(np.int32)
+    y1 = y0 + 1
+
+    x0 = np.clip(x0, 0, w - 1)
+    x1 = np.clip(x1, 0, w - 1)
+    y0 = np.clip(y0, 0, h - 1)
+    y1 = np.clip(y1, 0, h - 1)
+
+    wa = (x1 - x) * (y1 - y)
+    wb = (x1 - x) * (y - y0)
+    wc = (x - x0) * (y1 - y)
+    wd = (x - x0) * (y - y0)
+
+    wa = wa[None, :]
+    wb = wb[None, :]
+    wc = wc[None, :]
+    wd = wd[None, :]
+
+    Q_tl = desc_map[:, y0, x0]
+    Q_bl = desc_map[:, y1, x0]
+    Q_tr = desc_map[:, y0, x1]
+    Q_br = desc_map[:, y1, x1]
+
+    sampled = (Q_tl * wa + Q_bl * wb + Q_tr * wc + Q_br * wd)
+    descriptors = sampled.T
+    norm = np.linalg.norm(descriptors, axis=1, keepdims=True)
+    descriptors = descriptors / (norm + 1e-6)
+    
+    return descriptors.astype(np.float32)
+
+def infer(model: str, img1_path: str, img2_path: str, output: str, threshold: float, max_points: int):
+    session = axe.InferenceSession(model)
+
+    # superpoint only have one input
+    input_name = session.get_inputs()[0].name  # get model input node name
+    input_shape = session.get_inputs()[0].shape  # get model input shape (1, 1, H, W)
+    target_h, target_w = input_shape[2], input_shape[3]
+    print(f"Inference resolution: {target_w}x{target_h}")
+
+    # preprocess images
+    input_tensor1, img1, scale1 = preprocess_image(img1_path, target_h, target_w)
+    input_tensor2, img2, scale2 = preprocess_image(img2_path, target_h, target_w)
+
+    res1 = session.run(None, {input_name: input_tensor1})
+    res2 = session.run(None, {input_name: input_tensor2})
+
+    # [1,480,640], [1,256,60,80]
+    score_map1, desc1_map = res1[0], res1[1]
+    score_map2, desc2_map = res2[0], res2[1]
+
+    keypoints1, scores1 = get_keypoints(score_map1[0], threshold)
+    keypoints2, scores2 = get_keypoints(score_map2[0], threshold)
+
+    print(f"Found {len(keypoints1)} keypoints in image 1")
+    print(f"Found {len(keypoints2)} keypoints in image 2")
+
+    if len(keypoints1) > max_points:
+        idx = np.argsort(scores1)[::-1][:max_points]
+        keypoints1 = keypoints1[idx]
+        scores1 = scores1[idx]
+    if len(keypoints2) > max_points:
+        idx = np.argsort(scores2)[::-1][:max_points]
+        keypoints2 = keypoints2[idx]
+        scores2 = scores2[idx]
+
+    desc1 = get_descriptors(keypoints1, desc1_map[0])
+    desc2 = get_descriptors(keypoints2, desc2_map[0])
+
+    bf = cv2.BFMatcher(cv2.NORM_L2, crossCheck=True)
+    matches = bf.match(desc1, desc2)
+    matches = sorted(matches, key=lambda x: x.distance)
+
+    points1 = [cv2.KeyPoint(x, y, 1) for x, y in keypoints1]
+    points2 = [cv2.KeyPoint(x, y, 1) for x, y in keypoints2]
+
+    match_img = cv2.drawMatches(
+        img1, points1,
+        img2, points2,
+        matches, None,
+        flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS,
+        matchColor=(0, 255, 0)
+    )
+
+    # if len(matches) > 4:
+    #     pts1 = np.float32([keypoints1[m.queryIdx] for m in matches]).reshape(-1, 1, 2)
+    #     pts2 = np.float32([keypoints2[m.trainIdx] for m in matches]).reshape(-1, 1, 2)
+
+    #     H, mask = cv2.findHomography(pts1, pts2, cv2.RANSAC, 3.0)
+
+    #     if mask is not None:
+    #         matches_mask = mask.ravel().tolist()
+    #         inlier_matches = [m for i, m in enumerate(matches) if matches_mask[i]]
+    #         print(f"Inliers: {len(inlier_matches)} / {len(matches)}")
+
+    #         inlier_img = cv2.drawMatches(
+    #             img1, points1,
+    #             img2, points2,
+    #             inlier_matches, None,
+    #             flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS,
+    #             matchColor=(0, 255, 0)
+    #         )
+    #         cv2.imwrite("inliers_" + output, inlier_img)
+
+    cv2.imwrite(output, match_img)
+    print(f"Result saved to {output}")
+
+def main():
+    args = parse_args()
+    infer(args.model, args.img1, args.img2, args.output, args.threshold, args.max_points)
+
+if __name__ == '__main__':
+    main()