Init Kornia.onnx module

docs/source/index.rst

@@ -15,6 +15,35 @@ within the context of an Open Source community.
 
 Ready to use with state-of-the art Deep Learning models:
 
+DexiNed edge detection model.
+
+.. code-block:: python
+
+      image = kornia.utils.sample.get_sample_images()[0][None]
+      model = DexiNedBuilder.build()
+      model.save(image)
+
+RTDETRDetector for object detection.
+
+.. code-block:: python
+
+      image = kornia.utils.sample.get_sample_images()[0][None]
+      model = RTDETRDetectorBuilder.build()
+      model.save(image)
+
+BoxMotTracker for object tracking.
+
+.. code-block:: python
+
+      import kornia
+      image = kornia.utils.sample.get_sample_images()[0][None]
+      model = BoxMotTracker()
+      for i in range(4):
+         model.update(image)
+      model.save(image)
+
+Vision Transformer for image classification.
+
 .. code:: python
 
    >>> import torch.nn as nn
@@ -66,9 +95,11 @@ Join the community
    io
    image
    losses
+   models
    metrics
    morphology
    nerf
+   onnx
    tracking
    testing
    utils

docs/source/models.rst

@@ -0,0 +1,135 @@
+Models Overview
+===============
+
+This section covers several of Kornia's built-in models for key computer vision tasks. Each model is documented with its respective API and example usage.
+
+.. _RTDETRDetectorBuilder:
+
+RTDETRDetectorBuilder
+---------------------
+
+The `RTDETRDetectorBuilder` class is a builder for constructing a detection model based on the RT-DETR architecture, which is designed for real-time object detection. It is capable of detecting multiple objects within an image and provides efficient inference suitable for real-world applications.
+
+**Key Methods:**
+
+- `build`: Constructs and returns an instance of the RTDETR detection model.
+- `save`: Saves the processed image or results after applying the detection model.
+
+.. autoclass:: kornia.models.detection.rtdetr.RTDETRDetectorBuilder
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+   .. rubric:: Example
+
+   The following code demonstrates how to use `RTDETRDetectorBuilder` to detect objects in an image:
+
+   .. code-block:: python
+
+       import kornia
+       image = kornia.utils.sample.get_sample_images()[0][None]
+       model = kornia.models.detection.rtdetr.RTDETRDetectorBuilder.build()
+       model.save(image)
+
+.. _DexiNedBuilder:
+
+DexiNedBuilder
+--------------
+
+The `DexiNedBuilder` class implements a state-of-the-art edge detection model based on DexiNed, which excels at detecting fine-grained edges in images. This model is well-suited for tasks like medical imaging, object contour detection, and more.
+
+**Key Methods:**
+
+- `build`: Builds and returns an instance of the DexiNed edge detection model.
+- `save`: Saves the detected edges for further processing or visualization.
+
+.. autoclass:: kornia.models.edge_detection.dexined.DexiNedBuilder
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+   .. rubric:: Example
+
+   The following code shows how to use the `DexiNedBuilder` to detect edges in an image:
+
+   .. code-block:: python
+
+       import kornia
+       image = kornia.utils.sample.get_sample_images()[0][None]
+       model = kornia.models.edge_detection.dexined.DexiNedBuilder.build()
+       model.save(image)
+
+.. _SegmentationModels:
+
+SegmentationModels
+------------------
+
+The `SegmentationModels` class offers a flexible API for implementing and running various segmentation models. It supports a variety of architectures such as UNet, FPN, and others, making it highly adaptable for tasks like semantic segmentation, instance segmentation, and more.
+
+**Key Methods:**
+
+- `__init__`: Initializes a segmentation model based on the chosen architecture (e.g., UNet, DeepLabV3, etc.).
+- `forward`: Runs inference on an input tensor and returns segmented output.
+
+**Parameters:**
+
+- `model_name`: (str) Name of the segmentation architecture to use, e.g., `"Unet"`, `"DeepLabV3"`.
+- `classes`: (int) The number of output classes for segmentation.
+
+.. autoclass:: kornia.models.segmentation.segmentation_models.SegmentationModels
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+   .. rubric:: Example
+
+   Here's an example of how to use `SegmentationModels` for binary segmentation:
+
+   .. code-block:: python
+
+       import kornia
+       input_tensor = kornia.utils.sample.get_sample_images()[0][None]
+       model = kornia.models.segmentation.segmentation_models.SegmentationModels()
+       segmented_output = model(input_tensor)
+       print(segmented_output.shape)
+
+.. _BoxMotTracker:
+
+BoxMotTracker
+-------------
+
+The `BoxMotTracker` class is used for multi-object tracking in video streams. It is designed to track bounding boxes of objects across multiple frames, supporting various tracking algorithms for object detection and tracking continuity.
+
+**Key Methods:**
+
+- `__init__`: Initializes the multi-object tracker.
+- `update`: Updates the tracker with a new image frame.
+- `save`: Saves the tracked object data or visualization for post-processing.
+
+**Parameters:**
+
+- `max_lost`: (int) The maximum number of frames where an object can be lost before it is removed from the tracker.
+
+.. autoclass:: kornia.models.tracking.boxmot_tracker.BoxMotTracker
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+   .. rubric:: Example
+
+   The following example demonstrates how to track objects across multiple frames using `BoxMotTracker`:
+
+   .. code-block:: python
+
+       import kornia
+       image = kornia.utils.sample.get_sample_images()[0][None]
+       model = kornia.models.tracking.boxmot_tracker.BoxMotTracker()
+       for i in range(4):
+           model.update(image)  # Update the tracker with new frames
+       model.save(image)       # Save the tracking result
+
+---
+
+.. note::
+
+   This documentation provides detailed information about each model class, its methods, and usage examples. For further details on individual methods and arguments, refer to the respective code documentation.

