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
@@ -29,9 +29,9 @@ class OnnxLightGlue:
     for how to export the model to ONNX and optimize it.
 
     Args:
-        weights: Pretrained weights, or a path to your own exported ONNX model. Available pretrained weights are:
-        `disk`, `superpoint`, `disk_fp16`, and `superpoint_fp16`. Defaults to `disk_fp16`.
-        device: Device to run inference on. Currently, only `cuda` is supported. Defaults to `cuda`.
+        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'``.
     """
     MODEL_URLS: ClassVar[dict[str, str]] = {
         "disk": "https://github.com/fabio-sim/LightGlue-ONNX/releases/download/v1.0.0/disk_lightglue_fused.onnx",
@@ -78,21 +78,27 @@ def forward(self, data: dict[str, dict[str, Tensor]]) -> dict[str, Tensor]:
             data: Dictionary containing both images and the keypoints and descriptors thereof.
 
         Returns:
-            output: Dictionary containing the following matches and scores.
-
-        `data`:
-            image0: dict
-                keypoints (`float32`): [1 x M x 2]
-                descriptors (`float32`): [1 x M x D]
-                image: [1 x C x H x W] or image_size: [1 x 2]
-            image1: dict
-                keypoints (`float32`): [1 x N x 2]
-                descriptors (`float32`): [1 x N x D]
-                image: [1 x C x H x W] or image_size: [1 x 2]
-
-        `output`:
-                matches (`int64`): [S x 2]
-                scores (`float32`): [S]
+            Dictionary containing the matches and scores.
+
+        ``data`` (``dict``):
+            ``image0`` (``dict``):
+                ``keypoints`` (`float32`): :math:`(1, M, 2)`
+
+                ``descriptors`` (`float32`): :math:`(1, M, D)`
+
+                ``image``: :math:`(1, C, H, W)` or ``image_size``: :math:`(1, 2)`
+
+            ``image1`` (``dict``):
+                ``keypoints`` (`float32`): :math:`(1, N, 2)`
+
+                ``descriptors`` (`float32`): :math:`(1, N, D)`
+
+                ``image``: :math:`(1, C, H, W)` or ``image_size``: :math:`(1, 2)`
+
+        ``output`` (``dict``):
+            ``matches`` (`int64`): :math:`(S, 2)`
+
+            ``scores`` (`float32`): :math:`(S)`
         """
         # Input validation.
         for key in self.required_data_keys:


Test patch: