From 1f5f0aacdd886903c51e7de1b9ad3417c91fcd92 Mon Sep 17 00:00:00 2001 From: tavislocus Date: Tue, 2 Sep 2025 11:35:36 +0100 Subject: [PATCH 1/6] Added OpenCVs ORB feature extraction and matching --- hloc/extract_features.py | 8 ++++ hloc/extractors/orb.py | 96 ++++++++++++++++++++++++++++++++++++++ hloc/match_features.py | 7 +++ hloc/matchers/orb_match.py | 77 ++++++++++++++++++++++++++++++ 4 files changed, 188 insertions(+) create mode 100644 hloc/extractors/orb.py create mode 100644 hloc/matchers/orb_match.py 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..723f9afe --- /dev/null +++ b/hloc/extractors/orb.py @@ -0,0 +1,96 @@ +import numpy as np +import torch +import cv2 + +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..e955a44e --- /dev/null +++ b/hloc/matchers/orb_match.py @@ -0,0 +1,77 @@ +import torch +import cv2 +import numpy as np + +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} + + From dadd7cb7cb2fb782d990256805c62b0b8325b950 Mon Sep 17 00:00:00 2001 From: tavislocus Date: Tue, 2 Sep 2025 11:59:48 +0100 Subject: [PATCH 2/6] fixed some formatting issues --- hloc/extractors/orb.py | 22 ++++++++++++---------- hloc/matchers/orb_match.py | 15 +++++++-------- 2 files changed, 19 insertions(+), 18 deletions(-) diff --git a/hloc/extractors/orb.py b/hloc/extractors/orb.py index 723f9afe..b4a59ba8 100644 --- a/hloc/extractors/orb.py +++ b/hloc/extractors/orb.py @@ -6,6 +6,7 @@ EPS = 1e-6 + class ORB(BaseModel): default_conf = { "options": { @@ -19,7 +20,7 @@ class ORB(BaseModel): "patchSize": 31, "fastThreshold": 20, }, - "descriptor": "orb", + "descriptor": "orb", "max_keypoints": -1, } required_inputs = ["image"] @@ -29,7 +30,7 @@ class ORB(BaseModel): def _init(self, conf): if conf["descriptor"] != "orb": raise ValueError(f'Unknown descriptor: {conf["descriptor"]}') - self.orb = None + self.orb = None self.dummy_param = torch.nn.Parameter(torch.empty(0)) def _make_orb(self): @@ -49,7 +50,7 @@ def _make_orb(self): 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 @@ -76,9 +77,10 @@ def _forward(self, data): scales = torch.from_numpy(scales) oris = torch.from_numpy(angles) scores = torch.from_numpy(responses) - descriptors = torch.from_numpy(descriptors) # (N,32) uint8 + descriptors = torch.from_numpy(descriptors) # (N,32) uint8 - if self.conf["max_keypoints"] != -1 and len(keypoints) > self.conf["max_keypoints"]: + 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] @@ -88,9 +90,9 @@ def _forward(self, data): 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] + "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/matchers/orb_match.py b/hloc/matchers/orb_match.py index e955a44e..88eea1f1 100644 --- a/hloc/matchers/orb_match.py +++ b/hloc/matchers/orb_match.py @@ -21,8 +21,8 @@ def tens_to_cv(x): class BinaryNearestNeighbor(BaseModel): default_conf = { - "ratio_threshold": None, - "distance_threshold_bits": None, + "ratio_threshold": None, + "distance_threshold_bits": None, "do_mutual_check": True, } @@ -40,7 +40,7 @@ def _init(self, conf): def _forward(self, data): d0 = data["descriptors0"] d1 = data["descriptors1"] - + d0 = tens_to_cv(d0) d1 = tens_to_cv(d1) @@ -50,7 +50,8 @@ def _forward(self, data): 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), + "matching_scores0": torch.zeros((1, N0), dtype=torch.float32, + device=device), } matches = self.matcher.match(d0, d1) @@ -68,10 +69,8 @@ def _forward(self, data): matches0[q] = t matching_scores0[q] = 1.0 - dist / Dbits - matches0 = matches0.unsqueeze(0) # [1, N0] + matches0 = matches0.unsqueeze(0) # [1, N0] matching_scores0 = matching_scores0.unsqueeze(0) - return {"matches0": matches0, + return {"matches0": matches0, "matching_scores0": matching_scores0} - - From a04e113bcdc53845b188938656e798e652c5267e Mon Sep 17 00:00:00 2001 From: tavislocus Date: Tue, 2 Sep 2025 12:00:49 +0100 Subject: [PATCH 3/6] formatting issues --- hloc/extractors/orb.py | 4 ++-- hloc/matchers/orb_match.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/hloc/extractors/orb.py b/hloc/extractors/orb.py index b4a59ba8..89018e49 100644 --- a/hloc/extractors/orb.py +++ b/hloc/extractors/orb.py @@ -65,7 +65,7 @@ def _forward(self, data): 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 + 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) @@ -92,7 +92,7 @@ def _forward(self, data): return { "keypoints": keypoints[None], # [1, N, 2] (x, y) "scales": scales[None], # [1, N] - "oris": oris[None], # [1, N] + "oris": oris[None], # [1, N] "scores": scores[None], # [1, N] "descriptors": descriptors.T[None], # [1, 32, N] } diff --git a/hloc/matchers/orb_match.py b/hloc/matchers/orb_match.py index 88eea1f1..4ce18f89 100644 --- a/hloc/matchers/orb_match.py +++ b/hloc/matchers/orb_match.py @@ -50,7 +50,7 @@ def _forward(self, data): device = d0.device return { "matches0": torch.full((1, N0), -1, dtype=torch.long, device=device), - "matching_scores0": torch.zeros((1, N0), dtype=torch.float32, + "matching_scores0": torch.zeros((1, N0), dtype=torch.float32, device=device), } From 19e15ea6c80dc7758d7eef809470bb9a5f4cae7d Mon Sep 17 00:00:00 2001 From: tavislocus Date: Tue, 2 Sep 2025 12:02:16 +0100 Subject: [PATCH 4/6] rearrange imports --- hloc/extractors/orb.py | 2 +- hloc/matchers/orb_match.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/hloc/extractors/orb.py b/hloc/extractors/orb.py index 89018e49..8cc5bbcf 100644 --- a/hloc/extractors/orb.py +++ b/hloc/extractors/orb.py @@ -1,6 +1,6 @@ +import cv2 import numpy as np import torch -import cv2 from ..utils.base_model import BaseModel diff --git a/hloc/matchers/orb_match.py b/hloc/matchers/orb_match.py index 4ce18f89..5a885f90 100644 --- a/hloc/matchers/orb_match.py +++ b/hloc/matchers/orb_match.py @@ -1,6 +1,6 @@ -import torch import cv2 import numpy as np +import torch from ..utils.base_model import BaseModel From 7533deec7c8d11eec9ef3ef95d3dd11c66353810 Mon Sep 17 00:00:00 2001 From: tavislocus Date: Tue, 2 Sep 2025 12:05:37 +0100 Subject: [PATCH 5/6] more formatting issues --- hloc/extractors/orb.py | 14 ++++++++------ hloc/matchers/orb_match.py | 5 +++-- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/hloc/extractors/orb.py b/hloc/extractors/orb.py index 8cc5bbcf..8294ba9a 100644 --- a/hloc/extractors/orb.py +++ b/hloc/extractors/orb.py @@ -79,8 +79,10 @@ def _forward(self, data): 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"]): + 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] @@ -90,9 +92,9 @@ def _forward(self, data): 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] + "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/matchers/orb_match.py b/hloc/matchers/orb_match.py index 5a885f90..18f46345 100644 --- a/hloc/matchers/orb_match.py +++ b/hloc/matchers/orb_match.py @@ -50,8 +50,9 @@ def _forward(self, data): 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), + "matching_scores0": torch.zeros( + (1, N0), dtype=torch.float32, device=device + ), } matches = self.matcher.match(d0, d1) From 643f6b5a07d7047c01ff5c4979df619d43d77882 Mon Sep 17 00:00:00 2001 From: tavislocus Date: Tue, 2 Sep 2025 12:07:48 +0100 Subject: [PATCH 6/6] More formatting --- hloc/matchers/orb_match.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/hloc/matchers/orb_match.py b/hloc/matchers/orb_match.py index 18f46345..ebe26174 100644 --- a/hloc/matchers/orb_match.py +++ b/hloc/matchers/orb_match.py @@ -26,13 +26,20 @@ class BinaryNearestNeighbor(BaseModel): "do_mutual_check": True, } - required_inputs = ['descriptors0', 'scores0', - 'descriptors1', 'scores1'] + 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) + 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) @@ -73,5 +80,4 @@ def _forward(self, data): matches0 = matches0.unsqueeze(0) # [1, N0] matching_scores0 = matching_scores0.unsqueeze(0) - return {"matches0": matches0, - "matching_scores0": matching_scores0} + return {"matches0": matches0, "matching_scores0": matching_scores0}