Vincentqyw
add: files
9223079
raw
history blame
2.2 kB
import sys
from pathlib import Path
import torch
from PIL import Image
import subprocess
import logging
from ..utils.base_model import BaseModel
sys.path.append(str(Path(__file__).parent / "../../third_party"))
from DKM.dkm import DKMv3_outdoor
dkm_path = Path(__file__).parent / "../../third_party/DKM"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logger = logging.getLogger(__name__)
class DKMv3(BaseModel):
default_conf = {
"model_name": "DKMv3_outdoor.pth",
"match_threshold": 0.2,
"checkpoint_dir": dkm_path / "pretrained",
}
required_inputs = [
"image0",
"image1",
]
# Models exported using
dkm_models = {
"DKMv3_outdoor.pth": "https://github.com/Parskatt/storage/releases/download/dkmv3/DKMv3_outdoor.pth",
"DKMv3_indoor.pth": "https://github.com/Parskatt/storage/releases/download/dkmv3/DKMv3_indoor.pth",
}
def _init(self, conf):
model_path = dkm_path / "pretrained" / conf["model_name"]
# Download the model.
if not model_path.exists():
model_path.parent.mkdir(exist_ok=True)
link = self.dkm_models[conf["model_name"]]
cmd = ["wget", link, "-O", str(model_path)]
logger.info(f"Downloading the DKMv3 model with `{cmd}`.")
subprocess.run(cmd, check=True)
logger.info(f"Loading DKMv3 model...")
self.net = DKMv3_outdoor(path_to_weights=str(model_path), device=device)
def _forward(self, data):
img0 = data["image0"].cpu().numpy().squeeze() * 255
img1 = data["image1"].cpu().numpy().squeeze() * 255
img0 = img0.transpose(1, 2, 0)
img1 = img1.transpose(1, 2, 0)
img0 = Image.fromarray(img0.astype("uint8"))
img1 = Image.fromarray(img1.astype("uint8"))
W_A, H_A = img0.size
W_B, H_B = img1.size
warp, certainty = self.net.match(img0, img1, device=device)
matches, certainty = self.net.sample(warp, certainty)
kpts1, kpts2 = self.net.to_pixel_coordinates(matches, H_A, W_A, H_B, W_B)
pred = {}
pred["keypoints0"], pred["keypoints1"] = kpts1, kpts2
return pred