File size: 3,301 Bytes
9223079 |
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 |
import sys
from pathlib import Path
import subprocess
import logging
import torch
from PIL import Image
from ..utils.base_model import BaseModel
roma_path = Path(__file__).parent / "../../third_party/Roma"
sys.path.append(str(roma_path))
from roma.models.model_zoo.roma_models import roma_model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logger = logging.getLogger(__name__)
class Roma(BaseModel):
default_conf = {
"name": "two_view_pipeline",
"model_name": "roma_outdoor.pth",
"model_utils_name": "dinov2_vitl14_pretrain.pth",
"max_keypoints": 3000,
}
required_inputs = [
"image0",
"image1",
]
weight_urls = {
"roma": {
"roma_outdoor.pth": "https://github.com/Parskatt/storage/releases/download/roma/roma_outdoor.pth",
"roma_indoor.pth": "https://github.com/Parskatt/storage/releases/download/roma/roma_indoor.pth",
},
"dinov2_vitl14_pretrain.pth": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitl14/dinov2_vitl14_pretrain.pth",
}
# Initialize the line matcher
def _init(self, conf):
model_path = roma_path / "pretrained" / conf["model_name"]
dinov2_weights = roma_path / "pretrained" / conf["model_utils_name"]
# Download the model.
if not model_path.exists():
model_path.parent.mkdir(exist_ok=True)
link = self.weight_urls["roma"][conf["model_name"]]
cmd = ["wget", link, "-O", str(model_path)]
logger.info(f"Downloading the Roma model with `{cmd}`.")
subprocess.run(cmd, check=True)
if not dinov2_weights.exists():
dinov2_weights.parent.mkdir(exist_ok=True)
link = self.weight_urls[conf["model_utils_name"]]
cmd = ["wget", link, "-O", str(dinov2_weights)]
logger.info(f"Downloading the dinov2 model with `{cmd}`.")
subprocess.run(cmd, check=True)
logger.info(f"Loading Roma model...")
# load the model
weights = torch.load(model_path, map_location="cpu")
dinov2_weights = torch.load(dinov2_weights, map_location="cpu")
self.net = roma_model(
resolution=(14 * 8 * 6, 14 * 8 * 6),
upsample_preds=False,
weights=weights,
dinov2_weights=dinov2_weights,
device=device,
)
logger.info(f"Load Roma model done.")
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
# Match
warp, certainty = self.net.match(img0, img1, device=device)
# Sample matches for estimation
matches, certainty = self.net.sample(
warp, certainty, num=self.conf["max_keypoints"]
)
kpts1, kpts2 = self.net.to_pixel_coordinates(matches, H_A, W_A, H_B, W_B)
pred = {}
pred["keypoints0"], pred["keypoints1"] = kpts1, kpts2
pred["mconf"] = certainty
return pred
|