Spaces:
Sleeping
Sleeping
import sys | |
from pathlib import Path | |
import subprocess | |
import torch | |
from PIL import Image | |
from collections import OrderedDict, namedtuple | |
from ..utils.base_model import BaseModel | |
from ..utils import do_system | |
from .. import logger | |
sgmnet_path = Path(__file__).parent / "../../third_party/SGMNet" | |
sys.path.append(str(sgmnet_path)) | |
from sgmnet import matcher as SGM_Model | |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
class SGMNet(BaseModel): | |
default_conf = { | |
"name": "SGM", | |
"model_name": "model_best.pth", | |
"seed_top_k": [256, 256], | |
"seed_radius_coe": 0.01, | |
"net_channels": 128, | |
"layer_num": 9, | |
"head": 4, | |
"seedlayer": [0, 6], | |
"use_mc_seeding": True, | |
"use_score_encoding": False, | |
"conf_bar": [1.11, 0.1], | |
"sink_iter": [10, 100], | |
"detach_iter": 1000000, | |
"match_threshold": 0.2, | |
} | |
required_inputs = [ | |
"image0", | |
"image1", | |
] | |
weight_urls = { | |
"model_best.pth": "https://drive.google.com/uc?id=1Ca0WmKSSt2G6P7m8YAOlSAHEFar_TAWb&confirm=t", | |
} | |
proxy = "http://localhost:1080" | |
# Initialize the line matcher | |
def _init(self, conf): | |
sgmnet_weights = sgmnet_path / "weights/sgm/root" / conf["model_name"] | |
link = self.weight_urls[conf["model_name"]] | |
tar_path = sgmnet_path / "weights.tar.gz" | |
# Download the model. | |
if not sgmnet_weights.exists(): | |
if not tar_path.exists(): | |
cmd = [ | |
"gdown", | |
link, | |
"-O", | |
str(tar_path), | |
"--proxy", | |
self.proxy, | |
] | |
cmd_wo_proxy = ["gdown", link, "-O", str(tar_path)] | |
logger.info( | |
f"Downloading the SGMNet model with `{cmd_wo_proxy}`." | |
) | |
try: | |
subprocess.run(cmd_wo_proxy, check=True) | |
except subprocess.CalledProcessError as e: | |
logger.info(f"Downloading the SGMNet model with `{cmd}`.") | |
try: | |
subprocess.run(cmd, check=True) | |
except subprocess.CalledProcessError as e: | |
logger.error(f"Failed to download the SGMNet model.") | |
raise e | |
cmd = [f"cd {str(sgmnet_path)} & tar -xvf", str(tar_path)] | |
logger.info(f"Unzip model file `{cmd}`.") | |
do_system(f"cd {str(sgmnet_path)} & tar -xvf {str(tar_path)}") | |
# config | |
config = namedtuple("config", conf.keys())(*conf.values()) | |
self.net = SGM_Model(config) | |
checkpoint = torch.load(sgmnet_weights, map_location="cpu") | |
# for ddp model | |
if ( | |
list(checkpoint["state_dict"].items())[0][0].split(".")[0] | |
== "module" | |
): | |
new_stat_dict = OrderedDict() | |
for key, value in checkpoint["state_dict"].items(): | |
new_stat_dict[key[7:]] = value | |
checkpoint["state_dict"] = new_stat_dict | |
self.net.load_state_dict(checkpoint["state_dict"]) | |
logger.info(f"Load SGMNet model done.") | |
def _forward(self, data): | |
x1 = data["keypoints0"].squeeze() # N x 2 | |
x2 = data["keypoints1"].squeeze() | |
score1 = data["scores0"].reshape(-1, 1) # N x 1 | |
score2 = data["scores1"].reshape(-1, 1) | |
desc1 = data["descriptors0"].permute(0, 2, 1) # 1 x N x 128 | |
desc2 = data["descriptors1"].permute(0, 2, 1) | |
size1 = torch.tensor(data["image0"].shape[2:]).flip(0) # W x H -> x & y | |
size2 = torch.tensor(data["image1"].shape[2:]).flip(0) # W x H | |
norm_x1 = self.normalize_size(x1, size1) | |
norm_x2 = self.normalize_size(x2, size2) | |
x1 = torch.cat((norm_x1, score1), dim=-1) # N x 3 | |
x2 = torch.cat((norm_x2, score2), dim=-1) | |
input = {"x1": x1[None], "x2": x2[None], "desc1": desc1, "desc2": desc2} | |
input = { | |
k: v.to(device).float() if isinstance(v, torch.Tensor) else v | |
for k, v in input.items() | |
} | |
pred = self.net(input, test_mode=True) | |
p = pred["p"] # shape: N * M | |
indices0 = self.match_p(p[0, :-1, :-1]) | |
pred = { | |
"matches0": indices0.unsqueeze(0), | |
"matching_scores0": torch.zeros(indices0.size(0)).unsqueeze(0), | |
} | |
return pred | |
def match_p(self, p): | |
score, index = torch.topk(p, k=1, dim=-1) | |
_, index2 = torch.topk(p, k=1, dim=-2) | |
mask_th, index, index2 = ( | |
score[:, 0] > self.conf["match_threshold"], | |
index[:, 0], | |
index2.squeeze(0), | |
) | |
mask_mc = index2[index] == torch.arange(len(p)).to(device) | |
mask = mask_th & mask_mc | |
indices0 = torch.where(mask, index, index.new_tensor(-1)) | |
return indices0 | |
def normalize_size(self, x, size, scale=1): | |
norm_fac = size.max() | |
return (x - size / 2 + 0.5) / (norm_fac * scale) | |