From a8c70225b57ac2b3dcd568a898965059948034e9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=90=95=E5=A5=89=E5=9D=A4?= <50992921+lvfengkun@users.noreply.github.com> Date: Mon, 24 Jun 2024 10:02:21 +0800 Subject: [PATCH 1/5] Add files via upload --- lightglue/__init__.py | 7 +- lightglue/sift.py | 220 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 224 insertions(+), 3 deletions(-) create mode 100644 lightglue/sift.py diff --git a/lightglue/__init__.py b/lightglue/__init__.py index 97ad123..6aee2bf 100644 --- a/lightglue/__init__.py +++ b/lightglue/__init__.py @@ -1,4 +1,5 @@ -from .lightglue import LightGlue -from .superpoint import SuperPoint -from .disk import DISK +from .lightglue import LightGlue +from .superpoint import SuperPoint +from .disk import DISK +from .sift import SIFT # noqa from .utils import match_pair \ No newline at end of file diff --git a/lightglue/sift.py b/lightglue/sift.py new file mode 100644 index 0000000..7a53ecd --- /dev/null +++ b/lightglue/sift.py @@ -0,0 +1,220 @@ +import warnings + +import cv2 +import numpy as np +import torch +from kornia.color import rgb_to_grayscale +from packaging import version +import time +try: + import pycolmap +except ImportError: + pycolmap = None + +from .utils import Extractor + + +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 + """ + #start_time1=time.time() + 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)) + #end_time = time.time() + #elapsed_time = end_time - start_time1 + #print(f"sift特征提取时间：{elapsed_time}") + return points, scores, scales, angles, descriptors + + +class SIFT(Extractor): + default_conf = { + "rootsift": True, #True + "nms_radius": 0, # None to disable filtering entirely. + "max_num_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, + } + + preprocess_conf = { + "resize": 1024, + } + + required_data_keys = ["image"] + + def __init__(self, **conf): + super().__init__(**conf) # Update with default configuration. + 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_num_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_num_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)}}}." + ) + + 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["keypoint_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("keypoint_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_num_keypoints + if num_points is not None and len(pred["keypoints"]) > num_points: + indices = torch.topk(pred["keypoint_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"]) + return pred From c5830f3dcc85c4fa8eaedbe01ee8839d1f1c09e0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=90=95=E5=A5=89=E5=9D=A4?= <50992921+lvfengkun@users.noreply.github.com> Date: Mon, 24 Jun 2024 10:03:28 +0800 Subject: [PATCH 2/5] Add files via upload --- lightglue_onnx/__init__.py | 1 + lightglue_onnx/sift.py | 220 +++++++++++++++++++++++++++++++++++++ lightglue_onnx/utils.py | 50 +++++++++ 3 files changed, 271 insertions(+) create mode 100644 lightglue_onnx/sift.py diff --git a/lightglue_onnx/__init__.py b/lightglue_onnx/__init__.py index a3b8930..a467e01 100644 --- a/lightglue_onnx/__init__.py +++ b/lightglue_onnx/__init__.py @@ -2,4 +2,5 @@ from .end2end import LightGlueEnd2End from .lightglue import LightGlue from .superpoint import SuperPoint +from .sift import SIFT # noqa from .utils import match_pair diff --git a/lightglue_onnx/sift.py b/lightglue_onnx/sift.py new file mode 100644 index 0000000..7a53ecd --- /dev/null +++ b/lightglue_onnx/sift.py @@ -0,0 +1,220 @@ +import warnings + +import cv2 +import numpy as np +import torch +from kornia.color import rgb_to_grayscale +from packaging import version +import time +try: + import pycolmap +except ImportError: + pycolmap = None + +from .utils import Extractor + + +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 + """ + #start_time1=time.time() + 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)) + #end_time = time.time() + #elapsed_time = end_time - start_time1 + #print(f"sift特征提取时间：{elapsed_time}") + return points, scores, scales, angles, descriptors + + +class SIFT(Extractor): + default_conf = { + "rootsift": True, #True + "nms_radius": 0, # None to disable filtering entirely. + "max_num_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, + } + + preprocess_conf = { + "resize": 1024, + } + + required_data_keys = ["image"] + + def __init__(self, **conf): + super().__init__(**conf) # Update with default configuration. + 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_num_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_num_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)}}}." + ) + + 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["keypoint_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("keypoint_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_num_keypoints + if num_points is not None and len(pred["keypoints"]) > num_points: + indices = torch.topk(pred["keypoint_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"]) + return pred diff --git a/lightglue_onnx/utils.py b/lightglue_onnx/utils.py index 1bbcaeb..74b94d1 100644 --- a/lightglue_onnx/utils.py +++ b/lightglue_onnx/utils.py @@ -104,3 +104,53 @@ def match_pair(extractor, matcher, image0, image1, scales0=None, scales1=None): valid = matches0 > -1 matches = torch.stack([torch.where(valid)[0], matches0[valid]], -1) return {**pred, "matches": matches, "matching_scores": mscores0[valid]} + + +from types import SimpleNamespace +from typing import Callable, List, Optional, Tuple, Union +import kornia +class ImagePreprocessor: + default_conf = { + "resize": None, # target edge length, None for no resizing + "side": "long", + "interpolation": "bilinear", + "align_corners": None, + "antialias": True, + } + + def __init__(self, **conf) -> None: + super().__init__() + self.conf = {**self.default_conf, **conf} + self.conf = SimpleNamespace(**self.conf) + + def __call__(self, img: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Resize and preprocess an image, return image and resize scale""" + h, w = img.shape[-2:] + if self.conf.resize is not None: + img = kornia.geometry.transform.resize( + img, + self.conf.resize, + side=self.conf.side, + antialias=self.conf.antialias, + align_corners=self.conf.align_corners, + ) + scale = torch.Tensor([img.shape[-1] / w, img.shape[-2] / h]).to(img) + return img, scale + +class Extractor(torch.nn.Module): + def __init__(self, **conf): + super().__init__() + self.conf = SimpleNamespace(**{**self.default_conf, **conf}) + + @torch.no_grad() + def extract(self, img: torch.Tensor, **conf) -> dict: + """Perform extraction with online resizing""" + if img.dim() == 3: + img = img[None] # add batch dim + assert img.dim() == 4 and img.shape[0] == 1 + shape = img.shape[-2:][::-1] + img, scales = ImagePreprocessor(**{**self.preprocess_conf, **conf})(img) + feats = self.forward({"image": img}) + feats["image_size"] = torch.tensor(shape)[None].to(img).float() + feats["keypoints"] = (feats["keypoints"] + 0.5) / scales[None] - 0.5 + return feats From 9a4d9c18f3e6c226e6cfe5fee5a48de1cfafa58c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=90=95=E5=A5=89=E5=9D=A4?= <50992921+lvfengkun@users.noreply.github.com> Date: Mon, 24 Jun 2024 10:04:04 +0800 Subject: [PATCH 3/5] Add files via upload --- onnx_runner/__init__.py | 4 +- onnx_runner/lightglue_sift.py | 92 ++++++++++++++ onnx_runner/sift.py | 220 ++++++++++++++++++++++++++++++++++ 3 files changed, 315 insertions(+), 1 deletion(-) create mode 100644 onnx_runner/lightglue_sift.py create mode 100644 onnx_runner/sift.py diff --git a/onnx_runner/__init__.py b/onnx_runner/__init__.py index 1aa40a7..d529e21 100644 --- a/onnx_runner/__init__.py +++ b/onnx_runner/__init__.py @@ -1,2 +1,4 @@ from .lightglue import LightGlueRunner -from .utils import load_image, rgb_to_grayscale +from .lightglue_sift import LightGlueRunner_SIFT +from .utils import load_image, rgb_to_grayscale, load_image_sift +from .sift import SIFT # noqa \ No newline at end of file diff --git a/onnx_runner/lightglue_sift.py b/onnx_runner/lightglue_sift.py new file mode 100644 index 0000000..f1a52a8 --- /dev/null +++ b/onnx_runner/lightglue_sift.py @@ -0,0 +1,92 @@ +import numpy as np +import onnxruntime as ort +import time +from .sift import SIFT +from lightglue_onnx.end2end import normalize_keypoints +import torch +class LightGlueRunner_SIFT: + def __init__( + self, + lightglue_path: str, + extractor_path=None, + providers=["CUDAExecutionProvider", "CPUExecutionProvider"], + ): + self.extractor = SIFT(max_num_keypoints=128).eval() + sess_options = ort.SessionOptions() + self.lightglue = ort.InferenceSession( + lightglue_path, sess_options=sess_options, providers=providers + ) + + # Check for invalid models. + lightglue_inputs = [i.name for i in self.lightglue.get_inputs()] + if self.extractor is not None and "image0" in lightglue_inputs: + raise TypeError( + f"The specified LightGlue model at {lightglue_path} is end-to-end. Please do not pass the extractor_path argument." + ) + elif self.extractor is None and "image0" not in lightglue_inputs: + raise TypeError( + f"The specified LightGlue model at {lightglue_path} is not end-to-end. Please pass the extractor_path argument." + ) + + def run(self, image0: np.ndarray, image1: np.ndarray, scales0, scales1): + if self.extractor is None: + kpts0, kpts1, matches0, mscores0 = self.lightglue.run( + None, + { + "image0": image0, + "image1": image1, + }, + ) + m_kpts0, m_kpts1 = self.post_process( + kpts0, kpts1, matches0, scales0, scales1 + ) + return m_kpts0, m_kpts1 + else: + feats0 = self.extractor.extract(image0[None]) + feats1 = self.extractor.extract(image1[None]) + kpts0, scores0, desc0 = feats0["keypoints"],feats0["keypoint_scores"],feats0["descriptors"] + kpts1, scores1, desc1 = feats1["keypoints"],feats1["keypoint_scores"],feats1["descriptors"] + kpts0_cpoy=kpts0.clone() + kpts1_cpoy=kpts1.clone() + kpts0 = normalize_keypoints(kpts0, image0.shape[1], image0.shape[2]) + kpts1 = normalize_keypoints(kpts1, image1.shape[1], image1.shape[2]) + kpts0 = torch.cat( + [kpts0] + [feats0[k].unsqueeze(-1) for k in ("scales", "oris")], -1 + ) + kpts1 = torch.cat( + [kpts1] + [feats1[k].unsqueeze(-1) for k in ("scales", "oris")], -1 + ) + + start_time1 = time.time() + matches0, mscores0 = self.lightglue.run( + None, + { + "kpts0": kpts0.numpy(), + "kpts1": kpts1.numpy(), + "desc0": desc0.numpy(), + "desc1": desc1.numpy(), + }, + ) + end_time = time.time() + elapsed_time = end_time - start_time1 + print(f"程序运行时间: {elapsed_time} 秒") + m_kpts0, m_kpts1 = self.post_process( + kpts0_cpoy, kpts1_cpoy, matches0, scales0, scales1 + ) + return m_kpts0, m_kpts1 + + @staticmethod + def normalize_keypoints(kpts: np.ndarray, h: int, w: int) -> np.ndarray: + size = np.array([w, h]) + shift = size / 2 + scale = size.max() / 2 + kpts = (kpts - shift) / scale + return kpts.astype(np.float32) + + @staticmethod + def post_process(kpts0, kpts1, matches, scales0, scales1): + kpts0 = (kpts0 + 0.5) / scales0 - 0.5 + kpts1 = (kpts1 + 0.5) / scales1 - 0.5 + # create match indices + m_kpts0, m_kpts1 = kpts0[0][matches[..., 0]], kpts1[0][matches[..., 1]] + return m_kpts0, m_kpts1 diff --git a/onnx_runner/sift.py b/onnx_runner/sift.py new file mode 100644 index 0000000..7a53ecd --- /dev/null +++ b/onnx_runner/sift.py @@ -0,0 +1,220 @@ +import warnings + +import cv2 +import numpy as np +import torch +from kornia.color import rgb_to_grayscale +from packaging import version +import time +try: + import pycolmap +except ImportError: + pycolmap = None + +from .utils import Extractor + + +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 + """ + #start_time1=time.time() + 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)) + #end_time = time.time() + #elapsed_time = end_time - start_time1 + #print(f"sift特征提取时间：{elapsed_time}") + return points, scores, scales, angles, descriptors + + +class SIFT(Extractor): + default_conf = { + "rootsift": True, #True + "nms_radius": 0, # None to disable filtering entirely. + "max_num_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, + } + + preprocess_conf = { + "resize": 1024, + } + + required_data_keys = ["image"] + + def __init__(self, **conf): + super().__init__(**conf) # Update with default configuration. + 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_num_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_num_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)}}}." + ) + + 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["keypoint_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("keypoint_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_num_keypoints + if num_points is not None and len(pred["keypoints"]) > num_points: + indices = torch.topk(pred["keypoint_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"]) + return pred From a2474f270dcf2d107a2e143450f7c0bf134bcbc2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=90=95=E5=A5=89=E5=9D=A4?= <50992921+lvfengkun@users.noreply.github.com> Date: Mon, 24 Jun 2024 10:04:49 +0800 Subject: [PATCH 4/5] Add files via upload --- export_sift.py | 214 +++++++++++++++++++++++++++++++++++++++++++++++++ infer_sift.py | 144 +++++++++++++++++++++++++++++++++ 2 files changed, 358 insertions(+) create mode 100644 export_sift.py create mode 100644 infer_sift.py diff --git a/export_sift.py b/export_sift.py new file mode 100644 index 0000000..b140b6a --- /dev/null +++ b/export_sift.py @@ -0,0 +1,214 @@ +import argparse +from typing import List + +import torch + +from lightglue_onnx import DISK, LightGlue, LightGlueEnd2End, SuperPoint, SIFT +from lightglue_onnx.end2end import normalize_keypoints +from lightglue_onnx.utils import load_image, rgb_to_grayscale + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "--img_size", + nargs="+", + type=int, + default=512, + required=False, + help="Sample image size for ONNX tracing. If a single integer is given, resize the longer side of the image to this value. Otherwise, please provide two integers (height width).", + ) + parser.add_argument( + "--extractor_type", + type=str, + default="sift", + choices=["superpoint", "disk", "sift"], + required=False, + help="Type of feature extractor. Supported extractors are 'superpoint' and 'disk'. Defaults to 'superpoint'.", + ) + parser.add_argument( + "--extractor_path", + type=str, + default=None, + required=False, + help="Path to save the feature extractor ONNX model.", + ) + parser.add_argument( + "--lightglue_path", + type=str, + default=None, + required=False, + help="Path to save the LightGlue ONNX model.", + ) + parser.add_argument( + "--end2end", + action="store_true", + help="Whether to export an end-to-end pipeline instead of individual models.", + ) + parser.add_argument( + "--dynamic", action="store_true", help="Whether to allow dynamic image sizes." + ) + + # Extractor-specific args: + parser.add_argument( + "--max_num_keypoints", + type=int, + default=None, + required=False, + help="Maximum number of keypoints outputted by the extractor.", + ) + + return parser.parse_args() + + +def export_onnx( + img_size=512, + extractor_type="superpoint", + extractor_path=None, + lightglue_path=None, + img0_path="assets/sacre_coeur1.jpg", + img1_path="assets/sacre_coeur2.jpg", + end2end=False, + dynamic=False, + max_num_keypoints=128,#max_num_keypoints +): + # Handle args + if isinstance(img_size, List) and len(img_size) == 1: + img_size = img_size[0] + + if extractor_path is not None and end2end: + raise ValueError( + "Extractor will be combined with LightGlue when exporting end-to-end model." + ) + if extractor_path is None: + extractor_path = f"weights/{extractor_type}.onnx" + if max_num_keypoints is not None: + extractor_path = extractor_path.replace( + ".onnx", f"_{max_num_keypoints}.onnx" + ) + + if lightglue_path is None: + lightglue_path = ( + f"weights/{extractor_type}_lightglue" + f"{'_end2end' if end2end else ''}" + ".onnx" + ) + + # Sample images for tracing + image0, scales0 = load_image(img0_path, resize=img_size) + image1, scales1 = load_image(img1_path, resize=img_size) + # Models + extractor_type = extractor_type.lower() + if extractor_type == "superpoint": + # SuperPoint works on grayscale images. + image0 = rgb_to_grayscale(image0) + image1 = rgb_to_grayscale(image1) + extractor = SuperPoint(max_num_keypoints=max_num_keypoints).eval() + lightglue = LightGlue(extractor_type).eval() + elif extractor_type == "disk": + extractor = DISK(max_num_keypoints=max_num_keypoints).eval() + lightglue = LightGlue(extractor_type).eval() + elif extractor_type == "sift": + extractor = SIFT(max_num_keypoints=128).eval() + lightglue = LightGlue(extractor_type).eval() + else: + raise NotImplementedError( + f"LightGlue has not been trained on {extractor_type} features." + ) + + # ONNX Export + if end2end: + pipeline = LightGlueEnd2End(extractor, lightglue).eval() + + dynamic_axes = { + "kpts0": {1: "num_keypoints0"}, + "kpts1": {1: "num_keypoints1"}, + "matches0": {0: "num_matches0"}, + "mscores0": {0: "num_matches0"}, + } + if dynamic: + dynamic_axes.update( + { + "image0": {2: "height0", 3: "width0"}, + "image1": {2: "height1", 3: "width1"}, + } + ) + + torch.onnx.export( + pipeline, + (image0[None], image1[None]), + lightglue_path, + input_names=["image0", "image1"], + output_names=[ + "kpts0", + "kpts1", + "matches0", + "mscores0", + ], + opset_version=17, + dynamic_axes=dynamic_axes, + ) + else: + # Export Extractor + dynamic_axes = { + "keypoints": {1: "num_keypoints"}, + "scores": {1: "num_keypoints"}, + "descriptors": {1: "num_keypoints"}, + } + if dynamic: + dynamic_axes.update({"image": {2: "height", 3: "width"}}) + else: + print( + f"WARNING: Exporting without --dynamic implies that the {extractor_type} extractor's input image size will be locked to {image0.shape[-2:]}" + ) + extractor_path = extractor_path.replace( + ".onnx", f"_{image0.shape[-2]}x{image0.shape[-1]}.onnx" + ) + + '''torch.onnx.export( + extractor, + image0[None], + extractor_path, + input_names=["image"], + output_names=["keypoints", "scores", "descriptors"], + opset_version=17, + dynamic_axes=dynamic_axes, + )''' + feats0 = extractor.extract(image0[None]) + feats1 = extractor.extract(image1[None]) + + # Export LightGlue + #feats0, feats1 = extractor(image0[None]), extractor(image1[None]) + kpts0, scores0, desc0 = feats0["keypoints"],feats0["keypoint_scores"],feats0["descriptors"] + kpts1, scores1, desc1 = feats1["keypoints"],feats1["keypoint_scores"],feats1["descriptors"] + kpts0 = normalize_keypoints(kpts0, image0.shape[1], image0.shape[2]) + kpts1 = normalize_keypoints(kpts1, image1.shape[1], image1.shape[2]) + kpts0 = torch.cat( + [kpts0] + [feats0[k].unsqueeze(-1) for k in ("scales", "oris")], -1 + ) + kpts1 = torch.cat( + [kpts1] + [feats1[k].unsqueeze(-1) for k in ("scales", "oris")], -1 + ) + + + torch.onnx.export( + lightglue, + (kpts0, kpts1, desc0, desc1), + lightglue_path, + input_names=["kpts0", "kpts1", "desc0", "desc1"], + output_names=["matches0", "mscores0"], + opset_version=17, + dynamic_axes={ + "kpts0": {1: "num_keypoints0"}, + "kpts1": {1: "num_keypoints1"}, + "desc0": {1: "num_keypoints0"}, + "desc1": {1: "num_keypoints1"}, + "matches0": {0: "num_matches0"}, + "mscores0": {0: "num_matches0"}, + }, + ) + + +if __name__ == "__main__": + args = parse_args() + export_onnx(**vars(args)) diff --git a/infer_sift.py b/infer_sift.py new file mode 100644 index 0000000..ef7e780 --- /dev/null +++ b/infer_sift.py @@ -0,0 +1,144 @@ +import argparse +from typing import List +import time +from onnx_runner import LightGlueRunner, load_image, rgb_to_grayscale, viz2d,LightGlueRunner_SIFT, load_image_sift + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "--img_paths", + nargs=2, + default=["assets/remapR1.bmp", "assets/remapL1.bmp"], + required=False, + type=str, + ) + parser.add_argument( + "--lightglue_path", + type=str, + default="weights/sift_lightglue_fused_cpu.onnx", + required=False, + help="Path to the LightGlue ONNX model or end-to-end LightGlue pipeline.", + ) + parser.add_argument( + "--extractor_type", + type=str, + default="sift", + choices=["superpoint", "disk","sift"], + required=False, + help="Type of feature extractor. Supported extractors are 'superpoint' and 'disk'.", + ) + parser.add_argument( + "--extractor_path", + type=str, + default=" ", + required=False, + help="Path to the feature extractor ONNX model. If this argument is not provided, it is assumed that lightglue_path refers to an end-to-end model.", + ) + parser.add_argument( + "--img_size", + nargs="+", + type=int, + default=512, + required=False, + help="Sample image size for ONNX tracing. If a single integer is given, resize the longer side of the images to this value. Otherwise, please provide two integers (height width) to resize both images to this size, or four integers (height width height width).", + ) + parser.add_argument( + "--trt", + action="store_true", + help="Whether to use TensorRT (experimental).", + ) + parser.add_argument( + "--viz", action="store_true", default=True,help="Whether to visualize the results." + ) + return parser.parse_args() + + +def infer( + img_paths: List[str], + lightglue_path: str, + extractor_type: str, + extractor_path=None, + img_size=512, + trt=False, + viz=False, +): + # Handle args + img0_path = img_paths[0] + img1_path = img_paths[1] + if isinstance(img_size, List): + if len(img_size) == 1: + size0 = size1 = img_size[0] + elif len(img_size) == 2: + size0 = size1 = img_size + elif len(img_size) == 4: + size0, size1 = img_size[:2], img_size[2:] + else: + raise ValueError("Invalid img_size. Please provide 1, 2, or 4 integers.") + else: + size0 = size1 = img_size + + image0, scales0 = load_image(img0_path, resize=size0) + image1, scales1 = load_image(img1_path, resize=size1) + + extractor_type = extractor_type.lower() + if extractor_type == "superpoint": + image0 = rgb_to_grayscale(image0) + image1 = rgb_to_grayscale(image1) + elif extractor_type == "disk": + pass + elif extractor_type=="sift": + image0, scales0 = load_image_sift(img0_path, resize=size0) + image1, scales1 = load_image_sift(img1_path, resize=size1) + else: + raise NotImplementedError( + f"Unsupported feature extractor type: {extractor_type}." + ) + + # Load ONNX models + providers = ["CPUExecutionProvider"] #"CUDAExecutionProvider", + if trt: + providers = [ + ( + "TensorrtExecutionProvider", + { + "trt_fp16_enable": True, + "trt_engine_cache_enable": True, + "trt_engine_cache_path": "weights/cache", + }, + ) + ] + providers + + runner = LightGlueRunner_SIFT( + extractor_path=extractor_path, + lightglue_path=lightglue_path, + providers=providers, + ) + + # Run inference + start_time = time.time() + + m_kpts0, m_kpts1 = runner.run(image0, image1, scales0, scales1) + + end_time = time.time() + elapsed_time = end_time - start_time + print(f"程序运行时间: {elapsed_time} 秒") + + # Visualisation + if viz: + orig_image0, _ = load_image(img0_path) + orig_image1, _ = load_image(img1_path) + viz2d.plot_images( + [orig_image0[0].transpose(1, 2, 0), orig_image1[0].transpose(1, 2, 0)] + ) + + viz2d.plot_matches(m_kpts0, m_kpts1, color="lime", lw=0.2) + viz2d.plt.show() + + return m_kpts0, m_kpts1 + + +if __name__ == "__main__": + args = parse_args() + m_kpts0, m_kpts1 = infer(**vars(args)) + #print(m_kpts0, m_kpts1) From 024bdf56434c7c9860ad66e4a6c81e76970c49c1 Mon Sep 17 00:00:00 2001 From: AMDNO2022 <3048055004@qq.com> Date: Mon, 24 Jun 2024 10:33:53 +0800 Subject: [PATCH 5/5] SIFT support --- infer_sift.py | 2 +- lightglue_onnx/lightglue.py | 10 ++++- onnx_runner/utils.py | 77 +++++++++++++++++++++++++++++++++++-- optimize.py | 4 +- 4 files changed, 86 insertions(+), 7 deletions(-) diff --git a/infer_sift.py b/infer_sift.py index ef7e780..560e26b 100644 --- a/infer_sift.py +++ b/infer_sift.py @@ -9,7 +9,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--img_paths", nargs=2, - default=["assets/remapR1.bmp", "assets/remapL1.bmp"], + default=["assets/DSC_0410.JPG", "assets/DSC_0411.JPG"], required=False, type=str, ) diff --git a/lightglue_onnx/lightglue.py b/lightglue_onnx/lightglue.py index eec3e8a..4c60135 100644 --- a/lightglue_onnx/lightglue.py +++ b/lightglue_onnx/lightglue.py @@ -247,6 +247,7 @@ class LightGlue(nn.Module): features = { "superpoint": ("superpoint_lightglue", 256), "disk": ("disk_lightglue", 128), + "sift": ("sift_lightglue", 128), } def __init__(self, features="superpoint", **conf) -> None: @@ -263,7 +264,11 @@ def __init__(self, features="superpoint", **conf) -> None: self.input_proj = nn.Identity() head_dim = conf.descriptor_dim // conf.num_heads - self.posenc = LearnableFourierPositionalEncoding(2, head_dim) + if (features=="sift"): + bias=True + else: + bias=False + self.posenc = LearnableFourierPositionalEncoding(2+2*bias, head_dim) h, n, d = conf.num_heads, conf.n_layers, conf.descriptor_dim @@ -314,6 +319,9 @@ def forward( desc0 = self.input_proj(desc0) desc1 = self.input_proj(desc1) + + + # cache positional embeddings encoding0 = self.posenc(kpts0) encoding1 = self.posenc(kpts1) diff --git a/onnx_runner/utils.py b/onnx_runner/utils.py index 235efd0..8c78acf 100644 --- a/onnx_runner/utils.py +++ b/onnx_runner/utils.py @@ -1,8 +1,10 @@ -from typing import List, Optional, Union - +#from typing import List, Optional, Union +import torch import cv2 import numpy as np - +from types import SimpleNamespace +from typing import Callable, List, Optional, Tuple, Union +import kornia def read_image(path: str, grayscale: bool = False) -> np.ndarray: """Read an image from path as RGB or grayscale""" @@ -67,9 +69,78 @@ def load_image( img, scales = resize_image(img, resize, fn=fn, interp=interp) return normalize_image(img)[None].astype(np.float32), np.asarray(scales) +def numpy_image_to_torch(image: np.ndarray) -> torch.Tensor: + """Normalize the image tensor and reorder the dimensions.""" + if image.ndim == 3: + image = image.transpose((2, 0, 1)) # HxWxC to CxHxW + elif image.ndim == 2: + image = image[None] # add channel axis + else: + raise ValueError(f"Not an image: {image.shape}") + return torch.tensor(image / 255.0, dtype=torch.float) +def load_image_sift( + path: str, + grayscale: bool = False, + resize: int = None, + fn: str = "max", + interp: str = "area", +) -> torch.Tensor: + img = read_image(path, grayscale=grayscale) + scales = [1, 1] + if resize is not None: + img, scales = resize_image(img, resize, fn=fn, interp=interp) + return numpy_image_to_torch(img), torch.Tensor(scales) + def rgb_to_grayscale(image: np.ndarray) -> np.ndarray: """Convert an RGB image to grayscale.""" scale = np.array([0.299, 0.587, 0.114], dtype=image.dtype).reshape(3, 1, 1) image = (image * scale).sum(axis=-3, keepdims=True) return image + + +class ImagePreprocessor: + default_conf = { + "resize": None, # target edge length, None for no resizing + "side": "long", + "interpolation": "bilinear", + "align_corners": None, + "antialias": True, + } + + def __init__(self, **conf) -> None: + super().__init__() + self.conf = {**self.default_conf, **conf} + self.conf = SimpleNamespace(**self.conf) + + def __call__(self, img: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Resize and preprocess an image, return image and resize scale""" + h, w = img.shape[-2:] + if self.conf.resize is not None: + img = kornia.geometry.transform.resize( + img, + self.conf.resize, + side=self.conf.side, + antialias=self.conf.antialias, + align_corners=self.conf.align_corners, + ) + scale = torch.Tensor([img.shape[-1] / w, img.shape[-2] / h]).to(img) + return img, scale + +class Extractor(torch.nn.Module): + def __init__(self, **conf): + super().__init__() + self.conf = SimpleNamespace(**{**self.default_conf, **conf}) + + @torch.no_grad() + def extract(self, img: torch.Tensor, **conf) -> dict: + """Perform extraction with online resizing""" + if img.dim() == 3: + img = img[None] # add batch dim + assert img.dim() == 4 and img.shape[0] == 1 + shape = img.shape[-2:][::-1] + img, scales = ImagePreprocessor(**{**self.preprocess_conf, **conf})(img) + feats = self.forward({"image": img}) + feats["image_size"] = torch.tensor(shape)[None].to(img).float() + feats["keypoints"] = (feats["keypoints"] + 0.5) / scales[None] - 0.5 + return feats diff --git a/optimize.py b/optimize.py index ccc09eb..1300499 100644 --- a/optimize.py +++ b/optimize.py @@ -15,13 +15,13 @@ def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument( - "-i", "--input", type=str, required=True, help="Path to LightGlue ONNX model." + "-i", "--input", type=str, default="weights/sift_lightglue.onnx", help="Path to LightGlue ONNX model." ) parser.add_argument( "-o", "--output", type=str, help="Path to output fused LightGlue ONNX model." ) parser.add_argument( - "--cpu", action="store_true", help="Whether to optimize for CPU." + "--cpu", default=True , action="store_true", help="Whether to optimize for CPU." ) return parser.parse_args()