File size: 4,887 Bytes
0a82b18 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 |
from matching import get_matcher, available_models, get_default_device
from pathlib import Path
from argparse import ArgumentParser
import cv2
import time
from tqdm.auto import tqdm
import torch
import numpy as np
def parse_args():
parser = ArgumentParser()
parser.add_argument(
"--task", type=str, default="benchmark", help="run benchmark or unit tests"
)
parser.add_argument(
"--matcher",
type=str,
nargs="+",
default="all",
help="which model or list of models to benchmark",
)
parser.add_argument(
"--img-size",
type=int,
default=512,
help="image size to run matching on (resized to square)",
)
parser.add_argument(
"--device",
type=str,
default=get_default_device(),
help="Device to run benchmark on",
)
parser.add_argument(
"--num-iters",
type=int,
default=5,
help="number of interations to run benchmark and average over",
)
args = parser.parse_args()
if args.device == "cuda":
assert (
torch.cuda.is_available()
), "Chosen cuda as device but cuda unavailable! Try another device (cpu)"
if args.matcher == "all":
args.matcher = available_models
return args
def get_img_pairs():
asset_dir = Path(__file__).parent.joinpath("assets/example_pairs")
pairs = [
list(pair.iterdir()) for pair in list(asset_dir.iterdir()) if pair.is_dir()
]
return pairs
def test_H_est(matcher, img_size=512):
"""Given a matcher, compute a homography of two images with known ground
truth and its error. The error for sift-lg is 0.002 for img_size=500. So it
should roughly be below 0.01."""
img0_path = "assets/example_test/warped.jpg"
img1_path = "assets/example_test/original.jpg"
ground_truth = np.array(
[[0.1500, 0.3500], [0.9500, 0.1500], [0.9000, 0.7000], [0.2500, 0.7000]]
)
image0 = matcher.load_image(img0_path, resize=img_size)
image1 = matcher.load_image(img1_path, resize=img_size)
result = matcher(image0, image1)
pred_homog = np.array(
[[0, 0], [img_size, 0], [img_size, img_size], [0, img_size]], dtype=np.float32
)
pred_homog = np.reshape(pred_homog, (4, 1, 2))
prediction = cv2.perspectiveTransform(pred_homog, result["H"])[:, 0] / img_size
max_error = np.abs(ground_truth - prediction).max()
return max_error
def test(matcher, img_sizes=[512, 256], error_thresh=0.05):
passing = True
for img_size in img_sizes:
error = test_H_est(matcher, img_size=img_size)
if error > error_thresh:
passing = False
raise RuntimeError(
f"Large homography error in matcher (size={img_size} px): {error}"
)
return passing, error
def benchmark(matcher, num_iters=1, img_size=512):
runtime = []
for _ in range(num_iters):
for pair in get_img_pairs():
img0 = matcher.load_image(pair[0], resize=img_size)
img1 = matcher.load_image(pair[1], resize=img_size)
start = time.time()
_ = matcher(img0, img1)
duration = time.time() - start
runtime.append(duration)
return runtime, np.mean(runtime)
def main(args):
print(args)
if args.task == "benchmark":
with open("runtime_results.txt", "w") as f:
for model in tqdm(args.matcher):
try:
matcher = get_matcher(model, device=args.device)
runtimes, avg_runtime = benchmark(
matcher, num_iters=args.num_iters, img_size=args.img_size
)
runtime_str = f"{model: <15} OK {avg_runtime=:.3f}"
f.write(runtime_str + "\n")
tqdm.write(runtime_str)
except Exception as e:
tqdm.write(f"{model: <15} NOT OK - exception: {e}")
elif args.task == "test":
with open("test_results.txt", "w") as f:
test_str = "Matcher, Passing Tests, Error (px)"
f.write(test_str + "\n" + "-" * 40 + "\n")
tqdm.write(test_str)
for model in tqdm(args.matcher):
try:
matcher = get_matcher(model, device=args.device)
passing, error_val = test(matcher)
test_str = f"{model}, {passing}, {error_val}"
f.write(test_str + "\n")
tqdm.write(test_str)
except Exception as e:
f.write(f"Error with {model}: {e}")
tqdm.write(f"Error with {model}: {e}")
if __name__ == "__main__":
args = parse_args()
import warnings
warnings.filterwarnings("ignore")
print(f"Running with args: {args}")
main(args)
|