docs/source/onnx.rst

@@ -0,0 +1,137 @@
+ONNXSequential: Chain Multiple ONNX Models with Ease
+====================================================
+
+The `ONNXSequential` class is a powerful new feature that allows users to effortlessly combine and chain multiple ONNX models together. This is especially useful when you have several pre-trained models or custom ONNX operators that you want to execute sequentially as part of a larger pipeline.
+
+Whether you're working with models for inference, experimentation, or optimization, `ONNXSequential` makes it easier to manage, combine, and run ONNX models in a streamlined manner. It also supports flexibility in execution environments with ONNXRuntime’s execution providers (CPU, CUDA, etc.).
+
+Key Features
+------------
+
+- **Seamless Model Chaining**: Combine multiple ONNX models into a single computational graph.
+- **Flexible Input/Output Mapping**: Control how the outputs of one model are passed as inputs to the next.
+- **Optimized Execution**: Automatically create optimized `ONNXRuntime` sessions to speed up inference.
+- **Export to ONNX**: Save the combined model into a single ONNX file for easy deployment and sharing.
+- **Execution Providers Support**: Utilize ONNXRuntime's execution providers (e.g., `CUDAExecutionProvider`, `CPUExecutionProvider`) for accelerated inference on different hardware.
+- **PyTorch-like Interface**: Use the `ONNXSequential` class like a PyTorch `nn.Sequential` model, including calling it directly for inference.
+
+Quickstart Guide
+----------------
+
+Here's how you can quickly get started with `ONNXSequential`:
+
+1. **Install ONNX and ONNXRuntime**
+
+   If you haven't already installed `onnx` and `onnxruntime`, you can install them using `pip`:
+
+   .. code-block:: bash
+
+      pip install onnx onnxruntime
+
+2. **Combining ONNX Models**
+
+   You can initialize the `ONNXSequential` with a list of ONNX models or file paths. Models will be automatically chained together and optimized for inference.
+
+   .. code-block:: python
+
+      import numpy as np
+      from kornia.onnx import ONNXSequential
+
+      # Initialize ONNXSequential with two models, loading from our only repo
+      onnx_seq = ONNXSequential(
+         "hf://operators/kornia.color.gray.RgbToGrayscale",
+         "hf://operators/kornia.geometry.transform.affwarp.Resize_512x512"
+      )
+
+      # Prepare some input data
+      input_data = np.random.randn(1, 3, 256, 512).astype(np.float32)
+
+      # Perform inference
+      outputs = onnx_seq(input_data)
+
+      # Print the model outputs
+      print(outputs)
+
+   .. note::
+      By default, we assume each ONNX model contains only one input node named "input" and one output node named "output". For complex models, you may need to pass an `io_maps` argument.
+
+3. **Input/Output Mapping Between Models**
+
+   When combining models, you can specify how the outputs of one model are mapped to the inputs of the next. This allows you to chain models in custom ways.
+
+   .. code-block:: python
+
+      io_map = [("model1_output_0", "model2_input_0"), ("model1_output_1", "model2_input_1")]
+      onnx_seq = ONNXSequential("model1.onnx", "model2.onnx", io_map=io_map)
+
+4. **Exporting the Combined Model**
+
+   You can easily export the combined model to an ONNX file:
+
+   .. code-block:: python
+
+      # Export the combined model to a file
+      onnx_seq.export("combined_model.onnx")
+
+5. **Optimizing with Execution Providers**
+
+   Leverage ONNXRuntime's execution providers for optimized inference. For example, to run the model on a GPU:
+
+   .. code-block:: python
+
+      # Initialize with CUDA execution provider
+      onnx_seq = ONNXSequential(
+         "hf://operators/kornia.geometry.transform.flips.Hflip",
+         # Or you may use a local model with either a filepath "YOUR_OWN_MODEL.onnx" or a loaded ONNX model.
+         "hf://models/kornia.models.detection.rtdetr_r18vd_640x640",
+         providers=['CUDAExecutionProvider']
+      )
+
+      # Run inference
+      outputs = onnx_seq(input_data)
+
+
+Frequently Asked Questions (FAQ)
+--------------------------------
+
+**1. Can I chain models from different sources?**
+
+Yes! You can chain models from different ONNX files or directly from `onnx.ModelProto` objects. `ONNXSequential` handles the integration and merging of their graphs.
+
+**2. What happens if the input/output sizes of models don't match?**
+
+You can use the `io_map` parameter to control how outputs of one model are mapped to the inputs of the next. This allows for greater flexibility when chaining models with different architectures.
+
+**3. Can I use custom ONNXRuntime session options?**
+
+Absolutely! You can pass your own session options to the `create_session` method to fine-tune performance, memory usage, or logging.
+
+Why Choose ONNXSequential?
+--------------------------
+
+With the increasing adoption of ONNX for model interoperability and deployment, `ONNXSequential` provides a simple yet powerful interface for combining models and operators. By leveraging ONNXRuntime’s optimization and execution provider capabilities, it gives you the flexibility to:
+- Deploy on different hardware (CPU, GPU).
+- Run complex pipelines in production environments.
+- Combine and experiment with models effortlessly.
+
+Whether you're building an advanced deep learning pipeline or simply trying to chain pre-trained models, `ONNXSequential` makes it easy to manage, optimize, and execute ONNX models at scale.
+
+Get started today and streamline your ONNX workflows!
+
+
+API Documentation
+-----------------
+.. autoclass:: kornia.onnx.sequential.ONNXSequential
+    :members:
+
+.. autoclass:: kornia.onnx.utils.ONNXLoader
+
+    .. code-block:: python
+
+        onnx_loader = ONNXLoader()
+        # Load a HuggingFace operator
+        onnx_loader.load_model("hf://operators/kornia.color.gray.GrayscaleToRgb")  # doctest: +SKIP
+        # Load a local converted/downloaded operator
+        onnx_loader.load_model("operators/kornia.color.gray.GrayscaleToRgb")  # doctest: +SKIP
+
+    :members:

