PCGao/MatchAnything

1import cv22from pathlib import Path3from imcui.hloc import logger4from imcui.ui.utils import DEVICE, get_matcher_zoo, load_config5from imcui.api import ImageMatchingAPI6 7ROOT = Path(__file__).parents[1]8 9 10def test_all():11    config = load_config(ROOT / "config/config.yaml")12    img_path1 = ROOT / "tests/data/02928139_3448003521.jpg"13    img_path2 = ROOT / "tests/data/17295357_9106075285.jpg"14    image0 = cv2.imread(str(img_path1))[:, :, ::-1]  # RGB15    image1 = cv2.imread(str(img_path2))[:, :, ::-1]  # RGB16 17    matcher_zoo_restored = get_matcher_zoo(config["matcher_zoo"])18    for k, v in matcher_zoo_restored.items():19        if image0 is None or image1 is None:20            logger.error("Error: No images found! Please upload two images.")21        enable = config["matcher_zoo"][k].get("enable", True)22        skip_ci = config["matcher_zoo"][k].get("skip_ci", False)23        if enable and not skip_ci:24            logger.info(f"Testing {k} ...")25            api = ImageMatchingAPI(conf=v, device=DEVICE)26            pred = api(image0, image1)27            assert pred is not None28            log_path = ROOT / "experiments" / "all"29            log_path.mkdir(exist_ok=True, parents=True)30            api.visualize(log_path=log_path)31        else:32            logger.info(f"Skipping {k} ...")33 34 35def test_one():36    img_path1 = ROOT / "tests/data/02928139_3448003521.jpg"37    img_path2 = ROOT / "tests/data/17295357_9106075285.jpg"38 39    image0 = cv2.imread(str(img_path1))[:, :, ::-1]  # RGB40    image1 = cv2.imread(str(img_path2))[:, :, ::-1]  # RGB41    # sparse42    conf = {43        "feature": {44            "output": "feats-superpoint-n4096-rmax1600",45            "model": {46                "name": "superpoint",47                "nms_radius": 3,48                "max_keypoints": 4096,49                "keypoint_threshold": 0.005,50            },51            "preprocessing": {52                "grayscale": True,53                "force_resize": True,54                "resize_max": 1600,55                "width": 640,56                "height": 480,57                "dfactor": 8,58            },59        },60        "matcher": {61            "output": "matches-NN-mutual",62            "model": {63                "name": "nearest_neighbor",64                "do_mutual_check": True,65                "match_threshold": 0.2,66            },67        },68        "dense": False,69    }70    api = ImageMatchingAPI(conf=conf, device=DEVICE)71    pred = api(image0, image1)72    assert pred is not None73    log_path = ROOT / "experiments" / "one"74    log_path.mkdir(exist_ok=True, parents=True)75    api.visualize(log_path=log_path)76 77    # dense78    conf = {79        "matcher": {80            "output": "matches-loftr",81            "model": {82                "name": "loftr",83                "weights": "outdoor",84                "max_keypoints": 2000,85                "match_threshold": 0.2,86            },87            "preprocessing": {88                "grayscale": True,89                "resize_max": 1024,90                "dfactor": 8,91                "width": 640,92                "height": 480,93                "force_resize": True,94            },95            "max_error": 1,96            "cell_size": 1,97        },98        "dense": True,99    }100 101    api = ImageMatchingAPI(conf=conf, device=DEVICE)102    pred = api(image0, image1)103    assert pred is not None104    log_path = ROOT / "experiments" / "one"105    log_path.mkdir(exist_ok=True, parents=True)106    api.visualize(log_path=log_path)107 108 109if __name__ == "__main__":110    test_one()111    test_all()112