Spaces:
Runtime error
Runtime error
| from pathlib import Path | |
| import argparse | |
| import cv2 | |
| import matplotlib.cm as cm | |
| import torch | |
| import numpy as np | |
| from utils.nnmatching import NNMatching | |
| from utils.misc import ( | |
| AverageTimer, | |
| VideoStreamer, | |
| make_matching_plot_fast, | |
| frame2tensor, | |
| ) | |
| torch.set_grad_enabled(False) | |
| def compute_essential(matched_kp1, matched_kp2, K): | |
| pts1 = cv2.undistortPoints( | |
| matched_kp1, | |
| cameraMatrix=K, | |
| distCoeffs=(-0.117918271740560, 0.075246403574314, 0, 0), | |
| ) | |
| pts2 = cv2.undistortPoints( | |
| matched_kp2, | |
| cameraMatrix=K, | |
| distCoeffs=(-0.117918271740560, 0.075246403574314, 0, 0), | |
| ) | |
| K_1 = np.eye(3) | |
| # Estimate the homography between the matches using RANSAC | |
| ransac_model, ransac_inliers = cv2.findEssentialMat( | |
| pts1, pts2, K_1, method=cv2.RANSAC, prob=0.999, threshold=0.001, maxIters=10000 | |
| ) | |
| if ransac_inliers is None or ransac_model.shape != (3, 3): | |
| ransac_inliers = np.array([]) | |
| ransac_model = None | |
| return ransac_model, ransac_inliers, pts1, pts2 | |
| sizer = (960, 640) | |
| focallength_x = 4.504986436499113e03 / (6744 / sizer[0]) | |
| focallength_y = 4.513311442889859e03 / (4502 / sizer[1]) | |
| K = np.eye(3) | |
| K[0, 0] = focallength_x | |
| K[1, 1] = focallength_y | |
| K[0, 2] = 3.363322177533149e03 / (6744 / sizer[0]) # * 0.5 | |
| K[1, 2] = 2.291824660547715e03 / (4502 / sizer[1]) # * 0.5 | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser( | |
| description="DarkFeat demo", | |
| formatter_class=argparse.ArgumentDefaultsHelpFormatter, | |
| ) | |
| parser.add_argument("--input", type=str, help="path to an image directory") | |
| parser.add_argument( | |
| "--output_dir", | |
| type=str, | |
| default=None, | |
| help="Directory where to write output frames (If None, no output)", | |
| ) | |
| parser.add_argument( | |
| "--image_glob", | |
| type=str, | |
| nargs="+", | |
| default=["*.ARW"], | |
| help="Glob if a directory of images is specified", | |
| ) | |
| parser.add_argument( | |
| "--resize", | |
| type=int, | |
| nargs="+", | |
| default=[640, 480], | |
| help="Resize the input image before running inference. If two numbers, " | |
| "resize to the exact dimensions, if one number, resize the max " | |
| "dimension, if -1, do not resize", | |
| ) | |
| parser.add_argument( | |
| "--force_cpu", action="store_true", help="Force pytorch to run in CPU mode." | |
| ) | |
| parser.add_argument("--model_path", type=str, help="Path to the pretrained model") | |
| opt = parser.parse_args() | |
| print(opt) | |
| assert len(opt.resize) == 2 | |
| print("Will resize to {}x{} (WxH)".format(opt.resize[0], opt.resize[1])) | |
| device = "cuda" if torch.cuda.is_available() and not opt.force_cpu else "cpu" | |
| print('Running inference on device "{}"'.format(device)) | |
| matching = NNMatching(opt.model_path).eval().to(device) | |
| keys = ["keypoints", "scores", "descriptors"] | |
| vs = VideoStreamer(opt.input, opt.resize, opt.image_glob) | |
| frame, ret = vs.next_frame() | |
| assert ret, "Error when reading the first frame (try different --input?)" | |
| frame_tensor = frame2tensor(frame, device) | |
| last_data = matching.darkfeat({"image": frame_tensor}) | |
| last_data = {k + "0": [last_data[k]] for k in keys} | |
| last_data["image0"] = frame_tensor | |
| last_frame = frame | |
| last_image_id = 0 | |
| if opt.output_dir is not None: | |
| print("==> Will write outputs to {}".format(opt.output_dir)) | |
| Path(opt.output_dir).mkdir(exist_ok=True) | |
| timer = AverageTimer() | |
| while True: | |
| frame, ret = vs.next_frame() | |
| if not ret: | |
| print("Finished demo_darkfeat.py") | |
| break | |
| timer.update("data") | |
| stem0, stem1 = last_image_id, vs.i - 1 | |
| frame_tensor = frame2tensor(frame, device) | |
| pred = matching({**last_data, "image1": frame_tensor}) | |
| kpts0 = last_data["keypoints0"][0].cpu().numpy() | |
| kpts1 = pred["keypoints1"][0].cpu().numpy() | |
| matches = pred["matches0"][0].cpu().numpy() | |
| confidence = pred["matching_scores0"][0].cpu().numpy() | |
| timer.update("forward") | |
| valid = matches > -1 | |
| mkpts0 = kpts0[valid] | |
| mkpts1 = kpts1[matches[valid]] | |
| E, inliers, pts1, pts2 = compute_essential(mkpts0, mkpts1, K) | |
| color = cm.jet( | |
| np.clip(confidence[valid][inliers[:, 0].astype("bool")] * 2 - 1, -1, 1) | |
| ) | |
| text = ["DarkFeat", "Matches: {}".format(inliers.sum())] | |
| out = make_matching_plot_fast( | |
| last_frame, | |
| frame, | |
| mkpts0[inliers[:, 0].astype("bool")], | |
| mkpts1[inliers[:, 0].astype("bool")], | |
| color, | |
| text, | |
| path=None, | |
| small_text=" ", | |
| ) | |
| if opt.output_dir is not None: | |
| stem = "matches_{:06}_{:06}".format(stem0, stem1) | |
| out_file = str(Path(opt.output_dir, stem + ".png")) | |
| print("Writing image to {}".format(out_file)) | |
| cv2.imwrite(out_file, out) | |