Repository: kornia/kornia

Issue:
feat: LightGlue-ONNX
#### Changes
<!-- Please include a summary of the change and which issue is fixed. -->
This PR adds a wrapper class for loading and running LightGlue-ONNX models via ONNXRuntime.
<!-- Please also include relevant motivation and context. -->
Re: https://github.com/fabio-sim/LightGlue-ONNX/issues/40
<!-- List any dependencies that are required for this change. -->
`onnxruntime-gpu>=1.16` is required as a dependency to instantiate the new `OnnxLightGlue` class. (Importing it will still work without installing).
related to #2559 

#### Type of change
<!-- Please delete options that are not relevant. -->
- [x] 📚  Documentation Update
- [x] 🧪 Tests Cases
- [x] 🔬 New feature (non-breaking change which adds functionality)
- [x] 📝 This change requires a documentation update

#### Checklist

- [x] My code follows the style guidelines of this project
- [x] I have performed a self-review of my own code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [ ] Did you update CHANGELOG in case of a major change?

#### Example Usage
Sample images can be found [here](https://github.com/fabio-sim/LightGlue-ONNX/tree/main/assets).
```python
import torch

from kornia.feature import DISK, OnnxLightGlue
from kornia.io import ImageLoadType, load_image

device = torch.device("cuda")

img0 = load_image("sacre_coeur1.jpg", ImageLoadType.RGB32, device=device)[None]
img1 = load_image("sacre_coeur2.jpg", ImageLoadType.RGB32, device=device)[None]

extractor = DISK.from_pretrained("depth", device=device).eval().to(device)

data = {}
with torch.no_grad():
    for key, img in [("image0", img0), ("image1", img1)]:
        features = extractor(img, n=None, window_size=5, score_threshold=0.0, pad_if_not_divisible=True)
        data[key] = {
            "image": img,
            "keypoints": features[0].keypoints[None],
            "keypoint_scores": features[0].detection_scores[None],
            "descriptors": features[0].descriptors[None],
        }

matcher = OnnxLightGlue(weights="disk_fp16", device=device)
result = matcher(data)
print(result)
```



Discussion:
@ducha-aiki I believe some tests are failing as `onnxruntime-gpu` is not available on MacOS. Should it still be included as a dependency in `requirements-dev.txt`?
> @ducha-aiki I believe some tests are failing as `onnxruntime-gpu` is not available on MacOS. Should it still be included as a dependency in `requirements-dev.txt`?

Yes, just skip the test on the macOS then
> i'm getting an error on `TestOnnxLightGlue.test_forward`, someone can check locally too? @edgarriba @ducha-aiki -- the CI will not run this test, since it's just for cuda devices
> 
> ```
> E       RuntimeError: Error when binding input: There's no data transfer registered for copying tensors from Device:[DeviceType:1 MemoryType:0 DeviceId:0] to Device:[DeviceType:0 MemoryType:0 DeviceId:0]
> ```

Hi, @johnnv1, I am able to reproduce this issue if I have only `onnxruntime` installed. Can you check if you have installed only `onnxruntime-gpu`?
> Hi, @johnnv1, I am able to reproduce this issue if I have only `onnxruntime` installed. Can you check if you have installed only `onnxruntime-gpu`?

i see, I think I had installed both libraries, within just  `onnxruntime-gpu` it's fine 😄 

Resolving patch:
diff --git a/kornia/feature/lightglue_onnx/lightglue.py b/kornia/feature/lightglue_onnx/lightglue.py
--- a/kornia/feature/lightglue_onnx/lightglue.py
+++ b/kornia/feature/lightglue_onnx/lightglue.py
@@ -30,8 +30,9 @@ class OnnxLightGlue:
 
     Args:
         weights: Pretrained weights, or a path to your own exported ONNX model. Available pretrained weights
-          are ``'disk'``, ``'superpoint'``, ``'disk_fp16'``, and ``'superpoint_fp16'``.
-        device: Device to run inference on. Currently, only ``'cuda'`` is supported. Defaults to ``'cuda'``.
+          are ``'disk'``, ``'superpoint'``, ``'disk_fp16'``, and ``'superpoint_fp16'``. `Note that FP16 requires CUDA.`
+          Defaults to ``'disk_fp16'`` if ``device`` is CUDA, and ``'disk'`` if CPU.
+        device: Device to run inference on.
     """
 
     MODEL_URLS: ClassVar[dict[str, str]] = {
@@ -43,25 +44,32 @@ class OnnxLightGlue:
 
     required_data_keys: ClassVar[list[str]] = ["image0", "image1"]
 
-    def __init__(self, weights: str = "disk_fp16", device: Device = None) -> None:
+    def __init__(self, weights: str | None = None, device: Device = "cpu") -> None:
         KORNIA_CHECK(ort is not None, "onnxruntime is not installed.")
         KORNIA_CHECK(np is not None, "numpy is not installed.")
 
-        if device is None:
-            device = torch.device("cuda")
-        elif isinstance(device, str):
-            device = torch.device(device)
+        device = torch.device(device)  # type: ignore
         self.device = device
 
         if device.type == "cpu":
-            raise NotImplementedError("CPUExecutionProvider is not supported yet for Multihead-Attention op.")
+            providers = ["CPUExecutionProvider"]
         elif device.type == "cuda":
             providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
         else:
             raise ValueError(f"Unsupported device {device}")
 
+        if weights is None:
+            weights = "disk_fp16" if device.type == "cuda" else "disk"
+
         if weights in self.MODEL_URLS:
-            weights = download_onnx_from_url(self.MODEL_URLS[weights])
+            if "fp16" in weights:
+                KORNIA_CHECK(device.type == "cuda", "FP16 requires CUDA.")
+
+            url = self.MODEL_URLS[weights]
+            if device.type == "cpu":
+                url = url.replace(".onnx", "_cpu.onnx")
+
+            weights = download_onnx_from_url(url)
 
         self.session = ort.InferenceSession(weights, providers=providers)
 
diff --git a/kornia/geometry/depth.py b/kornia/geometry/depth.py
--- a/kornia/geometry/depth.py
+++ b/kornia/geometry/depth.py
@@ -112,10 +112,10 @@ def depth_to_3d_v2(
 
 @deprecated(
     replace_with="depth_to_3d_v2",
-    version='0.7.0',
+    version="0.7.0",
     extra_reason=(
-        ' This function will be replaced with the `depth_to_3d_v2` behaviour, where the that does not require the'
-        ' creation of a meshgrid. The return shape can be not backward compatible between these implementations.'
+        " This function will be replaced with the `depth_to_3d_v2` behaviour, where the that does not require the"
+        " creation of a meshgrid. The return shape can be not backward compatible between these implementations."
     ),
 )
 def depth_to_3d(depth: Tensor, camera_matrix: Tensor, normalize_points: bool = False) -> Tensor:


Test patch:
diff --git a/test/feature/test_lightglue_onnx.py b/test/feature/test_lightglue_onnx.py
--- a/test/feature/test_lightglue_onnx.py
+++ b/test/feature/test_lightglue_onnx.py
@@ -11,22 +11,28 @@
 
 
 @pytest.mark.skipif(ort is None, reason="OnnxLightGlue requires onnxruntime-gpu")
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="OnnxLightGlue requires CUDA")
 class TestOnnxLightGlue:
     @pytest.mark.slow
-    @pytest.mark.parametrize("weights", OnnxLightGlue.MODEL_URLS.keys())
-    def test_pretrained_weights(self, weights):
-        model = OnnxLightGlue(weights)
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Test requires CUDA")
+    @pytest.mark.parametrize("weights", ["disk", "superpoint", "disk_fp16", "superpoint_fp16"])
+    def test_pretrained_weights_cuda(self, weights):
+        model = OnnxLightGlue(weights, "cuda")
         assert model is not None
 
     @pytest.mark.slow
-    def test_forward(self):
-        model = OnnxLightGlue()
+    @pytest.mark.parametrize("weights", ["disk", "superpoint"])
+    def test_pretrained_weights_cpu(self, weights):
+        model = OnnxLightGlue(weights, "cpu")
+        assert model is not None
 
-        device = torch.device("cuda")
-        kpts = torch.zeros(1, 5, 2, device=device)
-        desc = torch.zeros(1, 5, 128, device=device)
-        image = torch.zeros(1, 3, 10, 10)
+    @pytest.mark.slow
+    @pytest.mark.parametrize("dtype", [torch.float32])
+    def test_forward(self, dtype, device):
+        model = OnnxLightGlue(device=device)
+
+        kpts = torch.zeros(1, 5, 2, dtype=dtype, device=device)
+        desc = torch.zeros(1, 5, 128, dtype=dtype, device=device)
+        image = torch.zeros(1, 3, 10, 10, dtype=dtype)
         outputs = model(
             {
                 "image0": {"keypoints": kpts, "descriptors": desc, "image": image},
@@ -44,14 +50,18 @@ def test_normalize_keypoints(self):
         assert torch.all(torch.abs(kpts) <= 1).item()
 
     @pytest.mark.slow
-    def test_exception(self):
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Test requires CUDA")
+    def test_exception(self, device):
         with pytest.raises(RuntimeError) as e:
             OnnxLightGlue(device="invalid device")
         assert "Invalid device string: 'invalid device'" in str(e)
 
-        model = OnnxLightGlue()
+        with pytest.raises(Exception) as e:
+            OnnxLightGlue("disk_fp16", "cpu")
+        assert "FP16 requires CUDA." in str(e)
+
+        model = OnnxLightGlue(device=device)
 
-        device = torch.device("cuda")
         kpts = torch.zeros(1, 5, 2, device=device)
         desc = torch.zeros(1, 5, 128, device=device)
         image = torch.zeros(1, 3, 10, 10)
@@ -76,10 +86,15 @@ def test_exception(self):
         assert "Wrong dtype" in str(e)
 
         # Wrong device
+        wrong_device = torch.device("cpu" if device.type == "cuda" else "cuda")
         with pytest.raises(Exception) as e:
             model(
                 {
-                    "image0": {"keypoints": torch.zeros(1, 5, 2, device="cpu"), "descriptors": desc, "image": image},
+                    "image0": {
+                        "keypoints": torch.zeros(1, 5, 2, device=wrong_device),
+                        "descriptors": desc,
+                        "image": image,
+                    },
                     "image1": {"keypoints": kpts, "descriptors": desc, "image": image},
                 }
             )