import sys from pathlib import Path import subprocess import logging import torch from ..utils.base_model import BaseModel logger = logging.getLogger(__name__) gluestick_path = Path(__file__).parent / "../../third_party/GlueStick" sys.path.append(str(gluestick_path)) from gluestick import batch_to_np from gluestick.models.two_view_pipeline import TwoViewPipeline device = torch.device("cuda" if torch.cuda.is_available() else "cpu") class GlueStick(BaseModel): default_conf = { "name": "two_view_pipeline", "model_name": "checkpoint_GlueStick_MD.tar", "use_lines": True, "max_keypoints": 1000, "max_lines": 300, "force_num_keypoints": False, } required_inputs = [ "image0", "image1", ] gluestick_models = { "checkpoint_GlueStick_MD.tar": "https://github.com/cvg/GlueStick/releases/download/v0.1_arxiv/checkpoint_GlueStick_MD.tar", } # Initialize the line matcher def _init(self, conf): model_path = gluestick_path / "resources" / "weights" / conf["model_name"] # Download the model. if not model_path.exists(): model_path.parent.mkdir(exist_ok=True) link = self.gluestick_models[conf["model_name"]] cmd = ["wget", link, "-O", str(model_path)] logger.info(f"Downloading the Gluestick model with `{cmd}`.") subprocess.run(cmd, check=True) logger.info(f"Loading GlueStick model...") gluestick_conf = { "name": "two_view_pipeline", "use_lines": True, "extractor": { "name": "wireframe", "sp_params": { "force_num_keypoints": False, "max_num_keypoints": 1000, }, "wireframe_params": { "merge_points": True, "merge_line_endpoints": True, }, "max_n_lines": 300, }, "matcher": { "name": "gluestick", "weights": str(model_path), "trainable": False, }, "ground_truth": { "from_pose_depth": False, }, } gluestick_conf["extractor"]["sp_params"]["max_num_keypoints"] = conf[ "max_keypoints" ] gluestick_conf["extractor"]["sp_params"]["force_num_keypoints"] = conf[ "force_num_keypoints" ] gluestick_conf["extractor"]["max_n_lines"] = conf["max_lines"] self.net = TwoViewPipeline(gluestick_conf) def _forward(self, data): pred = self.net(data) pred = batch_to_np(pred) kp0, kp1 = pred["keypoints0"], pred["keypoints1"] m0 = pred["matches0"] line_seg0, line_seg1 = pred["lines0"], pred["lines1"] line_matches = pred["line_matches0"] valid_matches = m0 != -1 match_indices = m0[valid_matches] matched_kps0 = kp0[valid_matches] matched_kps1 = kp1[match_indices] valid_matches = line_matches != -1 match_indices = line_matches[valid_matches] matched_lines0 = line_seg0[valid_matches] matched_lines1 = line_seg1[match_indices] pred["raw_lines0"], pred["raw_lines1"] = line_seg0, line_seg1 pred["lines0"], pred["lines1"] = matched_lines0, matched_lines1 pred["keypoints0"], pred["keypoints1"] = torch.from_numpy( matched_kps0 ), torch.from_numpy(matched_kps1) pred = {**pred, **data} return pred