Spaces:
Running
Running
| import subprocess | |
| import sys | |
| from pathlib import Path | |
| import gdown | |
| import torch | |
| from .. import logger | |
| from ..utils.base_model import BaseModel | |
| gim_path = Path(__file__).parent / "../../third_party/gim" | |
| sys.path.append(str(gim_path)) | |
| from dkm.models.model_zoo.DKMv3 import DKMv3 | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| class GIM(BaseModel): | |
| default_conf = { | |
| "model_name": "gim_dkm_100h.ckpt", | |
| "match_threshold": 0.2, | |
| "checkpoint_dir": gim_path / "weights", | |
| } | |
| required_inputs = [ | |
| "image0", | |
| "image1", | |
| ] | |
| model_dict = { | |
| "gim_lightglue_100h.ckpt": "https://github.com/xuelunshen/gim/blob/main/weights/gim_lightglue_100h.ckpt", | |
| "gim_dkm_100h.ckpt": "https://drive.google.com/file/d/1gk97V4IROnR1Nprq10W9NCFUv2mxXR_-/view", | |
| } | |
| def _init(self, conf): | |
| conf["model_name"] = str(conf["weights"]) | |
| if conf["model_name"] not in self.model_dict: | |
| raise ValueError(f"Unknown GIM model {conf['model_name']}.") | |
| model_path = conf["checkpoint_dir"] / conf["model_name"] | |
| # Download the model. | |
| if not model_path.exists(): | |
| model_path.parent.mkdir(exist_ok=True) | |
| model_link = self.model_dict[conf["model_name"]] | |
| if "drive.google.com" in model_link: | |
| gdown.download(model_link, output=str(model_path), fuzzy=True) | |
| else: | |
| cmd = ["wget", "--quiet", model_link, "-O", str(model_path)] | |
| subprocess.run(cmd, check=True) | |
| logger.info("Downloaded GIM model succeeed!") | |
| self.aspect_ratio = 896 / 672 | |
| model = DKMv3(None, 672, 896, upsample_preds=True) | |
| state_dict = torch.load(str(model_path), map_location="cpu") | |
| if "state_dict" in state_dict.keys(): | |
| state_dict = state_dict["state_dict"] | |
| for k in list(state_dict.keys()): | |
| if k.startswith("model."): | |
| state_dict[k.replace("model.", "", 1)] = state_dict.pop(k) | |
| if "encoder.net.fc" in k: | |
| state_dict.pop(k) | |
| model.load_state_dict(state_dict) | |
| self.net = model | |
| logger.info("Loaded GIM model") | |
| def pad_image(self, image, aspect_ratio): | |
| new_width = max(image.shape[3], int(image.shape[2] * aspect_ratio)) | |
| new_height = max(image.shape[2], int(image.shape[3] / aspect_ratio)) | |
| pad_width = new_width - image.shape[3] | |
| pad_height = new_height - image.shape[2] | |
| return torch.nn.functional.pad( | |
| image, | |
| ( | |
| pad_width // 2, | |
| pad_width - pad_width // 2, | |
| pad_height // 2, | |
| pad_height - pad_height // 2, | |
| ), | |
| ) | |
| def rescale_kpts(self, sparse_matches, shape0, shape1): | |
| kpts0 = torch.stack( | |
| ( | |
| shape0[1] * (sparse_matches[:, 0] + 1) / 2, | |
| shape0[0] * (sparse_matches[:, 1] + 1) / 2, | |
| ), | |
| dim=-1, | |
| ) | |
| kpts1 = torch.stack( | |
| ( | |
| shape1[1] * (sparse_matches[:, 2] + 1) / 2, | |
| shape1[0] * (sparse_matches[:, 3] + 1) / 2, | |
| ), | |
| dim=-1, | |
| ) | |
| return kpts0, kpts1 | |
| def compute_mask(self, kpts0, kpts1, orig_shape0, orig_shape1): | |
| mask = ( | |
| (kpts0[:, 0] > 0) | |
| & (kpts0[:, 1] > 0) | |
| & (kpts1[:, 0] > 0) | |
| & (kpts1[:, 1] > 0) | |
| ) | |
| mask &= ( | |
| (kpts0[:, 0] <= (orig_shape0[1] - 1)) | |
| & (kpts1[:, 0] <= (orig_shape1[1] - 1)) | |
| & (kpts0[:, 1] <= (orig_shape0[0] - 1)) | |
| & (kpts1[:, 1] <= (orig_shape1[0] - 1)) | |
| ) | |
| return mask | |
| def _forward(self, data): | |
| image0, image1 = self.pad_image( | |
| data["image0"], self.aspect_ratio | |
| ), self.pad_image(data["image1"], self.aspect_ratio) | |
| dense_matches, dense_certainty = self.net.match(image0, image1) | |
| sparse_matches, mconf = self.net.sample( | |
| dense_matches, dense_certainty, self.conf["max_keypoints"] | |
| ) | |
| kpts0, kpts1 = self.rescale_kpts( | |
| sparse_matches, image0.shape[-2:], image1.shape[-2:] | |
| ) | |
| mask = self.compute_mask( | |
| kpts0, kpts1, data["image0"].shape[-2:], data["image1"].shape[-2:] | |
| ) | |
| b_ids, i_ids = torch.where(mconf[None]) | |
| pred = { | |
| "keypoints0": kpts0[i_ids], | |
| "keypoints1": kpts1[i_ids], | |
| "confidence": mconf[i_ids], | |
| "batch_indexes": b_ids, | |
| } | |
| scores, b_ids = pred["confidence"], pred["batch_indexes"] | |
| kpts0, kpts1 = pred["keypoints0"], pred["keypoints1"] | |
| pred["confidence"], pred["batch_indexes"] = scores[mask], b_ids[mask] | |
| pred["keypoints0"], pred["keypoints1"] = kpts0[mask], kpts1[mask] | |
| out = { | |
| "keypoints0": pred["keypoints0"], | |
| "keypoints1": pred["keypoints1"], | |
| } | |
| return out | |