Spaces:
Running
Running
File size: 2,218 Bytes
8320ccc |
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 |
import subprocess
import sys
from pathlib import Path
import torch
from .. import logger
from ..utils.base_model import BaseModel
mickey_path = Path(__file__).parent / "../../third_party"
sys.path.append(str(mickey_path))
from mickey.config.default import cfg
from mickey.lib.models.builder import build_model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class Mickey(BaseModel):
default_conf = {
"config_path": "config.yaml",
"model_name": "mickey.ckpt",
"max_keypoints": 3000,
}
required_inputs = [
"image0",
"image1",
]
weight_urls = "https://storage.googleapis.com/niantic-lon-static/research/mickey/assets/mickey_weights.zip"
# Initialize the line matcher
def _init(self, conf):
model_path = mickey_path / "mickey/mickey_weights" / conf["model_name"]
zip_path = mickey_path / "mickey/mickey_weights.zip"
config_path = model_path.parent / self.conf["config_path"]
# Download the model.
if not model_path.exists():
model_path.parent.mkdir(exist_ok=True, parents=True)
link = self.weight_urls
if not zip_path.exists():
cmd = ["wget", "--quiet", link, "-O", str(zip_path)]
logger.info(f"Downloading the Mickey model with {cmd}.")
subprocess.run(cmd, check=True)
cmd = ["unzip", "-d", str(model_path.parent.parent), str(zip_path)]
logger.info(f"Running {cmd}.")
subprocess.run(cmd, check=True)
logger.info("Loading mickey model...")
cfg.merge_from_file(config_path)
self.net = build_model(cfg, checkpoint=model_path)
logger.info("Load Mickey model done.")
def _forward(self, data):
# data['K_color0'] = torch.from_numpy(K['im0.jpg']).unsqueeze(0).to(device)
# data['K_color1'] = torch.from_numpy(K['im1.jpg']).unsqueeze(0).to(device)
pred = self.net(data)
pred = {
**pred,
**data,
}
inliers = data["inliers_list"]
pred = {
"keypoints0": inliers[:, :2],
"keypoints1": inliers[:, 2:4],
}
return pred
|