PCGao/MatchAnything
0
1# api.py2import warnings3from pathlib import Path4from typing import Any, Dict, Optional5 6import cv27import matplotlib.pyplot as plt8import numpy as np9import torch10 11from ..hloc import extract_features, logger, match_dense, match_features12from ..hloc.utils.viz import add_text, plot_keypoints13from ..ui.utils import filter_matches, get_feature_model, get_model14from ..ui.viz import display_matches, fig2im, plot_images15 16warnings.simplefilter("ignore")17 18 19class ImageMatchingAPI(torch.nn.Module):20 default_conf = {21 "ransac": {22 "enable": True,23 "estimator": "poselib",24 "geometry": "homography",25 "method": "RANSAC",26 "reproj_threshold": 3,27 "confidence": 0.9999,28 "max_iter": 10000,29 },30 }31 32 def __init__(33 self,34 conf: dict = {},35 device: str = "cpu",36 detect_threshold: float = 0.015,37 max_keypoints: int = 1024,38 match_threshold: float = 0.2,39 ) -> None:40 """41 Initializes an instance of the ImageMatchingAPI class.42 43 Args:44 conf (dict): A dictionary containing the configuration parameters.45 device (str, optional): The device to use for computation. Defaults to "cpu".46 detect_threshold (float, optional): The threshold for detecting keypoints. Defaults to 0.015.47 max_keypoints (int, optional): The maximum number of keypoints to extract. Defaults to 1024.48 match_threshold (float, optional): The threshold for matching keypoints. Defaults to 0.2.49 50 Returns:51 None52 """53 super().__init__()54 self.device = device55 self.conf = {**self.default_conf, **conf}56 self._updata_config(detect_threshold, max_keypoints, match_threshold)57 self._init_models()58 if device == "cuda":59 memory_allocated = torch.cuda.memory_allocated(device)60 memory_reserved = torch.cuda.memory_reserved(device)61 logger.info(f"GPU memory allocated: {memory_allocated / 1024**2:.3f} MB")62 logger.info(f"GPU memory reserved: {memory_reserved / 1024**2:.3f} MB")63 self.pred = None64 65 def parse_match_config(self, conf):66 if conf["dense"]:67 return {68 **conf,69 "matcher": match_dense.confs.get(conf["matcher"]["model"]["name"]),70 "dense": True,71 }72 else:73 return {74 **conf,75 "feature": extract_features.confs.get(conf["feature"]["model"]["name"]),76 "matcher": match_features.confs.get(conf["matcher"]["model"]["name"]),77 "dense": False,78 }79 80 def _updata_config(81 self,82 detect_threshold: float = 0.015,83 max_keypoints: int = 1024,84 match_threshold: float = 0.2,85 ):86 self.dense = self.conf["dense"]87 if self.conf["dense"]:88 try:89 self.conf["matcher"]["model"]["match_threshold"] = match_threshold90 except TypeError as e:91 logger.error(e)92 else:93 self.conf["feature"]["model"]["max_keypoints"] = max_keypoints94 self.conf["feature"]["model"]["keypoint_threshold"] = detect_threshold95 self.extract_conf = self.conf["feature"]96 97 self.match_conf = self.conf["matcher"]98 99 def _init_models(self):100 # initialize matcher101 self.matcher = get_model(self.match_conf)102 # initialize extractor103 if self.dense:104 self.extractor = None105 else:106 self.extractor = get_feature_model(self.conf["feature"])107 108 def _forward(self, img0, img1):109 if self.dense:110 pred = match_dense.match_images(111 self.matcher,112 img0,113 img1,114 self.match_conf["preprocessing"],115 device=self.device,116 )117 last_fixed = "{}".format( # noqa: F841118 self.match_conf["model"]["name"]119 )120 else:121 pred0 = extract_features.extract(122 self.extractor, img0, self.extract_conf["preprocessing"]123 )124 pred1 = extract_features.extract(125 self.extractor, img1, self.extract_conf["preprocessing"]126 )127 pred = match_features.match_images(self.matcher, pred0, pred1)128 return pred129 130 def _convert_pred(self, pred):131 ret = {132 k: v.cpu().detach()[0].numpy() if isinstance(v, torch.Tensor) else v133 for k, v in pred.items()134 }135 ret = {136 k: v[0].cpu().detach().numpy() if isinstance(v, list) else v137 for k, v in ret.items()138 }139 return ret140 141 @torch.inference_mode()142 def extract(self, img0: np.ndarray, **kwargs) -> Dict[str, np.ndarray]:143 """Extract features from a single image.144 145 Args:146 img0 (np.ndarray): image147 148 Returns:149 Dict[str, np.ndarray]: feature dict150 """151 152 # setting prams153 self.extractor.conf["max_keypoints"] = kwargs.get("max_keypoints", 512)154 self.extractor.conf["keypoint_threshold"] = kwargs.get(155 "keypoint_threshold", 0.0156 )157 158 pred = extract_features.extract(159 self.extractor, img0, self.extract_conf["preprocessing"]160 )161 pred = self._convert_pred(pred)162 # back to origin scale163 s0 = pred["original_size"] / pred["size"]164 pred["keypoints_orig"] = (165 match_features.scale_keypoints(pred["keypoints"] + 0.5, s0) - 0.5166 )167 # TODO: rotate back168 binarize = kwargs.get("binarize", False)169 if binarize:170 assert "descriptors" in pred171 pred["descriptors"] = (pred["descriptors"] > 0).astype(np.uint8)172 pred["descriptors"] = pred["descriptors"].T # N x DIM173 return pred174 175 @torch.inference_mode()176 def forward(177 self,178 img0: np.ndarray,179 img1: np.ndarray,180 ) -> Dict[str, np.ndarray]:181 """182 Forward pass of the image matching API.183 184 Args:185 img0: A 3D NumPy array of shape (H, W, C) representing the first image.186 Values are in the range [0, 1] and are in RGB mode.187 img1: A 3D NumPy array of shape (H, W, C) representing the second image.188 Values are in the range [0, 1] and are in RGB mode.189 190 Returns:191 A dictionary containing the following keys:192 - image0_orig: The original image 0.193 - image1_orig: The original image 1.194 - keypoints0_orig: The keypoints detected in image 0.195 - keypoints1_orig: The keypoints detected in image 1.196 - mkeypoints0_orig: The raw matches between image 0 and image 1.197 - mkeypoints1_orig: The raw matches between image 1 and image 0.198 - mmkeypoints0_orig: The RANSAC inliers in image 0.199 - mmkeypoints1_orig: The RANSAC inliers in image 1.200 - mconf: The confidence scores for the raw matches.201 - mmconf: The confidence scores for the RANSAC inliers.202 """203 # Take as input a pair of images (not a batch)204 assert isinstance(img0, np.ndarray)205 assert isinstance(img1, np.ndarray)206 self.pred = self._forward(img0, img1)207 if self.conf["ransac"]["enable"]:208 self.pred = self._geometry_check(self.pred)209 return self.pred210 211 def _geometry_check(212 self,213 pred: Dict[str, Any],214 ) -> Dict[str, Any]:215 """216 Filter matches using RANSAC. If keypoints are available, filter by keypoints.217 If lines are available, filter by lines. If both keypoints and lines are218 available, filter by keypoints.219 220 Args:221 pred (Dict[str, Any]): dict of matches, including original keypoints.222 See :func:`filter_matches` for the expected keys.223 224 Returns:225 Dict[str, Any]: filtered matches226 """227 pred = filter_matches(228 pred,229 ransac_method=self.conf["ransac"]["method"],230 ransac_reproj_threshold=self.conf["ransac"]["reproj_threshold"],231 ransac_confidence=self.conf["ransac"]["confidence"],232 ransac_max_iter=self.conf["ransac"]["max_iter"],233 )234 return pred235 236 def visualize(237 self,238 log_path: Optional[Path] = None,239 ) -> None:240 """241 Visualize the matches.242 243 Args:244 log_path (Path, optional): The directory to save the images. Defaults to None.245 246 Returns:247 None248 """249 if self.conf["dense"]:250 postfix = str(self.conf["matcher"]["model"]["name"])251 else:252 postfix = "{}_{}".format(253 str(self.conf["feature"]["model"]["name"]),254 str(self.conf["matcher"]["model"]["name"]),255 )256 titles = [257 "Image 0 - Keypoints",258 "Image 1 - Keypoints",259 ]260 pred: Dict[str, Any] = self.pred261 image0: np.ndarray = pred["image0_orig"]262 image1: np.ndarray = pred["image1_orig"]263 output_keypoints: np.ndarray = plot_images(264 [image0, image1], titles=titles, dpi=300265 )266 if "keypoints0_orig" in pred.keys() and "keypoints1_orig" in pred.keys():267 plot_keypoints([pred["keypoints0_orig"], pred["keypoints1_orig"]])268 text: str = (269 f"# keypoints0: {len(pred['keypoints0_orig'])} \n"270 + f"# keypoints1: {len(pred['keypoints1_orig'])}"271 )272 add_text(0, text, fs=15)273 output_keypoints = fig2im(output_keypoints)274 # plot images with raw matches275 titles = [276 "Image 0 - Raw matched keypoints",277 "Image 1 - Raw matched keypoints",278 ]279 output_matches_raw, num_matches_raw = display_matches(280 pred, titles=titles, tag="KPTS_RAW"281 )282 # plot images with ransac matches283 titles = [284 "Image 0 - Ransac matched keypoints",285 "Image 1 - Ransac matched keypoints",286 ]287 output_matches_ransac, num_matches_ransac = display_matches(288 pred, titles=titles, tag="KPTS_RANSAC"289 )290 if log_path is not None:291 img_keypoints_path: Path = log_path / f"img_keypoints_{postfix}.png"292 img_matches_raw_path: Path = log_path / f"img_matches_raw_{postfix}.png"293 img_matches_ransac_path: Path = (294 log_path / f"img_matches_ransac_{postfix}.png"295 )296 cv2.imwrite(297 str(img_keypoints_path),298 output_keypoints[:, :, ::-1].copy(), # RGB -> BGR299 )300 cv2.imwrite(301 str(img_matches_raw_path),302 output_matches_raw[:, :, ::-1].copy(), # RGB -> BGR303 )304 cv2.imwrite(305 str(img_matches_ransac_path),306 output_matches_ransac[:, :, ::-1].copy(), # RGB -> BGR307 )308 plt.close("all")309