Spaces:
Running
Running
File size: 5,012 Bytes
4d4dd90 8320ccc 4d4dd90 8320ccc 4d4dd90 8320ccc 4d4dd90 8320ccc 4d4dd90 8320ccc 4d4dd90 8320ccc 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 144 |
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
|