import warnings import cv2 import numpy as np import torch from kornia.color import rgb_to_grayscale from omegaconf import OmegaConf from packaging import version try: import pycolmap except ImportError: pycolmap = None from hloc import logger from ..utils.base_model import BaseModel def filter_dog_point( points, scales, angles, image_shape, nms_radius, scores=None ): h, w = image_shape ij = np.round(points - 0.5).astype(int).T[::-1] # Remove duplicate points (identical coordinates). # Pick highest scale or score s = scales if scores is None else scores buffer = np.zeros((h, w)) np.maximum.at(buffer, tuple(ij), s) keep = np.where(buffer[tuple(ij)] == s)[0] # Pick lowest angle (arbitrary). ij = ij[:, keep] buffer[:] = np.inf o_abs = np.abs(angles[keep]) np.minimum.at(buffer, tuple(ij), o_abs) mask = buffer[tuple(ij)] == o_abs ij = ij[:, mask] keep = keep[mask] if nms_radius > 0: # Apply NMS on the remaining points buffer[:] = 0 buffer[tuple(ij)] = s[keep] # scores or scale local_max = torch.nn.functional.max_pool2d( torch.from_numpy(buffer).unsqueeze(0), kernel_size=nms_radius * 2 + 1, stride=1, padding=nms_radius, ).squeeze(0) is_local_max = buffer == local_max.numpy() keep = keep[is_local_max[tuple(ij)]] return keep def sift_to_rootsift(x: torch.Tensor, eps=1e-6) -> torch.Tensor: x = torch.nn.functional.normalize(x, p=1, dim=-1, eps=eps) x.clip_(min=eps).sqrt_() return torch.nn.functional.normalize(x, p=2, dim=-1, eps=eps) def run_opencv_sift(features: cv2.Feature2D, image: np.ndarray) -> np.ndarray: """ Detect keypoints using OpenCV Detector. Optionally, perform description. Args: features: OpenCV based keypoints detector and descriptor image: Grayscale image of uint8 data type Returns: keypoints: 1D array of detected cv2.KeyPoint scores: 1D array of responses descriptors: 1D array of descriptors """ detections, descriptors = features.detectAndCompute(image, None) points = np.array([k.pt for k in detections], dtype=np.float32) scores = np.array([k.response for k in detections], dtype=np.float32) scales = np.array([k.size for k in detections], dtype=np.float32) angles = np.deg2rad( np.array([k.angle for k in detections], dtype=np.float32) ) return points, scores, scales, angles, descriptors class SIFT(BaseModel): default_conf = { "rootsift": True, "nms_radius": 0, # None to disable filtering entirely. "max_keypoints": 4096, "backend": "opencv", # in {opencv, pycolmap, pycolmap_cpu, pycolmap_cuda} "detection_threshold": 0.0066667, # from COLMAP "edge_threshold": 10, "first_octave": -1, # only used by pycolmap, the default of COLMAP "num_octaves": 4, } required_data_keys = ["image"] def _init(self, conf): self.conf = OmegaConf.create(self.conf) backend = self.conf.backend if backend.startswith("pycolmap"): if pycolmap is None: raise ImportError( "Cannot find module pycolmap: install it with pip" "or use backend=opencv." ) options = { "peak_threshold": self.conf.detection_threshold, "edge_threshold": self.conf.edge_threshold, "first_octave": self.conf.first_octave, "num_octaves": self.conf.num_octaves, "normalization": pycolmap.Normalization.L2, # L1_ROOT is buggy. } device = ( "auto" if backend == "pycolmap" else backend.replace("pycolmap_", "") ) if ( backend == "pycolmap_cpu" or not pycolmap.has_cuda ) and pycolmap.__version__ < "0.5.0": warnings.warn( "The pycolmap CPU SIFT is buggy in version < 0.5.0, " "consider upgrading pycolmap or use the CUDA version.", stacklevel=1, ) else: options["max_num_features"] = self.conf.max_keypoints self.sift = pycolmap.Sift(options=options, device=device) elif backend == "opencv": self.sift = cv2.SIFT_create( contrastThreshold=self.conf.detection_threshold, nfeatures=self.conf.max_keypoints, edgeThreshold=self.conf.edge_threshold, nOctaveLayers=self.conf.num_octaves, ) else: backends = {"opencv", "pycolmap", "pycolmap_cpu", "pycolmap_cuda"} raise ValueError( f"Unknown backend: {backend} not in " f"{{{','.join(backends)}}}." ) logger.info("Load SIFT model done.") def extract_single_image(self, image: torch.Tensor): image_np = image.cpu().numpy().squeeze(0) if self.conf.backend.startswith("pycolmap"): if version.parse(pycolmap.__version__) >= version.parse("0.5.0"): detections, descriptors = self.sift.extract(image_np) scores = None # Scores are not exposed by COLMAP anymore. else: detections, scores, descriptors = self.sift.extract(image_np) keypoints = detections[:, :2] # Keep only (x, y). scales, angles = detections[:, -2:].T if scores is not None and ( self.conf.backend == "pycolmap_cpu" or not pycolmap.has_cuda ): # Set the scores as a combination of abs. response and scale. scores = np.abs(scores) * scales elif self.conf.backend == "opencv": # TODO: Check if opencv keypoints are already in corner convention keypoints, scores, scales, angles, descriptors = run_opencv_sift( self.sift, (image_np * 255.0).astype(np.uint8) ) pred = { "keypoints": keypoints, "scales": scales, "oris": angles, "descriptors": descriptors, } if scores is not None: pred["scores"] = scores # sometimes pycolmap returns points outside the image. We remove them if self.conf.backend.startswith("pycolmap"): is_inside = ( pred["keypoints"] + 0.5 < np.array([image_np.shape[-2:][::-1]]) ).all(-1) pred = {k: v[is_inside] for k, v in pred.items()} if self.conf.nms_radius is not None: keep = filter_dog_point( pred["keypoints"], pred["scales"], pred["oris"], image_np.shape, self.conf.nms_radius, scores=pred.get("scores"), ) pred = {k: v[keep] for k, v in pred.items()} pred = {k: torch.from_numpy(v) for k, v in pred.items()} if scores is not None: # Keep the k keypoints with highest score num_points = self.conf.max_keypoints if num_points is not None and len(pred["keypoints"]) > num_points: indices = torch.topk(pred["scores"], num_points).indices pred = {k: v[indices] for k, v in pred.items()} return pred def _forward(self, data: dict) -> dict: image = data["image"] if image.shape[1] == 3: image = rgb_to_grayscale(image) device = image.device image = image.cpu() pred = [] for k in range(len(image)): img = image[k] if "image_size" in data.keys(): # avoid extracting points in padded areas w, h = data["image_size"][k] img = img[:, :h, :w] p = self.extract_single_image(img) pred.append(p) pred = { k: torch.stack([p[k] for p in pred], 0).to(device) for k in pred[0] } if self.conf.rootsift: pred["descriptors"] = sift_to_rootsift(pred["descriptors"]) pred["descriptors"] = pred["descriptors"].permute(0, 2, 1) pred["keypoint_scores"] = pred["scores"].clone() return pred