import sys from pathlib import Path import subprocess import torch from ..utils.base_model import BaseModel rekd_path = Path(__file__).parent / "../../third_party/REKD" sys.path.append(str(rekd_path)) from training.model.REKD import REKD as REKD_ device = torch.device("cuda" if torch.cuda.is_available() else "cpu") class REKD(BaseModel): default_conf = { "model_name": "v0", "keypoint_threshold": 0.1, } required_inputs = ["image"] def _init(self, conf): model_path = ( rekd_path / "checkpoints" / f'PointModel_{conf["model_name"]}.pth' ) if not model_path.exists(): print(f"No model found at {model_path}") self.net = REKD_(is_test=True) state_dict = torch.load(model_path, map_location="cpu") self.net.load_state_dict(state_dict["model_state"]) def _forward(self, data): image = data["image"] keypoints, scores, descriptors = self.net(image) _, _, Hc, Wc = descriptors.shape # Scores & Descriptors kpts_score = ( torch.cat([keypoints, scores], dim=1) .view(3, -1) .t() .cpu() .detach() .numpy() ) descriptors = ( descriptors.view(256, Hc, Wc) .view(256, -1) .t() .cpu() .detach() .numpy() ) # Filter based on confidence threshold descriptors = descriptors[ kpts_score[:, 0] > self.conf["keypoint_threshold"], : ] kpts_score = kpts_score[ kpts_score[:, 0] > self.conf["keypoint_threshold"], : ] keypoints = kpts_score[:, 1:] scores = kpts_score[:, 0] return { "keypoints": torch.from_numpy(keypoints)[None], "scores": torch.from_numpy(scores)[None], "descriptors": torch.from_numpy(descriptors.T)[None], }