diff --git a/hloc/extract_features.py b/hloc/extract_features.py index ab9456a8..53e31c5f 100644 --- a/hloc/extract_features.py +++ b/hloc/extract_features.py @@ -95,6 +95,14 @@ "resize_max": 1600, }, }, + "orb": { + "output": "feats-orb", + "model": {"name": "orb"}, + "preprocessing": { + "grayscale": True, + "resize_max": 1600, + }, + }, "sosnet": { "output": "feats-sosnet", "model": {"name": "dog", "descriptor": "sosnet"}, diff --git a/hloc/extractors/orb.py b/hloc/extractors/orb.py new file mode 100644 index 00000000..8294ba9a --- /dev/null +++ b/hloc/extractors/orb.py @@ -0,0 +1,100 @@ +import cv2 +import numpy as np +import torch + +from ..utils.base_model import BaseModel + +EPS = 1e-6 + + +class ORB(BaseModel): + default_conf = { + "options": { + "nfeatures": 5000, + "scaleFactor": 1.2, + "nlevels": 8, + "edgeThreshold": 31, + "firstLevel": 0, + "WTA_K": 2, + "scoreType": cv2.ORB_HARRIS_SCORE, # or cv2.ORB_FAST_SCORE + "patchSize": 31, + "fastThreshold": 20, + }, + "descriptor": "orb", + "max_keypoints": -1, + } + required_inputs = ["image"] + detection_noise = 1.0 + max_batch_size = 4096 + + def _init(self, conf): + if conf["descriptor"] != "orb": + raise ValueError(f'Unknown descriptor: {conf["descriptor"]}') + self.orb = None + self.dummy_param = torch.nn.Parameter(torch.empty(0)) + + def _make_orb(self): + opts = self.conf["options"] + + self.orb = cv2.ORB_create( + nfeatures=int(opts.get("nfeatures", 5000)), + scaleFactor=float(opts.get("scaleFactor", 1.2)), + nlevels=int(opts.get("nlevels", 8)), + edgeThreshold=int(opts.get("edgeThreshold", 31)), + firstLevel=int(opts.get("firstLevel", 0)), + WTA_K=int(opts.get("WTA_K", 2)), + scoreType=int(opts.get("scoreType", cv2.ORB_HARRIS_SCORE)), + patchSize=int(opts.get("patchSize", 31)), + fastThreshold=int(opts.get("fastThreshold", 20)), + ) + + def _forward(self, data): + image = data["image"] + + image_np = image.cpu().numpy()[0, 0] + assert image.shape[1] == 1, "ORB expects a single-channel image" + assert image_np.min() >= -EPS and image_np.max() <= 1 + EPS + + if self.orb is None: + self._make_orb() + + # Greyscale + img_u8 = np.clip(image_np * 255.0 + 0.5, 0, 255).astype(np.uint8) + + keypoints, descriptors = self.orb.detectAndCompute(img_u8, None) + + pts = np.array([kp.pt for kp in keypoints], dtype=np.float32) + sizes = np.array([kp.size for kp in keypoints], dtype=np.float32) + scales = sizes / 2.0 + angles = np.array([kp.angle for kp in keypoints], dtype=np.float32) + responses = np.array([kp.response for kp in keypoints], dtype=np.float32) + + # [N, 32] binary ORB + if descriptors is None: + descriptors = np.empty((0, 32), dtype=np.uint8) + + keypoints = torch.from_numpy(pts) + scales = torch.from_numpy(scales) + oris = torch.from_numpy(angles) + scores = torch.from_numpy(responses) + descriptors = torch.from_numpy(descriptors) # (N,32) uint8 + + if ( + self.conf["max_keypoints"] != -1 + and len(keypoints) > self.conf["max_keypoints"] + ): + k = int(self.conf["max_keypoints"]) + vals, idxs = torch.topk(scores, k) + keypoints = keypoints[idxs] + scales = scales[idxs] + oris = oris[idxs] + scores = vals + descriptors = descriptors[idxs] + + return { + "keypoints": keypoints[None], # [1, N, 2] (x, y) + "scales": scales[None], # [1, N] + "oris": oris[None], # [1, N] + "scores": scores[None], # [1, N] + "descriptors": descriptors.T[None], # [1, 32, N] + } diff --git a/hloc/match_features.py b/hloc/match_features.py index 679e81e9..ed8d13bb 100644 --- a/hloc/match_features.py +++ b/hloc/match_features.py @@ -81,6 +81,13 @@ "do_mutual_check": True, }, }, + "orb": { + "output": "matches-orb", + "model": { + "name": "orb_match", + "do_mutual_check": True, + }, + }, "adalam": { "output": "matches-adalam", "model": {"name": "adalam"}, diff --git a/hloc/matchers/orb_match.py b/hloc/matchers/orb_match.py new file mode 100644 index 00000000..ebe26174 --- /dev/null +++ b/hloc/matchers/orb_match.py @@ -0,0 +1,83 @@ +import cv2 +import numpy as np +import torch + +from ..utils.base_model import BaseModel + + +def tens_to_cv(x): + if isinstance(x, torch.Tensor): + x = x.detach().cpu() + if x.ndim == 3 and x.shape[0] == 1: + x = x.squeeze(0) + if x.ndim == 2 and x.shape[0] in (8, 16, 32, 64, 128, 256): + x = x.transpose(0, 1) + if isinstance(x, torch.Tensor): + x = x.contiguous().to(torch.uint8).numpy() + else: + x = np.ascontiguousarray(x, dtype=np.uint8) + return x # shape (N, 32) + + +class BinaryNearestNeighbor(BaseModel): + default_conf = { + "ratio_threshold": None, + "distance_threshold_bits": None, + "do_mutual_check": True, + } + + required_inputs = ["descriptors0", "scores0", "descriptors1", "scores1"] + + def _init(self, conf): + lut = torch.arange(256, dtype=torch.uint8) + lut = ( + (lut & 1) + + ((lut >> 1) & 1) + + ((lut >> 2) & 1) + + ((lut >> 3) & 1) + + ((lut >> 4) & 1) + + ((lut >> 5) & 1) + + ((lut >> 6) & 1) + + ((lut >> 7) & 1) + ) + self.register_buffer("_popcnt8", lut.to(torch.uint8), persistent=False) + + self.matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=False) + + def _forward(self, data): + d0 = data["descriptors0"] + d1 = data["descriptors1"] + + d0 = tens_to_cv(d0) + d1 = tens_to_cv(d1) + + D0, N0 = d0.shape + _, N1 = d1.shape + if N0 == 0 or N1 == 0: + device = d0.device + return { + "matches0": torch.full((1, N0), -1, dtype=torch.long, device=device), + "matching_scores0": torch.zeros( + (1, N0), dtype=torch.float32, device=device + ), + } + + matches = self.matcher.match(d0, d1) + + N0, Dbytes = d0.shape + Dbits = 8 * Dbytes + + matches0 = torch.full((N0,), -1, dtype=torch.long) + matching_scores0 = torch.zeros((N0,), dtype=torch.float32) + + for m in matches: + q = m.queryIdx + t = m.trainIdx + dist = float(m.distance) + matches0[q] = t + matching_scores0[q] = 1.0 - dist / Dbits + + matches0 = matches0.unsqueeze(0) # [1, N0] + matching_scores0 = matching_scores0.unsqueeze(0) + + return {"matches0": matches0, "matching_scores0": matching_scores0}