third_party/SGMNet/demo/demo.py

import cv2
import yaml
import numpy as np
import os
import sys

ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, ROOT_DIR)
from components import load_component
from utils import evaluation_utils

import argparse

parser = argparse.ArgumentParser()
parser.add_argument(
    "--config_path",
    type=str,
    default="configs/sgm_config.yaml",
    help="number of processes.",
)
parser.add_argument(
    "--img1_path", type=str, default="demo_1.jpg", help="number of processes."
)
parser.add_argument(
    "--img2_path", type=str, default="demo_2.jpg", help="number of processes."
)


args = parser.parse_args()

if __name__ == "__main__":
    with open(args.config_path, "r") as f:
        demo_config = yaml.load(f)

    extractor = load_component(
        "extractor", demo_config["extractor"]["name"], demo_config["extractor"]
    )

    img1, img2 = cv2.imread(args.img1_path), cv2.imread(args.img2_path)
    size1, size2 = np.flip(np.asarray(img1.shape[:2])), np.flip(
        np.asarray(img2.shape[:2])
    )
    kpt1, desc1 = extractor.run(args.img1_path)
    kpt2, desc2 = extractor.run(args.img2_path)

    matcher = load_component(
        "matcher", demo_config["matcher"]["name"], demo_config["matcher"]
    )
    test_data = {
        "x1": kpt1,
        "x2": kpt2,
        "desc1": desc1,
        "desc2": desc2,
        "size1": size1,
        "size2": size2,
    }
    corr1, corr2 = matcher.run(test_data)

    # draw points
    dis_points_1 = evaluation_utils.draw_points(img1, kpt1)
    dis_points_2 = evaluation_utils.draw_points(img2, kpt2)

    # visualize match
    display = evaluation_utils.draw_match(dis_points_1, dis_points_2, corr1, corr2)
    cv2.imwrite("match.png", display)