Spaces:
Running
Running
File size: 5,011 Bytes
4d4dd90 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
import os
import sys
import torch
import subprocess
import gdown
from pathlib import Path
from ..utils.base_model import BaseModel
from .. import logger
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", model_link, "-O", str(model_path)]
subprocess.run(cmd, check=True)
logger.info(f"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(f"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
|