import sys from pathlib import Path import subprocess import torch from ..utils.base_model import BaseModel from .. import logger rord_path = Path(__file__).parent / "../../third_party" sys.path.append(str(rord_path)) from RoRD.lib.model_test import D2Net as _RoRD from RoRD.lib.pyramid import process_multiscale class RoRD(BaseModel): default_conf = { "model_name": "rord.pth", "checkpoint_dir": rord_path / "RoRD" / "models", "use_relu": True, "multiscale": False, "max_keypoints": 1024, } required_inputs = ["image"] weight_urls = { "rord.pth": "https://drive.google.com/uc?id=12414ZGKwgPAjNTGtNrlB4VV9l7W76B2o&confirm=t", } proxy = "http://localhost:1080" def _init(self, conf): model_path = conf["checkpoint_dir"] / conf["model_name"] link = self.weight_urls[conf["model_name"]] if not model_path.exists(): model_path.parent.mkdir(exist_ok=True) cmd_wo_proxy = ["gdown", link, "-O", str(model_path)] cmd = ["gdown", link, "-O", str(model_path), "--proxy", self.proxy] logger.info( f"Downloading the RoRD model with `{cmd_wo_proxy}`." ) try: subprocess.run(cmd_wo_proxy, check=True) except subprocess.CalledProcessError as e: logger.info(f"Downloading the RoRD model with `{cmd}`.") try: subprocess.run(cmd, check=True) except subprocess.CalledProcessError as e: logger.error(f"Failed to download the RoRD model.") raise e logger.info("RoRD model loaded.") self.net = _RoRD( model_file=model_path, use_relu=conf["use_relu"], use_cuda=False ) def _forward(self, data): image = data["image"] image = image.flip(1) # RGB -> BGR norm = image.new_tensor([103.939, 116.779, 123.68]) image = image * 255 - norm.view(1, 3, 1, 1) # caffe normalization if self.conf["multiscale"]: keypoints, scores, descriptors = process_multiscale(image, self.net) else: keypoints, scores, descriptors = process_multiscale( image, self.net, scales=[1] ) keypoints = keypoints[:, [1, 0]] # (x, y) and remove the scale idxs = scores.argsort()[-self.conf["max_keypoints"] or None :] keypoints = keypoints[idxs, :2] descriptors = descriptors[idxs] scores = scores[idxs] return { "keypoints": torch.from_numpy(keypoints)[None], "scores": torch.from_numpy(scores)[None], "descriptors": torch.from_numpy(descriptors.T)[None], }