|
|
""" |
|
|
This script performs image matching using a specified matcher model. It processes pairs of input images, |
|
|
detects keypoints, matches them, and performs RANSAC to find inliers. The results, including visualizations |
|
|
and metadata, are saved to the specified output directory. |
|
|
""" |
|
|
|
|
|
import torch |
|
|
import argparse |
|
|
from pathlib import Path |
|
|
|
|
|
from matching.utils import get_image_pairs_paths, get_default_device |
|
|
from matching import get_matcher, available_models |
|
|
from matching.viz import plot_matches |
|
|
from typing import Union |
|
|
|
|
|
|
|
|
def main(args): |
|
|
image_size = [args.im_size, args.im_size] |
|
|
args.out_dir.mkdir(exist_ok=True, parents=True) |
|
|
|
|
|
|
|
|
matcher = get_matcher(args.matcher, device=args.device, max_num_keypoints=args.n_kpts) |
|
|
|
|
|
pairs_of_paths = get_image_pairs_paths(args.input) |
|
|
for i, (img0_path, img1_path) in enumerate(pairs_of_paths): |
|
|
|
|
|
image0 = matcher.load_image(img0_path, resize=image_size) |
|
|
image1 = matcher.load_image(img1_path, resize=image_size) |
|
|
result = matcher(image0, image1) |
|
|
|
|
|
out_str = f"Paths: {str(img0_path), str(img1_path)}. Found {result['num_inliers']} inliers after RANSAC. " |
|
|
|
|
|
if not args.no_viz: |
|
|
viz_path = args.out_dir / f"output_{i}_matches.jpg" |
|
|
plot_matches(image0, image1, result, save_path=viz_path) |
|
|
out_str += f"Viz saved in {viz_path}. " |
|
|
|
|
|
result["img0_path"] = img0_path |
|
|
result["img1_path"] = img1_path |
|
|
result["matcher"] = args.matcher |
|
|
result["n_kpts"] = args.n_kpts |
|
|
result["im_size"] = args.im_size |
|
|
|
|
|
dict_path = args.out_dir / f"output_{i}_result.torch" |
|
|
torch.save(result, dict_path) |
|
|
out_str += f"Output saved in {dict_path}" |
|
|
|
|
|
print(out_str) |
|
|
|
|
|
|
|
|
def parse_args(): |
|
|
parser = argparse.ArgumentParser( |
|
|
description="Image Matching Models", |
|
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter, |
|
|
) |
|
|
|
|
|
parser.add_argument( |
|
|
"--matcher", |
|
|
type=str, |
|
|
default="sift-lg", |
|
|
choices=available_models, |
|
|
help="choose your matcher", |
|
|
) |
|
|
|
|
|
|
|
|
parser.add_argument("--im_size", type=int, default=512, help="resize img to im_size x im_size") |
|
|
parser.add_argument("--n_kpts", type=int, default=2048, help="max num keypoints") |
|
|
parser.add_argument("--device", type=str, default=get_default_device(), choices=["cpu", "cuda"]) |
|
|
parser.add_argument("--no_viz", action="store_true", help="avoid saving visualizations") |
|
|
|
|
|
parser.add_argument( |
|
|
"--input", |
|
|
type=Path, |
|
|
nargs="+", |
|
|
default=[Path("assets/example_pairs")], |
|
|
help="path to either (1) two image paths or (2) dir with two images or (3) dir with dirs with image pairs or " |
|
|
"(4) txt file with two image paths per line", |
|
|
) |
|
|
parser.add_argument("--out_dir", type=Path, default=None, help="path where outputs are saved") |
|
|
|
|
|
args = parser.parse_args() |
|
|
|
|
|
if args.out_dir is None: |
|
|
args.out_dir = Path(f"outputs_{args.matcher}") |
|
|
|
|
|
return args |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
args = parse_args() |
|
|
main(args) |
|
|
|
