image-matching-webui / hloc /match_dense.py
Realcat
add: xfeat+lightglue
848664a
import argparse
import pprint
from collections import Counter, defaultdict
from itertools import chain
from pathlib import Path
from types import SimpleNamespace
from typing import Dict, Iterable, List, Optional, Set, Tuple, Union
import cv2
import h5py
import numpy as np
import torch
import torchvision.transforms.functional as F
from scipy.spatial import KDTree
from tqdm import tqdm
from . import logger, matchers
from .extract_features import read_image, resize_image
from .match_features import find_unique_new_pairs
from .utils.base_model import dynamic_load
from .utils.io import list_h5_names
from .utils.parsers import names_to_pair, parse_retrieval
device = "cuda" if torch.cuda.is_available() else "cpu"
confs = {
# Best quality but loads of points. Only use for small scenes
"loftr": {
"output": "matches-loftr",
"model": {
"name": "loftr",
"weights": "outdoor",
"max_keypoints": 2000,
"match_threshold": 0.2,
},
"preprocessing": {
"grayscale": True,
"resize_max": 1024,
"dfactor": 8,
"width": 640,
"height": 480,
"force_resize": True,
},
"max_error": 1, # max error for assigned keypoints (in px)
"cell_size": 1, # size of quantization patch (max 1 kp/patch)
},
"eloftr": {
"output": "matches-eloftr",
"model": {
"name": "eloftr",
"weights": "weights/eloftr_outdoor.ckpt",
"max_keypoints": 2000,
"match_threshold": 0.2,
},
"preprocessing": {
"grayscale": True,
"resize_max": 1024,
"dfactor": 32,
"width": 640,
"height": 480,
"force_resize": True,
},
"max_error": 1, # max error for assigned keypoints (in px)
"cell_size": 1, # size of quantization patch (max 1 kp/patch)
},
# "loftr_quadtree": {
# "output": "matches-loftr-quadtree",
# "model": {
# "name": "quadtree",
# "weights": "outdoor",
# "max_keypoints": 2000,
# "match_threshold": 0.2,
# },
# "preprocessing": {
# "grayscale": True,
# "resize_max": 1024,
# "dfactor": 8,
# "width": 640,
# "height": 480,
# "force_resize": True,
# },
# "max_error": 1, # max error for assigned keypoints (in px)
# "cell_size": 1, # size of quantization patch (max 1 kp/patch)
# },
"cotr": {
"output": "matches-cotr",
"model": {
"name": "cotr",
"weights": "out/default",
"max_keypoints": 2000,
"match_threshold": 0.2,
},
"preprocessing": {
"grayscale": False,
"resize_max": 1024,
"dfactor": 8,
"width": 640,
"height": 480,
"force_resize": True,
},
"max_error": 1, # max error for assigned keypoints (in px)
"cell_size": 1, # size of quantization patch (max 1 kp/patch)
},
# Semi-scalable loftr which limits detected keypoints
"loftr_aachen": {
"output": "matches-loftr_aachen",
"model": {
"name": "loftr",
"weights": "outdoor",
"max_keypoints": 2000,
"match_threshold": 0.2,
},
"preprocessing": {
"grayscale": True,
"resize_max": 1024,
"dfactor": 8,
"width": 640,
"height": 480,
"force_resize": True,
},
"max_error": 2, # max error for assigned keypoints (in px)
"cell_size": 8, # size of quantization patch (max 1 kp/patch)
},
# Use for matching superpoint feats with loftr
"loftr_superpoint": {
"output": "matches-loftr_aachen",
"model": {
"name": "loftr",
"weights": "outdoor",
"max_keypoints": 2000,
"match_threshold": 0.2,
},
"preprocessing": {
"grayscale": True,
"resize_max": 1024,
"dfactor": 8,
"width": 640,
"height": 480,
"force_resize": True,
},
"max_error": 4, # max error for assigned keypoints (in px)
"cell_size": 4, # size of quantization patch (max 1 kp/patch)
},
# Use topicfm for matching feats
"topicfm": {
"output": "matches-topicfm",
"model": {
"name": "topicfm",
"weights": "outdoor",
"max_keypoints": 2000,
"match_threshold": 0.2,
},
"preprocessing": {
"grayscale": True,
"force_resize": True,
"resize_max": 1024,
"dfactor": 8,
"width": 640,
"height": 480,
},
},
# Use aspanformer for matching feats
"aspanformer": {
"output": "matches-aspanformer",
"model": {
"name": "aspanformer",
"weights": "outdoor",
"max_keypoints": 2000,
"match_threshold": 0.2,
},
"preprocessing": {
"grayscale": True,
"force_resize": True,
"resize_max": 1024,
"width": 640,
"height": 480,
"dfactor": 8,
},
},
"duster": {
"output": "matches-duster",
"model": {
"name": "duster",
"weights": "vit_large",
"max_keypoints": 2000,
"match_threshold": 0.2,
},
"preprocessing": {
"grayscale": False,
"resize_max": 512,
"dfactor": 16,
},
},
"mast3r": {
"output": "matches-mast3r",
"model": {
"name": "mast3r",
"weights": "vit_large",
"max_keypoints": 2000,
"match_threshold": 0.2,
},
"preprocessing": {
"grayscale": False,
"resize_max": 512,
"dfactor": 16,
},
},
"xfeat_lightglue": {
"output": "matches-xfeat_lightglue",
"model": {
"name": "xfeat_lightglue",
"max_keypoints": 8000,
},
"preprocessing": {
"grayscale": False,
"force_resize": False,
"resize_max": 1024,
"width": 640,
"height": 480,
"dfactor": 8,
},
},
"xfeat_dense": {
"output": "matches-xfeat_dense",
"model": {
"name": "xfeat_dense",
"max_keypoints": 8000,
},
"preprocessing": {
"grayscale": False,
"force_resize": False,
"resize_max": 1024,
"width": 640,
"height": 480,
"dfactor": 8,
},
},
"dkm": {
"output": "matches-dkm",
"model": {
"name": "dkm",
"weights": "outdoor",
"max_keypoints": 2000,
"match_threshold": 0.2,
},
"preprocessing": {
"grayscale": False,
"force_resize": True,
"resize_max": 1024,
"width": 80,
"height": 60,
"dfactor": 8,
},
},
"roma": {
"output": "matches-roma",
"model": {
"name": "roma",
"weights": "outdoor",
"max_keypoints": 2000,
"match_threshold": 0.2,
},
"preprocessing": {
"grayscale": False,
"force_resize": True,
"resize_max": 1024,
"width": 320,
"height": 240,
"dfactor": 8,
},
},
"gim(dkm)": {
"output": "matches-gim",
"model": {
"name": "gim",
"weights": "gim_dkm_100h.ckpt",
"max_keypoints": 2000,
"match_threshold": 0.2,
},
"preprocessing": {
"grayscale": False,
"force_resize": True,
"resize_max": 1024,
"width": 320,
"height": 240,
"dfactor": 8,
},
},
"omniglue": {
"output": "matches-omniglue",
"model": {
"name": "omniglue",
"match_threshold": 0.2,
"max_keypoints": 2000,
"features": "null",
},
"preprocessing": {
"grayscale": False,
"resize_max": 1024,
"dfactor": 8,
"force_resize": False,
"resize_max": 1024,
"width": 640,
"height": 480,
"dfactor": 8,
},
},
"sold2": {
"output": "matches-sold2",
"model": {
"name": "sold2",
"max_keypoints": 2000,
"match_threshold": 0.2,
},
"preprocessing": {
"grayscale": True,
"force_resize": True,
"resize_max": 1024,
"width": 640,
"height": 480,
"dfactor": 8,
},
},
"gluestick": {
"output": "matches-gluestick",
"model": {
"name": "gluestick",
"use_lines": True,
"max_keypoints": 1000,
"max_lines": 300,
"force_num_keypoints": False,
},
"preprocessing": {
"grayscale": True,
"force_resize": True,
"resize_max": 1024,
"width": 640,
"height": 480,
"dfactor": 8,
},
},
}
def to_cpts(kpts, ps):
if ps > 0.0:
kpts = np.round(np.round((kpts + 0.5) / ps) * ps - 0.5, 2)
return [tuple(cpt) for cpt in kpts]
def assign_keypoints(
kpts: np.ndarray,
other_cpts: Union[List[Tuple], np.ndarray],
max_error: float,
update: bool = False,
ref_bins: Optional[List[Counter]] = None,
scores: Optional[np.ndarray] = None,
cell_size: Optional[int] = None,
):
if not update:
# Without update this is just a NN search
if len(other_cpts) == 0 or len(kpts) == 0:
return np.full(len(kpts), -1)
dist, kpt_ids = KDTree(np.array(other_cpts)).query(kpts)
valid = dist <= max_error
kpt_ids[~valid] = -1
return kpt_ids
else:
ps = cell_size if cell_size is not None else max_error
ps = max(ps, max_error)
# With update we quantize and bin (optionally)
assert isinstance(other_cpts, list)
kpt_ids = []
cpts = to_cpts(kpts, ps)
bpts = to_cpts(kpts, int(max_error))
cp_to_id = {val: i for i, val in enumerate(other_cpts)}
for i, (cpt, bpt) in enumerate(zip(cpts, bpts)):
try:
kid = cp_to_id[cpt]
except KeyError:
kid = len(cp_to_id)
cp_to_id[cpt] = kid
other_cpts.append(cpt)
if ref_bins is not None:
ref_bins.append(Counter())
if ref_bins is not None:
score = scores[i] if scores is not None else 1
ref_bins[cp_to_id[cpt]][bpt] += score
kpt_ids.append(kid)
return np.array(kpt_ids)
def get_grouped_ids(array):
# Group array indices based on its values
# all duplicates are grouped as a set
idx_sort = np.argsort(array)
sorted_array = array[idx_sort]
_, ids, _ = np.unique(sorted_array, return_counts=True, return_index=True)
res = np.split(idx_sort, ids[1:])
return res
def get_unique_matches(match_ids, scores):
if len(match_ids.shape) == 1:
return [0]
isets1 = get_grouped_ids(match_ids[:, 0])
isets2 = get_grouped_ids(match_ids[:, 1])
uid1s = [ids[scores[ids].argmax()] for ids in isets1 if len(ids) > 0]
uid2s = [ids[scores[ids].argmax()] for ids in isets2 if len(ids) > 0]
uids = list(set(uid1s).intersection(uid2s))
return match_ids[uids], scores[uids]
def matches_to_matches0(matches, scores):
if len(matches) == 0:
return np.zeros(0, dtype=np.int32), np.zeros(0, dtype=np.float16)
n_kps0 = np.max(matches[:, 0]) + 1
matches0 = -np.ones((n_kps0,))
scores0 = np.zeros((n_kps0,))
matches0[matches[:, 0]] = matches[:, 1]
scores0[matches[:, 0]] = scores
return matches0.astype(np.int32), scores0.astype(np.float16)
def kpids_to_matches0(kpt_ids0, kpt_ids1, scores):
valid = (kpt_ids0 != -1) & (kpt_ids1 != -1)
matches = np.dstack([kpt_ids0[valid], kpt_ids1[valid]])
matches = matches.reshape(-1, 2)
scores = scores[valid]
# Remove n-to-1 matches
matches, scores = get_unique_matches(matches, scores)
return matches_to_matches0(matches, scores)
def scale_keypoints(kpts, scale):
if np.any(scale != 1.0):
kpts *= kpts.new_tensor(scale)
return kpts
class ImagePairDataset(torch.utils.data.Dataset):
default_conf = {
"grayscale": True,
"resize_max": 1024,
"dfactor": 8,
"cache_images": False,
}
def __init__(self, image_dir, conf, pairs):
self.image_dir = image_dir
self.conf = conf = SimpleNamespace(**{**self.default_conf, **conf})
self.pairs = pairs
if self.conf.cache_images:
image_names = set(sum(pairs, ())) # unique image names in pairs
logger.info(
f"Loading and caching {len(image_names)} unique images."
)
self.images = {}
self.scales = {}
for name in tqdm(image_names):
image = read_image(self.image_dir / name, self.conf.grayscale)
self.images[name], self.scales[name] = self.preprocess(image)
def preprocess(self, image: np.ndarray):
image = image.astype(np.float32, copy=False)
size = image.shape[:2][::-1]
scale = np.array([1.0, 1.0])
if self.conf.resize_max:
scale = self.conf.resize_max / max(size)
if scale < 1.0:
size_new = tuple(int(round(x * scale)) for x in size)
image = resize_image(image, size_new, "cv2_area")
scale = np.array(size) / np.array(size_new)
if self.conf.grayscale:
assert image.ndim == 2, image.shape
image = image[None]
else:
image = image.transpose((2, 0, 1)) # HxWxC to CxHxW
image = torch.from_numpy(image / 255.0).float()
# assure that the size is divisible by dfactor
size_new = tuple(
map(
lambda x: int(x // self.conf.dfactor * self.conf.dfactor),
image.shape[-2:],
)
)
image = F.resize(image, size=size_new)
scale = np.array(size) / np.array(size_new)[::-1]
return image, scale
def __len__(self):
return len(self.pairs)
def __getitem__(self, idx):
name0, name1 = self.pairs[idx]
if self.conf.cache_images:
image0, scale0 = self.images[name0], self.scales[name0]
image1, scale1 = self.images[name1], self.scales[name1]
else:
image0 = read_image(self.image_dir / name0, self.conf.grayscale)
image1 = read_image(self.image_dir / name1, self.conf.grayscale)
image0, scale0 = self.preprocess(image0)
image1, scale1 = self.preprocess(image1)
return image0, image1, scale0, scale1, name0, name1
@torch.no_grad()
def match_dense(
conf: Dict,
pairs: List[Tuple[str, str]],
image_dir: Path,
match_path: Path, # out
existing_refs: Optional[List] = [],
):
device = "cuda" if torch.cuda.is_available() else "cpu"
Model = dynamic_load(matchers, conf["model"]["name"])
model = Model(conf["model"]).eval().to(device)
dataset = ImagePairDataset(image_dir, conf["preprocessing"], pairs)
loader = torch.utils.data.DataLoader(
dataset, num_workers=16, batch_size=1, shuffle=False
)
logger.info("Performing dense matching...")
with h5py.File(str(match_path), "a") as fd:
for data in tqdm(loader, smoothing=0.1):
# load image-pair data
image0, image1, scale0, scale1, (name0,), (name1,) = data
scale0, scale1 = scale0[0].numpy(), scale1[0].numpy()
image0, image1 = image0.to(device), image1.to(device)
# match semi-dense
# for consistency with pairs_from_*: refine kpts of image0
if name0 in existing_refs:
# special case: flip to enable refinement in query image
pred = model({"image0": image1, "image1": image0})
pred = {
**pred,
"keypoints0": pred["keypoints1"],
"keypoints1": pred["keypoints0"],
}
else:
# usual case
pred = model({"image0": image0, "image1": image1})
# Rescale keypoints and move to cpu
kpts0, kpts1 = pred["keypoints0"], pred["keypoints1"]
kpts0 = scale_keypoints(kpts0 + 0.5, scale0) - 0.5
kpts1 = scale_keypoints(kpts1 + 0.5, scale1) - 0.5
kpts0 = kpts0.cpu().numpy()
kpts1 = kpts1.cpu().numpy()
scores = pred["scores"].cpu().numpy()
# Write matches and matching scores in hloc format
pair = names_to_pair(name0, name1)
if pair in fd:
del fd[pair]
grp = fd.create_group(pair)
# Write dense matching output
grp.create_dataset("keypoints0", data=kpts0)
grp.create_dataset("keypoints1", data=kpts1)
grp.create_dataset("scores", data=scores)
del model, loader
# default: quantize all!
def load_keypoints(
conf: Dict, feature_paths_refs: List[Path], quantize: Optional[set] = None
):
name2ref = {
n: i for i, p in enumerate(feature_paths_refs) for n in list_h5_names(p)
}
existing_refs = set(name2ref.keys())
if quantize is None:
quantize = existing_refs # quantize all
if len(existing_refs) > 0:
logger.info(f"Loading keypoints from {len(existing_refs)} images.")
# Load query keypoints
cpdict = defaultdict(list)
bindict = defaultdict(list)
for name in existing_refs:
with h5py.File(str(feature_paths_refs[name2ref[name]]), "r") as fd:
kps = fd[name]["keypoints"].__array__()
if name not in quantize:
cpdict[name] = kps
else:
if "scores" in fd[name].keys():
kp_scores = fd[name]["scores"].__array__()
else:
# we set the score to 1.0 if not provided
# increase for more weight on reference keypoints for
# stronger anchoring
kp_scores = [1.0 for _ in range(kps.shape[0])]
# bin existing keypoints of reference images for association
assign_keypoints(
kps,
cpdict[name],
conf["max_error"],
True,
bindict[name],
kp_scores,
conf["cell_size"],
)
return cpdict, bindict
def aggregate_matches(
conf: Dict,
pairs: List[Tuple[str, str]],
match_path: Path,
feature_path: Path,
required_queries: Optional[Set[str]] = None,
max_kps: Optional[int] = None,
cpdict: Dict[str, Iterable] = defaultdict(list),
bindict: Dict[str, List[Counter]] = defaultdict(list),
):
if required_queries is None:
required_queries = set(sum(pairs, ()))
# default: do not overwrite existing features in feature_path!
required_queries -= set(list_h5_names(feature_path))
# if an entry in cpdict is provided as np.ndarray we assume it is fixed
required_queries -= set(
[k for k, v in cpdict.items() if isinstance(v, np.ndarray)]
)
# sort pairs for reduced RAM
pairs_per_q = Counter(list(chain(*pairs)))
pairs_score = [min(pairs_per_q[i], pairs_per_q[j]) for i, j in pairs]
pairs = [p for _, p in sorted(zip(pairs_score, pairs))]
if len(required_queries) > 0:
logger.info(
f"Aggregating keypoints for {len(required_queries)} images."
)
n_kps = 0
with h5py.File(str(match_path), "a") as fd:
for name0, name1 in tqdm(pairs, smoothing=0.1):
pair = names_to_pair(name0, name1)
grp = fd[pair]
kpts0 = grp["keypoints0"].__array__()
kpts1 = grp["keypoints1"].__array__()
scores = grp["scores"].__array__()
# Aggregate local features
update0 = name0 in required_queries
update1 = name1 in required_queries
# in localization we do not want to bin the query kp
# assumes that the query is name0!
if update0 and not update1 and max_kps is None:
max_error0 = cell_size0 = 0.0
else:
max_error0 = conf["max_error"]
cell_size0 = conf["cell_size"]
# Get match ids and extend query keypoints (cpdict)
mkp_ids0 = assign_keypoints(
kpts0,
cpdict[name0],
max_error0,
update0,
bindict[name0],
scores,
cell_size0,
)
mkp_ids1 = assign_keypoints(
kpts1,
cpdict[name1],
conf["max_error"],
update1,
bindict[name1],
scores,
conf["cell_size"],
)
# Build matches from assignments
matches0, scores0 = kpids_to_matches0(mkp_ids0, mkp_ids1, scores)
assert kpts0.shape[0] == scores.shape[0]
grp.create_dataset("matches0", data=matches0)
grp.create_dataset("matching_scores0", data=scores0)
# Convert bins to kps if finished, and store them
for name in (name0, name1):
pairs_per_q[name] -= 1
if pairs_per_q[name] > 0 or name not in required_queries:
continue
kp_score = [c.most_common(1)[0][1] for c in bindict[name]]
cpdict[name] = [c.most_common(1)[0][0] for c in bindict[name]]
cpdict[name] = np.array(cpdict[name], dtype=np.float32)
# Select top-k query kps by score (reassign matches later)
if max_kps:
top_k = min(max_kps, cpdict[name].shape[0])
top_k = np.argsort(kp_score)[::-1][:top_k]
cpdict[name] = cpdict[name][top_k]
kp_score = np.array(kp_score)[top_k]
# Write query keypoints
with h5py.File(feature_path, "a") as kfd:
if name in kfd:
del kfd[name]
kgrp = kfd.create_group(name)
kgrp.create_dataset("keypoints", data=cpdict[name])
kgrp.create_dataset("score", data=kp_score)
n_kps += cpdict[name].shape[0]
del bindict[name]
if len(required_queries) > 0:
avg_kp_per_image = round(n_kps / len(required_queries), 1)
logger.info(
f"Finished assignment, found {avg_kp_per_image} "
f"keypoints/image (avg.), total {n_kps}."
)
return cpdict
def assign_matches(
pairs: List[Tuple[str, str]],
match_path: Path,
keypoints: Union[List[Path], Dict[str, np.array]],
max_error: float,
):
if isinstance(keypoints, list):
keypoints = load_keypoints({}, keypoints, kpts_as_bin=set([]))
assert len(set(sum(pairs, ())) - set(keypoints.keys())) == 0
with h5py.File(str(match_path), "a") as fd:
for name0, name1 in tqdm(pairs):
pair = names_to_pair(name0, name1)
grp = fd[pair]
kpts0 = grp["keypoints0"].__array__()
kpts1 = grp["keypoints1"].__array__()
scores = grp["scores"].__array__()
# NN search across cell boundaries
mkp_ids0 = assign_keypoints(kpts0, keypoints[name0], max_error)
mkp_ids1 = assign_keypoints(kpts1, keypoints[name1], max_error)
matches0, scores0 = kpids_to_matches0(mkp_ids0, mkp_ids1, scores)
# overwrite matches0 and matching_scores0
del grp["matches0"], grp["matching_scores0"]
grp.create_dataset("matches0", data=matches0)
grp.create_dataset("matching_scores0", data=scores0)
@torch.no_grad()
def match_and_assign(
conf: Dict,
pairs_path: Path,
image_dir: Path,
match_path: Path, # out
feature_path_q: Path, # out
feature_paths_refs: Optional[List[Path]] = [],
max_kps: Optional[int] = 8192,
overwrite: bool = False,
) -> Path:
for path in feature_paths_refs:
if not path.exists():
raise FileNotFoundError(f"Reference feature file {path}.")
pairs = parse_retrieval(pairs_path)
pairs = [(q, r) for q, rs in pairs.items() for r in rs]
pairs = find_unique_new_pairs(pairs, None if overwrite else match_path)
required_queries = set(sum(pairs, ()))
name2ref = {
n: i for i, p in enumerate(feature_paths_refs) for n in list_h5_names(p)
}
existing_refs = required_queries.intersection(set(name2ref.keys()))
# images which require feature extraction
required_queries = required_queries - existing_refs
if feature_path_q.exists():
existing_queries = set(list_h5_names(feature_path_q))
feature_paths_refs.append(feature_path_q)
existing_refs = set.union(existing_refs, existing_queries)
if not overwrite:
required_queries = required_queries - existing_queries
if len(pairs) == 0 and len(required_queries) == 0:
logger.info("All pairs exist. Skipping dense matching.")
return
# extract semi-dense matches
match_dense(conf, pairs, image_dir, match_path, existing_refs=existing_refs)
logger.info("Assigning matches...")
# Pre-load existing keypoints
cpdict, bindict = load_keypoints(
conf, feature_paths_refs, quantize=required_queries
)
# Reassign matches by aggregation
cpdict = aggregate_matches(
conf,
pairs,
match_path,
feature_path=feature_path_q,
required_queries=required_queries,
max_kps=max_kps,
cpdict=cpdict,
bindict=bindict,
)
# Invalidate matches that are far from selected bin by reassignment
if max_kps is not None:
logger.info(f'Reassign matches with max_error={conf["max_error"]}.')
assign_matches(pairs, match_path, cpdict, max_error=conf["max_error"])
def scale_lines(lines, scale):
if np.any(scale != 1.0):
lines *= lines.new_tensor(scale)
return lines
def match(model, path_0, path_1, conf):
default_conf = {
"grayscale": True,
"resize_max": 1024,
"dfactor": 8,
"cache_images": False,
"force_resize": False,
"width": 320,
"height": 240,
}
def preprocess(image: np.ndarray):
image = image.astype(np.float32, copy=False)
size = image.shape[:2][::-1]
scale = np.array([1.0, 1.0])
if conf.resize_max:
scale = conf.resize_max / max(size)
if scale < 1.0:
size_new = tuple(int(round(x * scale)) for x in size)
image = resize_image(image, size_new, "cv2_area")
scale = np.array(size) / np.array(size_new)
if conf.force_resize:
size = image.shape[:2][::-1]
image = resize_image(image, (conf.width, conf.height), "cv2_area")
size_new = (conf.width, conf.height)
scale = np.array(size) / np.array(size_new)
if conf.grayscale:
assert image.ndim == 2, image.shape
image = image[None]
else:
image = image.transpose((2, 0, 1)) # HxWxC to CxHxW
image = torch.from_numpy(image / 255.0).float()
# assure that the size is divisible by dfactor
size_new = tuple(
map(
lambda x: int(x // conf.dfactor * conf.dfactor),
image.shape[-2:],
)
)
image = F.resize(image, size=size_new, antialias=True)
scale = np.array(size) / np.array(size_new)[::-1]
return image, scale
conf = SimpleNamespace(**{**default_conf, **conf})
image0 = read_image(path_0, conf.grayscale)
image1 = read_image(path_1, conf.grayscale)
image0, scale0 = preprocess(image0)
image1, scale1 = preprocess(image1)
image0 = image0.to(device)[None]
image1 = image1.to(device)[None]
pred = model({"image0": image0, "image1": image1})
# Rescale keypoints and move to cpu
kpts0, kpts1 = pred["keypoints0"], pred["keypoints1"]
kpts0 = scale_keypoints(kpts0 + 0.5, scale0) - 0.5
kpts1 = scale_keypoints(kpts1 + 0.5, scale1) - 0.5
ret = {
"image0": image0.squeeze().cpu().numpy(),
"image1": image1.squeeze().cpu().numpy(),
"keypoints0": kpts0.cpu().numpy(),
"keypoints1": kpts1.cpu().numpy(),
}
if "mconf" in pred.keys():
ret["mconf"] = pred["mconf"].cpu().numpy()
return ret
@torch.no_grad()
def match_images(model, image_0, image_1, conf, device="cpu"):
default_conf = {
"grayscale": True,
"resize_max": 1024,
"dfactor": 8,
"cache_images": False,
"force_resize": False,
"width": 320,
"height": 240,
}
def preprocess(image: np.ndarray):
image = image.astype(np.float32, copy=False)
size = image.shape[:2][::-1]
scale = np.array([1.0, 1.0])
if conf.resize_max:
scale = conf.resize_max / max(size)
if scale < 1.0:
size_new = tuple(int(round(x * scale)) for x in size)
image = resize_image(image, size_new, "cv2_area")
scale = np.array(size) / np.array(size_new)
if conf.force_resize:
size = image.shape[:2][::-1]
image = resize_image(image, (conf.width, conf.height), "cv2_area")
size_new = (conf.width, conf.height)
scale = np.array(size) / np.array(size_new)
if conf.grayscale:
assert image.ndim == 2, image.shape
image = image[None]
else:
image = image.transpose((2, 0, 1)) # HxWxC to CxHxW
image = torch.from_numpy(image / 255.0).float()
# assure that the size is divisible by dfactor
size_new = tuple(
map(
lambda x: int(x // conf.dfactor * conf.dfactor),
image.shape[-2:],
)
)
image = F.resize(image, size=size_new)
scale = np.array(size) / np.array(size_new)[::-1]
return image, scale
conf = SimpleNamespace(**{**default_conf, **conf})
if len(image_0.shape) == 3 and conf.grayscale:
image0 = cv2.cvtColor(image_0, cv2.COLOR_RGB2GRAY)
else:
image0 = image_0
if len(image_0.shape) == 3 and conf.grayscale:
image1 = cv2.cvtColor(image_1, cv2.COLOR_RGB2GRAY)
else:
image1 = image_1
# comment following lines, image is always RGB mode
# if not conf.grayscale and len(image0.shape) == 3:
# image0 = image0[:, :, ::-1] # BGR to RGB
# if not conf.grayscale and len(image1.shape) == 3:
# image1 = image1[:, :, ::-1] # BGR to RGB
image0, scale0 = preprocess(image0)
image1, scale1 = preprocess(image1)
image0 = image0.to(device)[None]
image1 = image1.to(device)[None]
pred = model({"image0": image0, "image1": image1})
s0 = np.array(image_0.shape[:2][::-1]) / np.array(image0.shape[-2:][::-1])
s1 = np.array(image_1.shape[:2][::-1]) / np.array(image1.shape[-2:][::-1])
# Rescale keypoints and move to cpu
if "keypoints0" in pred.keys() and "keypoints1" in pred.keys():
kpts0, kpts1 = pred["keypoints0"], pred["keypoints1"]
kpts0_origin = scale_keypoints(kpts0 + 0.5, s0) - 0.5
kpts1_origin = scale_keypoints(kpts1 + 0.5, s1) - 0.5
ret = {
"image0": image0.squeeze().cpu().numpy(),
"image1": image1.squeeze().cpu().numpy(),
"image0_orig": image_0,
"image1_orig": image_1,
"keypoints0": kpts0.cpu().numpy(),
"keypoints1": kpts1.cpu().numpy(),
"keypoints0_orig": kpts0_origin.cpu().numpy(),
"keypoints1_orig": kpts1_origin.cpu().numpy(),
"mkeypoints0": kpts0.cpu().numpy(),
"mkeypoints1": kpts1.cpu().numpy(),
"mkeypoints0_orig": kpts0_origin.cpu().numpy(),
"mkeypoints1_orig": kpts1_origin.cpu().numpy(),
"original_size0": np.array(image_0.shape[:2][::-1]),
"original_size1": np.array(image_1.shape[:2][::-1]),
"new_size0": np.array(image0.shape[-2:][::-1]),
"new_size1": np.array(image1.shape[-2:][::-1]),
"scale0": s0,
"scale1": s1,
}
if "mconf" in pred.keys():
ret["mconf"] = pred["mconf"].cpu().numpy()
elif "scores" in pred.keys(): # adapting loftr
ret["mconf"] = pred["scores"].cpu().numpy()
else:
ret["mconf"] = np.ones_like(kpts0.cpu().numpy()[:, 0])
if "lines0" in pred.keys() and "lines1" in pred.keys():
if "keypoints0" in pred.keys() and "keypoints1" in pred.keys():
kpts0, kpts1 = pred["keypoints0"], pred["keypoints1"]
kpts0_origin = scale_keypoints(kpts0 + 0.5, s0) - 0.5
kpts1_origin = scale_keypoints(kpts1 + 0.5, s1) - 0.5
kpts0_origin = kpts0_origin.cpu().numpy()
kpts1_origin = kpts1_origin.cpu().numpy()
else:
kpts0_origin, kpts1_origin = (
None,
None,
) # np.zeros([0]), np.zeros([0])
lines0, lines1 = pred["lines0"], pred["lines1"]
lines0_raw, lines1_raw = pred["raw_lines0"], pred["raw_lines1"]
lines0_raw = torch.from_numpy(lines0_raw.copy())
lines1_raw = torch.from_numpy(lines1_raw.copy())
lines0_raw = scale_lines(lines0_raw + 0.5, s0) - 0.5
lines1_raw = scale_lines(lines1_raw + 0.5, s1) - 0.5
lines0 = torch.from_numpy(lines0.copy())
lines1 = torch.from_numpy(lines1.copy())
lines0 = scale_lines(lines0 + 0.5, s0) - 0.5
lines1 = scale_lines(lines1 + 0.5, s1) - 0.5
ret = {
"image0_orig": image_0,
"image1_orig": image_1,
"line0": lines0_raw.cpu().numpy(),
"line1": lines1_raw.cpu().numpy(),
"line0_orig": lines0.cpu().numpy(),
"line1_orig": lines1.cpu().numpy(),
"line_keypoints0_orig": kpts0_origin,
"line_keypoints1_orig": kpts1_origin,
}
del pred
torch.cuda.empty_cache()
return ret
@torch.no_grad()
def main(
conf: Dict,
pairs: Path,
image_dir: Path,
export_dir: Optional[Path] = None,
matches: Optional[Path] = None, # out
features: Optional[Path] = None, # out
features_ref: Optional[Path] = None,
max_kps: Optional[int] = 8192,
overwrite: bool = False,
) -> Path:
logger.info(
"Extracting semi-dense features with configuration:"
f"\n{pprint.pformat(conf)}"
)
if features is None:
features = "feats_"
if isinstance(features, Path):
features_q = features
if matches is None:
raise ValueError(
"Either provide both features and matches as Path"
" or both as names."
)
else:
if export_dir is None:
raise ValueError(
"Provide an export_dir if features and matches"
f" are not file paths: {features}, {matches}."
)
features_q = Path(export_dir, f'{features}{conf["output"]}.h5')
if matches is None:
matches = Path(export_dir, f'{conf["output"]}_{pairs.stem}.h5')
if features_ref is None:
features_ref = []
elif isinstance(features_ref, list):
features_ref = list(features_ref)
elif isinstance(features_ref, Path):
features_ref = [features_ref]
else:
raise TypeError(str(features_ref))
match_and_assign(
conf,
pairs,
image_dir,
matches,
features_q,
features_ref,
max_kps,
overwrite,
)
return features_q, matches
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--pairs", type=Path, required=True)
parser.add_argument("--image_dir", type=Path, required=True)
parser.add_argument("--export_dir", type=Path, required=True)
parser.add_argument(
"--matches", type=Path, default=confs["loftr"]["output"]
)
parser.add_argument(
"--features", type=str, default="feats_" + confs["loftr"]["output"]
)
parser.add_argument(
"--conf", type=str, default="loftr", choices=list(confs.keys())
)
args = parser.parse_args()
main(
confs[args.conf],
args.pairs,
args.image_dir,
args.export_dir,
args.matches,
args.features,
)