kornia/__init__.py

@@ -10,12 +10,15 @@
     color,
     contrib,
     core,
+    config,
     enhance,
     feature,
     io,
     losses,
     metrics,
+    models,
     morphology,
+    onnx,
     tracking,
     utils,
     x,

kornia/augmentation/base.py

@@ -53,6 +53,10 @@ class _BasicAugmentationBase(Module):
           the batch form ``False``.
     """
 
+    # TODO: Hard to support. Many codes are not ONNX-friendly that contains lots of if-else blocks, etc.
+    # Please contribute if anyone interested.
+    ONNX_EXPORTABLE = False
+
     def __init__(
         self,
         p: float = 0.5,

kornia/color/gray.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Optional
+from typing import ClassVar, Optional
 
 import torch
 
@@ -125,6 +125,9 @@ class GrayscaleToRgb(Module):
         >>> output = rgb(input)  # 2x3x4x5
     """
 
+    ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 1, -1, -1]
+    ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1]
+
     def forward(self, image: Tensor) -> Tensor:
         return grayscale_to_rgb(image)
 
@@ -147,6 +150,9 @@ class RgbToGrayscale(Module):
         >>> output = gray(input)  # 2x1x4x5
     """
 
+    ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1]
+    ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 1, -1, -1]
+
     def __init__(self, rgb_weights: Optional[Tensor] = None) -> None:
         super().__init__()
         if rgb_weights is None:
@@ -175,5 +181,8 @@ class BgrToGrayscale(Module):
         >>> output = gray(input)  # 2x1x4x5
     """
 
+    ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1]
+    ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 1, -1, -1]
+
     def forward(self, image: Tensor) -> Tensor:
         return bgr_to_grayscale(image)