From 7ff2a0a509b9c80b56ee6c2e1b99ebda432cce67 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Sat, 7 Sep 2024 14:42:22 +0300 Subject: [PATCH 01/84] init --- docs/source/index.rst | 1 + docs/source/onnx.rst | 117 +++++++++++++++++++++ kornia/__init__.py | 1 + kornia/core/external.py | 2 + kornia/geometry/transform/affwarp.py | 15 +-- kornia/geometry/transform/flips.py | 3 +- kornia/onnx/__init__.py | 1 + kornia/onnx/sequential.py | 152 +++++++++++++++++++++++++++ 8 files changed, 284 insertions(+), 8 deletions(-) create mode 100644 docs/source/onnx.rst create mode 100644 kornia/onnx/__init__.py create mode 100644 kornia/onnx/sequential.py diff --git a/docs/source/index.rst b/docs/source/index.rst index 2fbc49c79c..4f18cee559 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -69,6 +69,7 @@ Join the community metrics morphology nerf + onnx tracking testing utils diff --git a/docs/source/onnx.rst b/docs/source/onnx.rst new file mode 100644 index 0000000000..c9bc314151 --- /dev/null +++ b/docs/source/onnx.rst @@ -0,0 +1,117 @@ +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 + onnx_seq = ONNXSequential("model1.onnx", "model2.onnx") + + # Prepare some input data + input_data = np.random.randn(1, 3, 224, 224).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` arguement. + +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("model1.onnx", "model2.onnx", 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: diff --git a/kornia/__init__.py b/kornia/__init__.py index 84f22f2051..5896d56f52 100644 --- a/kornia/__init__.py +++ b/kornia/__init__.py @@ -16,6 +16,7 @@ losses, metrics, morphology, + onnx, tracking, utils, x, diff --git a/kornia/core/external.py b/kornia/core/external.py index 1e0160035d..96de033f91 100644 --- a/kornia/core/external.py +++ b/kornia/core/external.py @@ -94,3 +94,5 @@ def __dir__(self) -> List[str]: numpy = LazyLoader("numpy") PILImage = LazyLoader("PIL.Image") diffusers = LazyLoader("diffusers") +onnx = LazyLoader("onnx") +onnxruntime = LazyLoader("onnxruntime") diff --git a/kornia/geometry/transform/affwarp.py b/kornia/geometry/transform/affwarp.py index c85c5ddf3d..41f726a8b0 100644 --- a/kornia/geometry/transform/affwarp.py +++ b/kornia/geometry/transform/affwarp.py @@ -5,6 +5,7 @@ from torch import nn from kornia.core import ones, ones_like, zeros +from kornia.core import ImageModule as Module from kornia.filters import gaussian_blur2d from kornia.utils import _extract_device_dtype from kornia.utils.image import perform_keep_shape_image @@ -648,7 +649,7 @@ def rescale( return resize(input, size, interpolation=interpolation, align_corners=align_corners, antialias=antialias) -class Resize(nn.Module): +class Resize(Module): r"""Resize the input torch.Tensor to the given size. Args: @@ -704,7 +705,7 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: ) -class Affine(nn.Module): +class Affine(Module): r"""Apply multiple elementary affine transforms simultaneously. Args: @@ -800,7 +801,7 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: return affine(input, matrix[..., :2, :3], self.mode, self.padding_mode, self.align_corners) -class Rescale(nn.Module): +class Rescale(Module): r"""Rescale the input torch.Tensor with the given factor. Args: @@ -843,7 +844,7 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: ) -class Rotate(nn.Module): +class Rotate(Module): r"""Rotate the tensor anti-clockwise about the centre. Args: @@ -888,7 +889,7 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: return rotate(input, self.angle, self.center, self.mode, self.padding_mode, self.align_corners) -class Translate(nn.Module): +class Translate(Module): r"""Translate the tensor in pixel units. Args: @@ -925,7 +926,7 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: return translate(input, self.translation, self.mode, self.padding_mode, self.align_corners) -class Scale(nn.Module): +class Scale(Module): r"""Scale the tensor by a factor. Args: @@ -972,7 +973,7 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: return scale(input, self.scale_factor, self.center, self.mode, self.padding_mode, self.align_corners) -class Shear(nn.Module): +class Shear(Module): r"""Shear the tensor. Args: diff --git a/kornia/geometry/transform/flips.py b/kornia/geometry/transform/flips.py index 6519725dd5..87df3bd0d1 100644 --- a/kornia/geometry/transform/flips.py +++ b/kornia/geometry/transform/flips.py @@ -1,6 +1,7 @@ import torch -from kornia.core import Module, Tensor +from kornia.core import Tensor +from kornia.core import ImageModule as Module __all__ = ["Vflip", "Hflip", "Rot180", "rot180", "hflip", "vflip"] diff --git a/kornia/onnx/__init__.py b/kornia/onnx/__init__.py new file mode 100644 index 0000000000..bce2460a4f --- /dev/null +++ b/kornia/onnx/__init__.py @@ -0,0 +1 @@ +from .sequential import * diff --git a/kornia/onnx/sequential.py b/kornia/onnx/sequential.py new file mode 100644 index 0000000000..4d30fea5dd --- /dev/null +++ b/kornia/onnx/sequential.py @@ -0,0 +1,152 @@ +from typing import Optional, Union + +from kornia.core.external import numpy as np +from kornia.core.external import onnx +from kornia.core.external import onnxruntime as ort + +__all__ = ["ONNXSequential"] + + +class ONNXSequential: + """ONNXSequential to chain multiple ONNX operators together. + + Args: + *args: + A variable number of ONNX models (either ONNX ModelProto objects or file paths). + providers: + A list of execution providers for ONNXRuntime (e.g., ['CUDAExecutionProvider', 'CPUExecutionProvider']). + session_options: + Optional ONNXRuntime session options for optimizing the session. + io_maps: + An optional list of list of tuples specifying input-output mappings for combining models. + If None, we assume the default input name and output name are "input" and "output" accordingly, and + only one input and output node for each graph. + If not None, `io_maps[0]` shall represent the `io_map` for combining the first and second ONNX models. + """ + def __init__( + self, + *args: Union[onnx.ModelProto, str], # type:ignore + providers: Optional[list[str]] = None, + session_options: Optional[ort.SessionOptions] = None, # type:ignore + io_maps: Optional[list[tuple[str, str]]] = None + ) -> None: + self.operators = args + self._combined_op = self._combine(io_maps) + self._session = self.create_session() + + def _load_op(self, arg: Union[onnx.ModelProto, str]) -> onnx.ModelProto: # type:ignore + """Loads an ONNX model, either from a file path or use the provided ONNX ModelProto. + + Args: + arg: Either an ONNX ModelProto object or a file path to an ONNX model. + + Returns: + onnx.ModelProto: The loaded ONNX model. + """ + if isinstance(arg, str): + return onnx.load(arg) # type:ignore + return arg + + def _combine(self, io_maps: Optional[list[tuple[str, str]]] = None) -> onnx.ModelProto: # type:ignore + """ Combine the provided ONNX models into a single ONNX graph. Optionally, map inputs and outputs + between operators using the `io_map`. + + Args: + io_maps: + A list of list of tuples representing input-output mappings for combining the models. + Example: [[(model1_output_name, model2_input_name)], [(model2_output_name, model3_input_name)]]. + + Returns: + onnx.ModelProto: The combined ONNX model as a single ONNX graph. + + Raises: + ValueError: If no operators are provided for combination. + """ + if len(self.operators) == 0: + raise ValueError("No operators found.") + + combined_op = self._load_op(self.operators[0]) + combined_op = onnx.compose.add_prefix(combined_op, prefix=f"K{str(0).zfill(2)}-") + + for i, op in enumerate(self.operators[1:]): + next_op = onnx.compose.add_prefix(self._load_op(op), prefix=f"K{str(i + 1).zfill(2)}-") + if io_maps is None: + io_map = [(f"K{str(i).zfill(2)}-output", f"K{str(i + 1).zfill(2)}-input")] + else: + io_map = [(f"K{str(i).zfill(2)}-{it[0]}", f"K{str(i + 1).zfill(2)}-{it[1]}") for it in io_maps[i]] + combined_op = onnx.compose.merge_models(combined_op, next_op, io_map=io_map) + + return combined_op + + def export(self, file_path: str) -> None: + """Export the combined ONNX model to a file. + + Args: + file_path: str + The file path to export the combined ONNX model. + """ + onnx.save(self._combined_op, file_path) + + def create_session( + self, + providers: Optional[list[str]] = None, + session_options: Optional[ort.SessionOptions] = None + ) -> ort.InferenceSession: # type:ignore + """Create an optimized ONNXRuntime InferenceSession for the combined model. + + Args: + providers: + Execution providers for ONNXRuntime (e.g., ['CUDAExecutionProvider', 'CPUExecutionProvider']). + session_options: + Optional ONNXRuntime session options for session configuration and optimizations. + + Returns: + ort.InferenceSession: The ONNXRuntime session optimized for inference. + """ + if providers is None: + sess_options = ort.SessionOptions() # type:ignore + sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED # type:ignore + if session_options is None: + sess_options = ort.SessionOptions() + sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED + session = ort.InferenceSession( + self._combined_op.SerializeToString(), + sess_options=sess_options, + providers=providers or ['CPUExecutionProvider'] + ) + return session + + def set_session(self, session: ort.InferenceSession) -> None: # type: ignore + """Set a custom ONNXRuntime InferenceSession. + + Args: + session: ort.InferenceSession + The custom ONNXRuntime session to be set for inference. + """ + self._session = session + + def get_session(self) -> ort.InferenceSession: # type: ignore + """Get the current ONNXRuntime InferenceSession. + + Returns: + ort.InferenceSession: The current ONNXRuntime session. + """ + return self._session + + def __call__(self, *inputs: np.ndarray) -> list[np.ndarray]: # type:ignore + """Perform inference using the combined ONNX model. + + Args: + *inputs: Inputs to the ONNX model. The number of inputs must match the expected inputs of the session. + + Returns: + List: The outputs from the ONNX model inference. + """ + ort_inputs = self._session.get_inputs() + if len(ort_inputs) != len(inputs): + raise ValueError(f"Expected {len(ort_inputs)} for the session while only {len(inputs)} received.") + + ort_input_values = {ort_inputs[i].name: inputs[i] for i in range(len(ort_inputs))} + outputs = self._session.run(None, ort_input_values) + + return outputs From a162c61b8d0e056554d0738f4f5ccccd2a155e9a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 7 Sep 2024 11:43:41 +0000 Subject: [PATCH 02/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/geometry/transform/affwarp.py | 3 +-- kornia/geometry/transform/flips.py | 2 +- kornia/onnx/sequential.py | 19 +++++++++---------- 3 files changed, 11 insertions(+), 13 deletions(-) diff --git a/kornia/geometry/transform/affwarp.py b/kornia/geometry/transform/affwarp.py index 41f726a8b0..3abb418ec4 100644 --- a/kornia/geometry/transform/affwarp.py +++ b/kornia/geometry/transform/affwarp.py @@ -2,10 +2,9 @@ from typing import Optional, Tuple, Union import torch -from torch import nn -from kornia.core import ones, ones_like, zeros from kornia.core import ImageModule as Module +from kornia.core import ones, ones_like, zeros from kornia.filters import gaussian_blur2d from kornia.utils import _extract_device_dtype from kornia.utils.image import perform_keep_shape_image diff --git a/kornia/geometry/transform/flips.py b/kornia/geometry/transform/flips.py index 87df3bd0d1..cd0030347d 100644 --- a/kornia/geometry/transform/flips.py +++ b/kornia/geometry/transform/flips.py @@ -1,7 +1,7 @@ import torch -from kornia.core import Tensor from kornia.core import ImageModule as Module +from kornia.core import Tensor __all__ = ["Vflip", "Hflip", "Rot180", "rot180", "hflip", "vflip"] diff --git a/kornia/onnx/sequential.py b/kornia/onnx/sequential.py index 4d30fea5dd..535639086c 100644 --- a/kornia/onnx/sequential.py +++ b/kornia/onnx/sequential.py @@ -23,12 +23,13 @@ class ONNXSequential: only one input and output node for each graph. If not None, `io_maps[0]` shall represent the `io_map` for combining the first and second ONNX models. """ + def __init__( self, *args: Union[onnx.ModelProto, str], # type:ignore providers: Optional[list[str]] = None, session_options: Optional[ort.SessionOptions] = None, # type:ignore - io_maps: Optional[list[tuple[str, str]]] = None + io_maps: Optional[list[tuple[str, str]]] = None, ) -> None: self.operators = args self._combined_op = self._combine(io_maps) @@ -48,8 +49,8 @@ def _load_op(self, arg: Union[onnx.ModelProto, str]) -> onnx.ModelProto: # type return arg def _combine(self, io_maps: Optional[list[tuple[str, str]]] = None) -> onnx.ModelProto: # type:ignore - """ Combine the provided ONNX models into a single ONNX graph. Optionally, map inputs and outputs - between operators using the `io_map`. + """Combine the provided ONNX models into a single ONNX graph. Optionally, map inputs and outputs between + operators using the `io_map`. Args: io_maps: @@ -58,7 +59,7 @@ def _combine(self, io_maps: Optional[list[tuple[str, str]]] = None) -> onnx.Mode Returns: onnx.ModelProto: The combined ONNX model as a single ONNX graph. - + Raises: ValueError: If no operators are provided for combination. """ @@ -88,12 +89,10 @@ def export(self, file_path: str) -> None: onnx.save(self._combined_op, file_path) def create_session( - self, - providers: Optional[list[str]] = None, - session_options: Optional[ort.SessionOptions] = None + self, providers: Optional[list[str]] = None, session_options: Optional[ort.SessionOptions] = None ) -> ort.InferenceSession: # type:ignore """Create an optimized ONNXRuntime InferenceSession for the combined model. - + Args: providers: Execution providers for ONNXRuntime (e.g., ['CUDAExecutionProvider', 'CPUExecutionProvider']). @@ -112,7 +111,7 @@ def create_session( session = ort.InferenceSession( self._combined_op.SerializeToString(), sess_options=sess_options, - providers=providers or ['CPUExecutionProvider'] + providers=providers or ["CPUExecutionProvider"], ) return session @@ -148,5 +147,5 @@ def __call__(self, *inputs: np.ndarray) -> list[np.ndarray]: # type:ignore ort_input_values = {ort_inputs[i].name: inputs[i] for i in range(len(ort_inputs))} outputs = self._session.run(None, ort_input_values) - + return outputs From 772716e422e8e8815659b07c3771060bd6cf3fe8 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Sat, 7 Sep 2024 14:50:46 +0300 Subject: [PATCH 03/84] update --- kornia/onnx/sequential.py | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/kornia/onnx/sequential.py b/kornia/onnx/sequential.py index 535639086c..a2a9480f2f 100644 --- a/kornia/onnx/sequential.py +++ b/kornia/onnx/sequential.py @@ -67,15 +67,15 @@ def _combine(self, io_maps: Optional[list[tuple[str, str]]] = None) -> onnx.Mode raise ValueError("No operators found.") combined_op = self._load_op(self.operators[0]) - combined_op = onnx.compose.add_prefix(combined_op, prefix=f"K{str(0).zfill(2)}-") + combined_op = onnx.compose.add_prefix(combined_op, prefix=f"K{str(0).zfill(2)}-") # type:ignore for i, op in enumerate(self.operators[1:]): - next_op = onnx.compose.add_prefix(self._load_op(op), prefix=f"K{str(i + 1).zfill(2)}-") + next_op = onnx.compose.add_prefix(self._load_op(op), prefix=f"K{str(i + 1).zfill(2)}-") # type:ignore if io_maps is None: io_map = [(f"K{str(i).zfill(2)}-output", f"K{str(i + 1).zfill(2)}-input")] else: io_map = [(f"K{str(i).zfill(2)}-{it[0]}", f"K{str(i + 1).zfill(2)}-{it[1]}") for it in io_maps[i]] - combined_op = onnx.compose.merge_models(combined_op, next_op, io_map=io_map) + combined_op = onnx.compose.merge_models(combined_op, next_op, io_map=io_map) # type:ignore return combined_op @@ -86,10 +86,12 @@ def export(self, file_path: str) -> None: file_path: str The file path to export the combined ONNX model. """ - onnx.save(self._combined_op, file_path) + onnx.save(self._combined_op, file_path) # type:ignore def create_session( - self, providers: Optional[list[str]] = None, session_options: Optional[ort.SessionOptions] = None + self, + providers: Optional[list[str]] = None, + session_options: Optional[ort.SessionOptions] = None # type:ignore ) -> ort.InferenceSession: # type:ignore """Create an optimized ONNXRuntime InferenceSession for the combined model. @@ -102,13 +104,10 @@ def create_session( Returns: ort.InferenceSession: The ONNXRuntime session optimized for inference. """ - if providers is None: + if session_options is None: sess_options = ort.SessionOptions() # type:ignore sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED # type:ignore - if session_options is None: - sess_options = ort.SessionOptions() - sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED - session = ort.InferenceSession( + session = ort.InferenceSession( # type:ignore self._combined_op.SerializeToString(), sess_options=sess_options, providers=providers or ["CPUExecutionProvider"], From 0fe5c04b52a95865b09d465b00d595729a4f67ac Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 7 Sep 2024 11:51:30 +0000 Subject: [PATCH 04/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/onnx/sequential.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kornia/onnx/sequential.py b/kornia/onnx/sequential.py index a2a9480f2f..223294912b 100644 --- a/kornia/onnx/sequential.py +++ b/kornia/onnx/sequential.py @@ -91,7 +91,7 @@ def export(self, file_path: str) -> None: def create_session( self, providers: Optional[list[str]] = None, - session_options: Optional[ort.SessionOptions] = None # type:ignore + session_options: Optional[ort.SessionOptions] = None, # type:ignore ) -> ort.InferenceSession: # type:ignore """Create an optimized ONNXRuntime InferenceSession for the combined model. From 50a999f037871a81ca75a8c1558c9a39ee79d2bb Mon Sep 17 00:00:00 2001 From: shijianjian Date: Sat, 7 Sep 2024 18:07:06 +0300 Subject: [PATCH 05/84] update --- docs/source/onnx.rst | 17 +++++-- kornia/augmentation/base.py | 4 ++ kornia/color/gray.py | 9 ++++ kornia/color/hls.py | 6 +++ kornia/color/hsv.py | 6 +++ kornia/color/lab.py | 6 +++ kornia/color/luv.py | 6 +++ kornia/color/raw.py | 6 +++ kornia/color/rgb.py | 24 +++++++++ kornia/color/xyz.py | 6 +++ kornia/color/ycbcr.py | 6 +++ kornia/color/yuv.py | 18 +++++++ kornia/core/module.py | 50 ++++++++++++++++++- kornia/enhance/adjust.py | 9 ++++ kornia/filters/canny.py | 3 ++ kornia/filters/dexined.py | 3 ++ kornia/filters/motion.py | 3 ++ kornia/filters/sobel.py | 6 +++ kornia/onnx/sequential.py | 17 ++++++- kornia/onnx/utils.py | 97 +++++++++++++++++++++++++++++++++++++ 20 files changed, 295 insertions(+), 7 deletions(-) create mode 100644 kornia/onnx/utils.py diff --git a/docs/source/onnx.rst b/docs/source/onnx.rst index c9bc314151..093bd694d1 100644 --- a/docs/source/onnx.rst +++ b/docs/source/onnx.rst @@ -37,11 +37,14 @@ Here's how you can quickly get started with `ONNXSequential`: import numpy as np from kornia.onnx import ONNXSequential - # Initialize ONNXSequential with two models - onnx_seq = ONNXSequential("model1.onnx", "model2.onnx") + # 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, 224, 224).astype(np.float32) + input_data = np.random.randn(1, 3, 256, 512).astype(np.float32) # Perform inference outputs = onnx_seq(input_data) @@ -50,7 +53,7 @@ Here's how you can quickly get started with `ONNXSequential`: 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` arguement. + 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** @@ -77,7 +80,11 @@ Here's how you can quickly get started with `ONNXSequential`: .. code-block:: python # Initialize with CUDA execution provider - onnx_seq = ONNXSequential("model1.onnx", "model2.onnx", providers=['CUDAExecutionProvider']) + onnx_seq = ONNXSequential( + "hf://operators/kornia.color.gray.RgbToGrayscale", + "hf://operators/kornia.geometry.transform.affwarp.Resize_512x512", + providers=['CUDAExecutionProvider'] + ) # Run inference outputs = onnx_seq(input_data) diff --git a/kornia/augmentation/base.py b/kornia/augmentation/base.py index e912768373..a4cfa33e2b 100644 --- a/kornia/augmentation/base.py +++ b/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, diff --git a/kornia/color/gray.py b/kornia/color/gray.py index 96223a7a69..2fe6fd7404 100644 --- a/kornia/color/gray.py +++ b/kornia/color/gray.py @@ -125,6 +125,9 @@ class GrayscaleToRgb(Module): >>> output = rgb(input) # 2x3x4x5 """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 1, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, 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: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, 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: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 1, -1, -1] + def forward(self, image: Tensor) -> Tensor: return bgr_to_grayscale(image) diff --git a/kornia/color/hls.py b/kornia/color/hls.py index 39c457a7c2..0f498dc66d 100644 --- a/kornia/color/hls.py +++ b/kornia/color/hls.py @@ -142,6 +142,9 @@ class RgbToHls(Module): >>> output = hls(input) # 2x3x4x5 """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def forward(self, image: Tensor) -> Tensor: return rgb_to_hls(image) @@ -167,5 +170,8 @@ class HlsToRgb(Module): >>> output = rgb(input) # 2x3x4x5 """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def forward(self, image: Tensor) -> Tensor: return hls_to_rgb(image) diff --git a/kornia/color/hsv.py b/kornia/color/hsv.py index ce5cc4c55c..561e585ec9 100644 --- a/kornia/color/hsv.py +++ b/kornia/color/hsv.py @@ -116,6 +116,9 @@ class RgbToHsv(Module): >>> output = hsv(input) # 2x3x4x5 """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def __init__(self, eps: float = 1e-6) -> None: super().__init__() self.eps = eps @@ -142,5 +145,8 @@ class HsvToRgb(Module): >>> output = rgb(input) # 2x3x4x5 """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def forward(self, image: torch.Tensor) -> torch.Tensor: return hsv_to_rgb(image) diff --git a/kornia/color/lab.py b/kornia/color/lab.py index c9bda065f9..0231a85b2d 100644 --- a/kornia/color/lab.py +++ b/kornia/color/lab.py @@ -150,6 +150,9 @@ class RgbToLab(Module): [3] https://github.com/torch/image/blob/dc061b98fb7e946e00034a5fc73e883a299edc7f/generic/image.c#L1467 """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def forward(self, image: torch.Tensor) -> torch.Tensor: return rgb_to_lab(image) @@ -177,5 +180,8 @@ class LabToRgb(Module): [3] https://github.com/torch/image/blob/dc061b98fb7e946e00034a5fc73e883a299edc7f/generic/image.c#L1518 """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def forward(self, image: torch.Tensor, clip: bool = True) -> torch.Tensor: return lab_to_rgb(image, clip) diff --git a/kornia/color/luv.py b/kornia/color/luv.py index 4af051e4f8..f5055d1ebe 100644 --- a/kornia/color/luv.py +++ b/kornia/color/luv.py @@ -141,6 +141,9 @@ class RgbToLuv(Module): [3] http://www.poynton.com/ColorFAQ.html """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def forward(self, image: torch.Tensor) -> torch.Tensor: return rgb_to_luv(image) @@ -168,5 +171,8 @@ class LuvToRgb(Module): [3] http://www.poynton.com/ColorFAQ.html """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def forward(self, image: torch.Tensor) -> torch.Tensor: return luv_to_rgb(image) diff --git a/kornia/color/raw.py b/kornia/color/raw.py index 7dca550375..a0d68d0d50 100644 --- a/kornia/color/raw.py +++ b/kornia/color/raw.py @@ -288,6 +288,9 @@ class RawToRgb(Module): >>> output = rgb(rawinput) # 2x3x4x5 """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 1, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def __init__(self, cfa: CFA) -> None: super().__init__() self.cfa = cfa @@ -314,6 +317,9 @@ class RgbToRaw(Module): >>> output = raw(rgbinput) # 2x1x4x6 """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 1, -1, -1] + def __init__(self, cfa: CFA) -> None: super().__init__() self.cfa = cfa diff --git a/kornia/color/rgb.py b/kornia/color/rgb.py index 1ac9d6761e..5003ccbafb 100644 --- a/kornia/color/rgb.py +++ b/kornia/color/rgb.py @@ -248,6 +248,9 @@ class BgrToRgb(Module): >>> output = rgb(input) # 2x3x4x5 """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def forward(self, image: Tensor) -> Tensor: return bgr_to_rgb(image) @@ -270,6 +273,9 @@ class RgbToBgr(Module): >>> output = bgr(input) # 2x3x4x5 """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def forward(self, image: Tensor) -> Tensor: return rgb_to_bgr(image) @@ -298,6 +304,9 @@ class RgbToRgba(Module): >>> output = rgba(input) # 2x4x4x5 """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 4, -1, -1] + def __init__(self, alpha_val: Union[float, Tensor]) -> None: super().__init__() self.alpha_val = alpha_val @@ -330,6 +339,9 @@ class BgrToRgba(Module): >>> output = rgba(input) # 2x4x4x5 """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 4, -1, -1] + def __init__(self, alpha_val: Union[float, Tensor]) -> None: super().__init__() self.alpha_val = alpha_val @@ -356,6 +368,9 @@ class RgbaToRgb(Module): >>> output = rgba(input) # 2x3x4x5 """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 4, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def forward(self, image: Tensor) -> Tensor: return rgba_to_rgb(image) @@ -378,6 +393,9 @@ class RgbaToBgr(Module): >>> output = rgba(input) # 2x3x4x5 """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 4, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def forward(self, image: Tensor) -> Tensor: return rgba_to_bgr(image) @@ -408,6 +426,9 @@ class RgbToLinearRgb(Module): [3] https://en.wikipedia.org/wiki/SRGB """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def forward(self, image: Tensor) -> Tensor: return rgb_to_linear_rgb(image) @@ -437,5 +458,8 @@ class LinearRgbToRgb(Module): [3] https://en.wikipedia.org/wiki/SRGB """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def forward(self, image: Tensor) -> Tensor: return linear_rgb_to_rgb(image) diff --git a/kornia/color/xyz.py b/kornia/color/xyz.py index 8a611e16bf..0b38d48782 100644 --- a/kornia/color/xyz.py +++ b/kornia/color/xyz.py @@ -91,6 +91,9 @@ class RgbToXyz(Module): [1] https://docs.opencv.org/4.0.1/de/d25/imgproc_color_conversions.html """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def forward(self, image: Tensor) -> Tensor: return rgb_to_xyz(image) @@ -114,5 +117,8 @@ class XyzToRgb(Module): [1] https://docs.opencv.org/4.0.1/de/d25/imgproc_color_conversions.html """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def forward(self, image: Tensor) -> Tensor: return xyz_to_rgb(image) diff --git a/kornia/color/ycbcr.py b/kornia/color/ycbcr.py index 918dad1068..6d839a86fb 100644 --- a/kornia/color/ycbcr.py +++ b/kornia/color/ycbcr.py @@ -121,6 +121,9 @@ class RgbToYcbcr(Module): >>> output = ycbcr(input) # 2x3x4x5 """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def forward(self, image: Tensor) -> Tensor: return rgb_to_ycbcr(image) @@ -143,5 +146,8 @@ class YcbcrToRgb(Module): >>> output = rgb(input) # 2x3x4x5 """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def forward(self, image: Tensor) -> Tensor: return ycbcr_to_rgb(image) diff --git a/kornia/color/yuv.py b/kornia/color/yuv.py index bf620c83ed..493a58123a 100644 --- a/kornia/color/yuv.py +++ b/kornia/color/yuv.py @@ -293,6 +293,9 @@ class RgbToYuv(Module): [1] https://es.wikipedia.org/wiki/YUV#RGB_a_Y'UV """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def forward(self, input: Tensor) -> Tensor: return rgb_to_yuv(input) @@ -324,6 +327,9 @@ class RgbToYuv420(Module): [1] https://es.wikipedia.org/wiki/YUV#RGB_a_Y'UV """ + # TODO: Handle multiple inputs and outputs models later + ONNX_EXPORTABLE = False + def forward(self, yuvinput: Tensor) -> Tuple[Tensor, Tensor]: # skipcq: PYL-R0201 return rgb_to_yuv420(yuvinput) @@ -355,6 +361,9 @@ class RgbToYuv422(Module): [1] https://es.wikipedia.org/wiki/YUV#RGB_a_Y'UV """ + # TODO: Handle multiple inputs and outputs models later + ONNX_EXPORTABLE = False + def forward(self, yuvinput: Tensor) -> Tuple[Tensor, Tensor]: # skipcq: PYL-R0201 return rgb_to_yuv422(yuvinput) @@ -382,6 +391,9 @@ class YuvToRgb(Module): >>> output = rgb(input) # 2x3x4x5 """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def forward(self, input: Tensor) -> Tensor: return yuv_to_rgb(input) @@ -413,6 +425,9 @@ class Yuv420ToRgb(Module): >>> output = rgb(inputy, inputuv) # 2x3x4x6 """ + # TODO: Handle multiple inputs and outputs models later + ONNX_EXPORTABLE = False + def forward(self, inputy: Tensor, inputuv: Tensor) -> Tensor: # skipcq: PYL-R0201 return yuv420_to_rgb(inputy, inputuv) @@ -444,5 +459,8 @@ class Yuv422ToRgb(Module): >>> output = rgb(inputy, inputuv) # 2x3x4x6 """ + # TODO: Handle multiple inputs and outputs models later + ONNX_EXPORTABLE = False + def forward(self, inputy: Tensor, inputuv: Tensor) -> Tensor: # skipcq: PYL-R0201 return yuv422_to_rgb(inputy, inputuv) diff --git a/kornia/core/module.py b/kornia/core/module.py index 97132caf52..648de4f5bd 100644 --- a/kornia/core/module.py +++ b/kornia/core/module.py @@ -4,6 +4,7 @@ from functools import wraps from typing import Any, Callable, List, Optional, Tuple, Union +import torch import kornia from ._backend import Module, Tensor, from_numpy @@ -11,6 +12,53 @@ from .external import numpy as np +class ONNXExportMixin: + + ONNX_EXPORTABLE: bool = True + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, -1, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, -1, -1, -1] + + def to_onnx( + self, + onnx_name: Optional[str] = None, + input_shape: Optional[tuple[int, int, int, int]] = None, + output_shape: Optional[tuple[int, int, int, int]] = None, + ) -> None: + if not self.ONNX_EXPORTABLE: + raise RuntimeError("This object cannot be exported to ONNX.") + + if input_shape is None: + input_shape = self.ONNX_DEFAULT_INPUTSHAPE + if output_shape is None: + output_shape = self.ONNX_DEFAULT_OUTPUTSHAPE + + if onnx_name is None: + onnx_name = f"Kornia-{self.__class__.__name__}.onnx" + + # Creating a dummy input with the given shape + psuedo_shape = (1, 3, 256, 256) + dummy_input = torch.randn(*[ + (psuedo_shape[i] if dim == -1 else dim) for i, dim in enumerate(input_shape)]) + + # Dynamic axis configuration for input and output + dynamic_axes = { + 'input': {i: 'dim_' + str(i) for i, dim in enumerate(input_shape) if dim == -1}, + 'output': {i: 'dim_' + str(i) for i, dim in enumerate(output_shape) if dim == -1} + } + + torch.onnx.export( + self, + dummy_input, + onnx_name, + export_params=True, + opset_version=17, + do_constant_folding=True, + input_names=["input"], + output_names=["output"], + dynamic_axes=dynamic_axes, + ) + + class ImageModuleMixIn: """A MixIn that handles image-based operations. @@ -219,7 +267,7 @@ def save(self, name: Optional[str] = None, n_row: Optional[int] = None) -> None: kornia.io.write_image(name, out_image.mul(255.0).byte()) -class ImageModule(Module, ImageModuleMixIn): +class ImageModule(Module, ImageModuleMixIn, ONNXExportMixin): """Handles image-based operations. This modules accepts multiple input and output data types, provides end-to-end diff --git a/kornia/enhance/adjust.py b/kornia/enhance/adjust.py index 82b2bb0d0b..ccf9581184 100644 --- a/kornia/enhance/adjust.py +++ b/kornia/enhance/adjust.py @@ -1042,6 +1042,9 @@ class AdjustSaturation(Module): tensor(0.) """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def __init__(self, saturation_factor: Union[float, Tensor]) -> None: super().__init__() self.saturation_factor: Union[float, Tensor] = saturation_factor @@ -1089,6 +1092,9 @@ class AdjustSaturationWithGraySubtraction(Module): tensor(0.) """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def __init__(self, saturation_factor: Union[float, Tensor]) -> None: super().__init__() self.saturation_factor: Union[float, Tensor] = saturation_factor @@ -1136,6 +1142,9 @@ class AdjustHue(Module): torch.Size([2, 3, 3, 3]) """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + def __init__(self, hue_factor: Union[float, Tensor]) -> None: super().__init__() self.hue_factor: Union[float, Tensor] = hue_factor diff --git a/kornia/filters/canny.py b/kornia/filters/canny.py index 35515bfbff..ff4f315b70 100644 --- a/kornia/filters/canny.py +++ b/kornia/filters/canny.py @@ -170,6 +170,9 @@ class Canny(Module): torch.Size([5, 1, 4, 4]) """ + # TODO: Handle multiple inputs and outputs models later + ONNX_EXPORTABLE = False + def __init__( self, low_threshold: float = 0.1, diff --git a/kornia/filters/dexined.py b/kornia/filters/dexined.py index d28d3f6930..7c34fa9cf7 100644 --- a/kornia/filters/dexined.py +++ b/kornia/filters/dexined.py @@ -180,6 +180,9 @@ class DexiNed(Module): torch.Size([1, 1, 320, 320]) """ + # TODO: Handle multiple inputs and outputs models later + ONNX_EXPORTABLE = False + def __init__(self, pretrained: bool) -> None: super().__init__() self.block_1 = DoubleConvBlock(3, 32, 64, stride=2) diff --git a/kornia/filters/motion.py b/kornia/filters/motion.py index 985004101b..d7f60f2024 100644 --- a/kornia/filters/motion.py +++ b/kornia/filters/motion.py @@ -84,6 +84,9 @@ class MotionBlur3D(Module): >>> output = motion_blur(input) # 2x4x5x7x9 """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, -1, -1, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, -1, -1, -1, -1] + def __init__( self, kernel_size: int, diff --git a/kornia/filters/sobel.py b/kornia/filters/sobel.py index f18e61cb71..0b247f0f6d 100644 --- a/kornia/filters/sobel.py +++ b/kornia/filters/sobel.py @@ -172,6 +172,9 @@ class SpatialGradient(Module): >>> output = SpatialGradient()(input) # 1x3x2x4x4 """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, -1, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, -1, 2, -1, -1] + def __init__(self, mode: str = "sobel", order: int = 1, normalized: bool = True) -> None: super().__init__() self.normalized: bool = normalized @@ -206,6 +209,9 @@ class SpatialGradient3d(Module): torch.Size([1, 4, 3, 2, 4, 4]) """ + ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, -1, -1, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, -1, -1, -1, -1, -1] + def __init__(self, mode: str = "diff", order: int = 1) -> None: super().__init__() self.order: int = order diff --git a/kornia/onnx/sequential.py b/kornia/onnx/sequential.py index 223294912b..3c0090a558 100644 --- a/kornia/onnx/sequential.py +++ b/kornia/onnx/sequential.py @@ -4,6 +4,8 @@ from kornia.core.external import onnx from kornia.core.external import onnxruntime as ort +from .utils import ONNXLoader + __all__ = ["ONNXSequential"] @@ -22,6 +24,17 @@ class ONNXSequential: If None, we assume the default input name and output name are "input" and "output" accordingly, and only one input and output node for each graph. If not None, `io_maps[0]` shall represent the `io_map` for combining the first and second ONNX models. + cache_dir: + cache_dir: The directory where ONNX models are cached locally (only for downloading from HuggingFace). + Defaults to None, which will use a default `.kornia_onnx_models` directory. + + .. code-block:: python + # Load ops from HuggingFace repos then chain to your own model! + model = kornia.onnx.ONNXSequential( + "hf://operators/kornia.color.gray.RgbToGrayscale", + "hf://operators/kornia.geometry.transform.affwarp.Resize_512x512", + "MY_OTHER_MODEL.onnx" + ) """ def __init__( @@ -30,7 +43,9 @@ def __init__( providers: Optional[list[str]] = None, session_options: Optional[ort.SessionOptions] = None, # type:ignore io_maps: Optional[list[tuple[str, str]]] = None, + cache_dir: str = None ) -> None: + self.onnx_loader = ONNXLoader(cache_dir) self.operators = args self._combined_op = self._combine(io_maps) self._session = self.create_session() @@ -45,7 +60,7 @@ def _load_op(self, arg: Union[onnx.ModelProto, str]) -> onnx.ModelProto: # type onnx.ModelProto: The loaded ONNX model. """ if isinstance(arg, str): - return onnx.load(arg) # type:ignore + return self.onnx_loader.load_model(arg) # type:ignore return arg def _combine(self, io_maps: Optional[list[tuple[str, str]]] = None) -> onnx.ModelProto: # type:ignore diff --git a/kornia/onnx/utils.py b/kornia/onnx/utils.py new file mode 100644 index 0000000000..3ba4aea4b3 --- /dev/null +++ b/kornia/onnx/utils.py @@ -0,0 +1,97 @@ +from typing import Optional +import os +import urllib.request + +from kornia.core.external import onnx + + +class ONNXLoader: + """Manages ONNX models, handling local caching, downloading from Hugging Face, and loading models. + + Attributes: + cache_dir: The directory where ONNX models are cached locally. + Defaults to None, which will use a default `.kornia_onnx_models` directory. + + .. code-block:: python + onnx_loader = ONNXLoader() + # Load a HuggingFace operator + onnx_loader.load_model("hf://operators/kornia.color.gray.GrayscaleToRgb") + # Load a local converted/downloaded operator + onnx_loader.load_model("operators/kornia.color.gray.GrayscaleToRgb") + """ + def __init__(self, cache_dir: Optional[str] = None): + self.cache_dir = cache_dir + + def _get_file_path(self, model_name: str, cache_dir: str) -> str: + """Constructs the file path for the ONNX model based on the model name and cache directory. + + Args: + model_name: The name of the model or operator, typically in the format 'operators/model_name'. + cache_dir: The directory where the model should be cached. + + Returns: + str: The full local path where the model should be stored or loaded from. + """ + # Determine the local file path + if cache_dir is None: + if self.cache_dir is not None: + cache_dir = self.cache_dir + else: + cache_dir = ".kornia_onnx_models" + + # The filename is the model name (without directory path) + file_name = f"{model_name.split('/')[-1]}.onnx" + file_path = os.path.join(cache_dir, '/'.join(model_name.split('/')[:-1]), file_name) + return file_path + + def load_model( + self, model_name: str, download: bool = False, **kwargs + ) -> onnx.ModelProto: + """Loads an ONNX model from the local cache or downloads it from Hugging Face if necessary. + + Args: + model_name: The name of the ONNX model or operator. For Hugging Face-hosted models, + use the format 'hf://model_name'. Valid `model_name` can be found on + https://huggingface.co/kornia/ONNX_models. + download: If True, the model will be downloaded from Hugging Face if it's not already in the local cache. + **kwargs: Additional arguments to pass to the download method, if needed. + + Returns: + onnx.ModelProto: The loaded ONNX model. + """ + if model_name.startswith("hf://"): + model_name = model_name[len("hf://"):] + file_path = self._get_file_path(model_name, self.cache_dir) + if not os.path.exists(file_path): + # Construct the raw URL for the ONNX file + url = f"https://huggingface.co/kornia/ONNX_models/resolve/main/{model_name}.onnx" + self.download(url, file_path, **kwargs) + return onnx.load(file_path) # type:ignore + + if os.path.exists(model_name): + return onnx.load(model_name) # type:ignore + + raise ValueError(f"File {model_name} not found") + + def download( + self, + url: str, + file_path: str, + cache_dir: Optional[str] = None, + ) -> None: + """Downloads an ONNX model from the specified URL and saves it to the specified file path. + + Args: + url: The URL of the ONNX model to download. + file_path: The local path where the downloaded model should be saved. + cache_dir: The directory to use for caching the file, defaults to the instance cache + directory if not provided. + """ + + os.makedirs(os.path.dirname(file_path), exist_ok=True) # Create the cache directory if it doesn't exist + + # Download the file and save it + try: + urllib.request.urlretrieve(url, file_path) + except urllib.error.HTTPError as e: + raise ValueError(f"Error in resolving `{url}`. {e}.") From fa864ca1a66995f548cfecb0b63db1ddfa4e9156 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 7 Sep 2024 15:07:59 +0000 Subject: [PATCH 06/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/core/module.py | 9 ++++----- kornia/onnx/sequential.py | 2 +- kornia/onnx/utils.py | 13 ++++++------- 3 files changed, 11 insertions(+), 13 deletions(-) diff --git a/kornia/core/module.py b/kornia/core/module.py index 648de4f5bd..feb6b56f44 100644 --- a/kornia/core/module.py +++ b/kornia/core/module.py @@ -5,6 +5,7 @@ from typing import Any, Callable, List, Optional, Tuple, Union import torch + import kornia from ._backend import Module, Tensor, from_numpy @@ -13,7 +14,6 @@ class ONNXExportMixin: - ONNX_EXPORTABLE: bool = True ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, -1, -1, -1] ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, -1, -1, -1] @@ -37,13 +37,12 @@ def to_onnx( # Creating a dummy input with the given shape psuedo_shape = (1, 3, 256, 256) - dummy_input = torch.randn(*[ - (psuedo_shape[i] if dim == -1 else dim) for i, dim in enumerate(input_shape)]) + dummy_input = torch.randn(*[(psuedo_shape[i] if dim == -1 else dim) for i, dim in enumerate(input_shape)]) # Dynamic axis configuration for input and output dynamic_axes = { - 'input': {i: 'dim_' + str(i) for i, dim in enumerate(input_shape) if dim == -1}, - 'output': {i: 'dim_' + str(i) for i, dim in enumerate(output_shape) if dim == -1} + "input": {i: "dim_" + str(i) for i, dim in enumerate(input_shape) if dim == -1}, + "output": {i: "dim_" + str(i) for i, dim in enumerate(output_shape) if dim == -1}, } torch.onnx.export( diff --git a/kornia/onnx/sequential.py b/kornia/onnx/sequential.py index 3c0090a558..d3e8e9cb31 100644 --- a/kornia/onnx/sequential.py +++ b/kornia/onnx/sequential.py @@ -43,7 +43,7 @@ def __init__( providers: Optional[list[str]] = None, session_options: Optional[ort.SessionOptions] = None, # type:ignore io_maps: Optional[list[tuple[str, str]]] = None, - cache_dir: str = None + cache_dir: str = None, ) -> None: self.onnx_loader = ONNXLoader(cache_dir) self.operators = args diff --git a/kornia/onnx/utils.py b/kornia/onnx/utils.py index 3ba4aea4b3..687dd490e4 100644 --- a/kornia/onnx/utils.py +++ b/kornia/onnx/utils.py @@ -1,6 +1,6 @@ -from typing import Optional import os import urllib.request +from typing import Optional from kornia.core.external import onnx @@ -19,6 +19,7 @@ class ONNXLoader: # Load a local converted/downloaded operator onnx_loader.load_model("operators/kornia.color.gray.GrayscaleToRgb") """ + def __init__(self, cache_dir: Optional[str] = None): self.cache_dir = cache_dir @@ -41,12 +42,10 @@ def _get_file_path(self, model_name: str, cache_dir: str) -> str: # The filename is the model name (without directory path) file_name = f"{model_name.split('/')[-1]}.onnx" - file_path = os.path.join(cache_dir, '/'.join(model_name.split('/')[:-1]), file_name) + file_path = os.path.join(cache_dir, "/".join(model_name.split("/")[:-1]), file_name) return file_path - def load_model( - self, model_name: str, download: bool = False, **kwargs - ) -> onnx.ModelProto: + def load_model(self, model_name: str, download: bool = False, **kwargs) -> onnx.ModelProto: """Loads an ONNX model from the local cache or downloads it from Hugging Face if necessary. Args: @@ -60,7 +59,7 @@ def load_model( onnx.ModelProto: The loaded ONNX model. """ if model_name.startswith("hf://"): - model_name = model_name[len("hf://"):] + model_name = model_name[len("hf://") :] file_path = self._get_file_path(model_name, self.cache_dir) if not os.path.exists(file_path): # Construct the raw URL for the ONNX file @@ -70,7 +69,7 @@ def load_model( if os.path.exists(model_name): return onnx.load(model_name) # type:ignore - + raise ValueError(f"File {model_name} not found") def download( From 68d5b777bb135235b132c974388c3e1e4015156f Mon Sep 17 00:00:00 2001 From: shijianjian Date: Sat, 7 Sep 2024 18:29:27 +0300 Subject: [PATCH 07/84] update --- kornia/onnx/__init__.py | 1 + kornia/onnx/utils.py | 52 ++++++++++++++++++++++++++++++++++++++++- 2 files changed, 52 insertions(+), 1 deletion(-) diff --git a/kornia/onnx/__init__.py b/kornia/onnx/__init__.py index bce2460a4f..40358f85b2 100644 --- a/kornia/onnx/__init__.py +++ b/kornia/onnx/__init__.py @@ -1 +1,2 @@ from .sequential import * +from .utils import * diff --git a/kornia/onnx/utils.py b/kornia/onnx/utils.py index 687dd490e4..28c8430821 100644 --- a/kornia/onnx/utils.py +++ b/kornia/onnx/utils.py @@ -1,9 +1,13 @@ +from typing import Optional +import requests +import pprint import os import urllib.request -from typing import Optional from kornia.core.external import onnx +__all__ = ["ONNXLoader"] + class ONNXLoader: """Manages ONNX models, handling local caching, downloading from Hugging Face, and loading models. @@ -94,3 +98,49 @@ def download( urllib.request.urlretrieve(url, file_path) except urllib.error.HTTPError as e: raise ValueError(f"Error in resolving `{url}`. {e}.") + + @staticmethod + def _fetch_repo_contents(folder: str) -> list[dict]: + """ + Fetches the contents of the Hugging Face repository using the Hugging Face API. + + Returns: + List[dict]: A list of all files in the repository as dictionaries containing file details. + """ + url = f"https://huggingface.co/api/models/kornia/ONNX_models/tree/main/{folder}" + response = requests.get(url) + + if response.status_code == 200: + return response.json() # Returns the JSON content of the repo + else: + raise ValueError(f"Failed to fetch repository contents: {response.status_code}") + + @staticmethod + def list_operators() -> list[str]: + """ + Lists all available ONNX operators in the 'operators' folder of the Hugging Face repository. + + Returns: + List[str]: A list of operator file paths. + """ + repo_contents = ONNXLoader._fetch_repo_contents("operators") + + # Filter for operators in the 'operators' directory + operators = [file['path'] for file in repo_contents] + + pprint.pp(operators) + + @staticmethod + def list_models() -> list[str]: + """ + Lists all available ONNX models in the 'models' folder of the Hugging Face repository. + + Returns: + List[str]: A list of model file paths. + """ + repo_contents = ONNXLoader._fetch_repo_contents("models") + + # Filter for models in the 'models' directory + models = [file['path'] for file in repo_contents] + + pprint.pp(models) From c41ef8d1db92e7f7ddff3ff15f248daa59f17070 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 7 Sep 2024 15:30:43 +0000 Subject: [PATCH 08/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/onnx/utils.py | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/kornia/onnx/utils.py b/kornia/onnx/utils.py index 28c8430821..cc22edf1c9 100644 --- a/kornia/onnx/utils.py +++ b/kornia/onnx/utils.py @@ -1,8 +1,9 @@ -from typing import Optional -import requests -import pprint import os +import pprint import urllib.request +from typing import Optional + +import requests from kornia.core.external import onnx @@ -101,8 +102,7 @@ def download( @staticmethod def _fetch_repo_contents(folder: str) -> list[dict]: - """ - Fetches the contents of the Hugging Face repository using the Hugging Face API. + """Fetches the contents of the Hugging Face repository using the Hugging Face API. Returns: List[dict]: A list of all files in the repository as dictionaries containing file details. @@ -117,8 +117,7 @@ def _fetch_repo_contents(folder: str) -> list[dict]: @staticmethod def list_operators() -> list[str]: - """ - Lists all available ONNX operators in the 'operators' folder of the Hugging Face repository. + """Lists all available ONNX operators in the 'operators' folder of the Hugging Face repository. Returns: List[str]: A list of operator file paths. @@ -126,14 +125,13 @@ def list_operators() -> list[str]: repo_contents = ONNXLoader._fetch_repo_contents("operators") # Filter for operators in the 'operators' directory - operators = [file['path'] for file in repo_contents] + operators = [file["path"] for file in repo_contents] pprint.pp(operators) @staticmethod def list_models() -> list[str]: - """ - Lists all available ONNX models in the 'models' folder of the Hugging Face repository. + """Lists all available ONNX models in the 'models' folder of the Hugging Face repository. Returns: List[str]: A list of model file paths. @@ -141,6 +139,6 @@ def list_models() -> list[str]: repo_contents = ONNXLoader._fetch_repo_contents("models") # Filter for models in the 'models' directory - models = [file['path'] for file in repo_contents] + models = [file["path"] for file in repo_contents] pprint.pp(models) From 3765390979e6617b175bf2002d7acc8ce6fca70d Mon Sep 17 00:00:00 2001 From: shijianjian Date: Sat, 7 Sep 2024 18:58:10 +0300 Subject: [PATCH 09/84] fix typing --- kornia/color/gray.py | 12 ++++++------ kornia/color/hls.py | 8 ++++---- kornia/color/hsv.py | 8 ++++---- kornia/color/lab.py | 8 ++++---- kornia/color/luv.py | 8 ++++---- kornia/color/raw.py | 8 ++++---- kornia/color/rgb.py | 32 +++++++++++++++---------------- kornia/color/xyz.py | 8 ++++---- kornia/color/ycbcr.py | 8 ++++---- kornia/color/yuv.py | 8 ++++---- kornia/core/module.py | 10 +++++----- kornia/enhance/adjust.py | 12 ++++++------ kornia/filters/motion.py | 4 ++-- kornia/filters/sobel.py | 8 ++++---- kornia/onnx/sequential.py | 4 ++-- kornia/onnx/utils.py | 23 ++++++++++++++++++++-- requirements/requirements-dev.txt | 3 ++- 17 files changed, 96 insertions(+), 76 deletions(-) diff --git a/kornia/color/gray.py b/kornia/color/gray.py index 2fe6fd7404..96d2b63dac 100644 --- a/kornia/color/gray.py +++ b/kornia/color/gray.py @@ -125,8 +125,8 @@ class GrayscaleToRgb(Module): >>> output = rgb(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 1, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 1, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return grayscale_to_rgb(image) @@ -150,8 +150,8 @@ class RgbToGrayscale(Module): >>> output = gray(input) # 2x1x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 1, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 1, -1, -1] def __init__(self, rgb_weights: Optional[Tensor] = None) -> None: super().__init__() @@ -181,8 +181,8 @@ class BgrToGrayscale(Module): >>> output = gray(input) # 2x1x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 1, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 1, -1, -1] def forward(self, image: Tensor) -> Tensor: return bgr_to_grayscale(image) diff --git a/kornia/color/hls.py b/kornia/color/hls.py index 0f498dc66d..f06448fd38 100644 --- a/kornia/color/hls.py +++ b/kornia/color/hls.py @@ -142,8 +142,8 @@ class RgbToHls(Module): >>> output = hls(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return rgb_to_hls(image) @@ -170,8 +170,8 @@ class HlsToRgb(Module): >>> output = rgb(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return hls_to_rgb(image) diff --git a/kornia/color/hsv.py b/kornia/color/hsv.py index 561e585ec9..ce295beae8 100644 --- a/kornia/color/hsv.py +++ b/kornia/color/hsv.py @@ -116,8 +116,8 @@ class RgbToHsv(Module): >>> output = hsv(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def __init__(self, eps: float = 1e-6) -> None: super().__init__() @@ -145,8 +145,8 @@ class HsvToRgb(Module): >>> output = rgb(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def forward(self, image: torch.Tensor) -> torch.Tensor: return hsv_to_rgb(image) diff --git a/kornia/color/lab.py b/kornia/color/lab.py index 0231a85b2d..3defcb2478 100644 --- a/kornia/color/lab.py +++ b/kornia/color/lab.py @@ -150,8 +150,8 @@ class RgbToLab(Module): [3] https://github.com/torch/image/blob/dc061b98fb7e946e00034a5fc73e883a299edc7f/generic/image.c#L1467 """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def forward(self, image: torch.Tensor) -> torch.Tensor: return rgb_to_lab(image) @@ -180,8 +180,8 @@ class LabToRgb(Module): [3] https://github.com/torch/image/blob/dc061b98fb7e946e00034a5fc73e883a299edc7f/generic/image.c#L1518 """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def forward(self, image: torch.Tensor, clip: bool = True) -> torch.Tensor: return lab_to_rgb(image, clip) diff --git a/kornia/color/luv.py b/kornia/color/luv.py index f5055d1ebe..bfec8c41b3 100644 --- a/kornia/color/luv.py +++ b/kornia/color/luv.py @@ -141,8 +141,8 @@ class RgbToLuv(Module): [3] http://www.poynton.com/ColorFAQ.html """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def forward(self, image: torch.Tensor) -> torch.Tensor: return rgb_to_luv(image) @@ -171,8 +171,8 @@ class LuvToRgb(Module): [3] http://www.poynton.com/ColorFAQ.html """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def forward(self, image: torch.Tensor) -> torch.Tensor: return luv_to_rgb(image) diff --git a/kornia/color/raw.py b/kornia/color/raw.py index a0d68d0d50..aabaa03289 100644 --- a/kornia/color/raw.py +++ b/kornia/color/raw.py @@ -288,8 +288,8 @@ class RawToRgb(Module): >>> output = rgb(rawinput) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 1, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 1, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def __init__(self, cfa: CFA) -> None: super().__init__() @@ -317,8 +317,8 @@ class RgbToRaw(Module): >>> output = raw(rgbinput) # 2x1x4x6 """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 1, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 1, -1, -1] def __init__(self, cfa: CFA) -> None: super().__init__() diff --git a/kornia/color/rgb.py b/kornia/color/rgb.py index 5003ccbafb..0217d98507 100644 --- a/kornia/color/rgb.py +++ b/kornia/color/rgb.py @@ -248,8 +248,8 @@ class BgrToRgb(Module): >>> output = rgb(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return bgr_to_rgb(image) @@ -273,8 +273,8 @@ class RgbToBgr(Module): >>> output = bgr(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return rgb_to_bgr(image) @@ -304,8 +304,8 @@ class RgbToRgba(Module): >>> output = rgba(input) # 2x4x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 4, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 4, -1, -1] def __init__(self, alpha_val: Union[float, Tensor]) -> None: super().__init__() @@ -339,8 +339,8 @@ class BgrToRgba(Module): >>> output = rgba(input) # 2x4x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 4, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 4, -1, -1] def __init__(self, alpha_val: Union[float, Tensor]) -> None: super().__init__() @@ -368,8 +368,8 @@ class RgbaToRgb(Module): >>> output = rgba(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 4, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 4, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return rgba_to_rgb(image) @@ -393,8 +393,8 @@ class RgbaToBgr(Module): >>> output = rgba(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 4, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 4, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return rgba_to_bgr(image) @@ -426,8 +426,8 @@ class RgbToLinearRgb(Module): [3] https://en.wikipedia.org/wiki/SRGB """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return rgb_to_linear_rgb(image) @@ -458,8 +458,8 @@ class LinearRgbToRgb(Module): [3] https://en.wikipedia.org/wiki/SRGB """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return linear_rgb_to_rgb(image) diff --git a/kornia/color/xyz.py b/kornia/color/xyz.py index 0b38d48782..c5493840db 100644 --- a/kornia/color/xyz.py +++ b/kornia/color/xyz.py @@ -91,8 +91,8 @@ class RgbToXyz(Module): [1] https://docs.opencv.org/4.0.1/de/d25/imgproc_color_conversions.html """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return rgb_to_xyz(image) @@ -117,8 +117,8 @@ class XyzToRgb(Module): [1] https://docs.opencv.org/4.0.1/de/d25/imgproc_color_conversions.html """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return xyz_to_rgb(image) diff --git a/kornia/color/ycbcr.py b/kornia/color/ycbcr.py index 6d839a86fb..aac42300eb 100644 --- a/kornia/color/ycbcr.py +++ b/kornia/color/ycbcr.py @@ -121,8 +121,8 @@ class RgbToYcbcr(Module): >>> output = ycbcr(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return rgb_to_ycbcr(image) @@ -146,8 +146,8 @@ class YcbcrToRgb(Module): >>> output = rgb(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return ycbcr_to_rgb(image) diff --git a/kornia/color/yuv.py b/kornia/color/yuv.py index 493a58123a..3e94b6615e 100644 --- a/kornia/color/yuv.py +++ b/kornia/color/yuv.py @@ -293,8 +293,8 @@ class RgbToYuv(Module): [1] https://es.wikipedia.org/wiki/YUV#RGB_a_Y'UV """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def forward(self, input: Tensor) -> Tensor: return rgb_to_yuv(input) @@ -391,8 +391,8 @@ class YuvToRgb(Module): >>> output = rgb(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def forward(self, input: Tensor) -> Tensor: return yuv_to_rgb(input) diff --git a/kornia/core/module.py b/kornia/core/module.py index feb6b56f44..00e315561e 100644 --- a/kornia/core/module.py +++ b/kornia/core/module.py @@ -15,14 +15,14 @@ class ONNXExportMixin: ONNX_EXPORTABLE: bool = True - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, -1, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, -1, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, -1, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, -1, -1, -1] def to_onnx( self, onnx_name: Optional[str] = None, - input_shape: Optional[tuple[int, int, int, int]] = None, - output_shape: Optional[tuple[int, int, int, int]] = None, + input_shape: Optional[list[int]] = None, + output_shape: Optional[list[int]] = None, ) -> None: if not self.ONNX_EXPORTABLE: raise RuntimeError("This object cannot be exported to ONNX.") @@ -46,7 +46,7 @@ def to_onnx( } torch.onnx.export( - self, + self, # type: ignore dummy_input, onnx_name, export_params=True, diff --git a/kornia/enhance/adjust.py b/kornia/enhance/adjust.py index ccf9581184..41c2ced013 100644 --- a/kornia/enhance/adjust.py +++ b/kornia/enhance/adjust.py @@ -1042,8 +1042,8 @@ class AdjustSaturation(Module): tensor(0.) """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def __init__(self, saturation_factor: Union[float, Tensor]) -> None: super().__init__() @@ -1092,8 +1092,8 @@ class AdjustSaturationWithGraySubtraction(Module): tensor(0.) """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def __init__(self, saturation_factor: Union[float, Tensor]) -> None: super().__init__() @@ -1142,8 +1142,8 @@ class AdjustHue(Module): torch.Size([2, 3, 3, 3]) """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] def __init__(self, hue_factor: Union[float, Tensor]) -> None: super().__init__() diff --git a/kornia/filters/motion.py b/kornia/filters/motion.py index d7f60f2024..d6eb4d1781 100644 --- a/kornia/filters/motion.py +++ b/kornia/filters/motion.py @@ -84,8 +84,8 @@ class MotionBlur3D(Module): >>> output = motion_blur(input) # 2x4x5x7x9 """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, -1, -1, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, -1, -1, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, -1, -1, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, -1, -1, -1, -1] def __init__( self, diff --git a/kornia/filters/sobel.py b/kornia/filters/sobel.py index 0b247f0f6d..897cb3d9a6 100644 --- a/kornia/filters/sobel.py +++ b/kornia/filters/sobel.py @@ -172,8 +172,8 @@ class SpatialGradient(Module): >>> output = SpatialGradient()(input) # 1x3x2x4x4 """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, -1, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, -1, 2, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, -1, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, -1, 2, -1, -1] def __init__(self, mode: str = "sobel", order: int = 1, normalized: bool = True) -> None: super().__init__() @@ -209,8 +209,8 @@ class SpatialGradient3d(Module): torch.Size([1, 4, 3, 2, 4, 4]) """ - ONNX_DEFAULT_INPUTSHAPE: tuple[int, int, int, int] = [-1, -1, -1, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: tuple[int, int, int, int] = [-1, -1, -1, -1, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, -1, -1, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, -1, -1, -1, -1, -1] def __init__(self, mode: str = "diff", order: int = 1) -> None: super().__init__() diff --git a/kornia/onnx/sequential.py b/kornia/onnx/sequential.py index d3e8e9cb31..98b9573e8a 100644 --- a/kornia/onnx/sequential.py +++ b/kornia/onnx/sequential.py @@ -43,7 +43,7 @@ def __init__( providers: Optional[list[str]] = None, session_options: Optional[ort.SessionOptions] = None, # type:ignore io_maps: Optional[list[tuple[str, str]]] = None, - cache_dir: str = None, + cache_dir: Optional[str] = None, ) -> None: self.onnx_loader = ONNXLoader(cache_dir) self.operators = args @@ -60,7 +60,7 @@ def _load_op(self, arg: Union[onnx.ModelProto, str]) -> onnx.ModelProto: # type onnx.ModelProto: The loaded ONNX model. """ if isinstance(arg, str): - return self.onnx_loader.load_model(arg) # type:ignore + return self.onnx_loader.load_model(arg) return arg def _combine(self, io_maps: Optional[list[tuple[str, str]]] = None) -> onnx.ModelProto: # type:ignore diff --git a/kornia/onnx/utils.py b/kornia/onnx/utils.py index cc22edf1c9..fb55d352ac 100644 --- a/kornia/onnx/utils.py +++ b/kornia/onnx/utils.py @@ -1,3 +1,6 @@ +from typing import Any, Optional +import requests +import pprint import os import pprint import urllib.request @@ -28,7 +31,7 @@ class ONNXLoader: def __init__(self, cache_dir: Optional[str] = None): self.cache_dir = cache_dir - def _get_file_path(self, model_name: str, cache_dir: str) -> str: + def _get_file_path(self, model_name: str, cache_dir: Optional[str]) -> str: """Constructs the file path for the ONNX model based on the model name and cache directory. Args: @@ -50,7 +53,7 @@ def _get_file_path(self, model_name: str, cache_dir: str) -> str: file_path = os.path.join(cache_dir, "/".join(model_name.split("/")[:-1]), file_name) return file_path - def load_model(self, model_name: str, download: bool = False, **kwargs) -> onnx.ModelProto: + def load_model(self, model_name: str, download: bool = False, **kwargs) -> onnx.ModelProto: # type:ignore """Loads an ONNX model from the local cache or downloads it from Hugging Face if necessary. Args: @@ -101,8 +104,14 @@ def download( raise ValueError(f"Error in resolving `{url}`. {e}.") @staticmethod +<<<<<<< HEAD def _fetch_repo_contents(folder: str) -> list[dict]: """Fetches the contents of the Hugging Face repository using the Hugging Face API. +======= + def _fetch_repo_contents(folder: str) -> list[dict[str, Any]]: + """ + Fetches the contents of the Hugging Face repository using the Hugging Face API. +>>>>>>> 1b3c0dbb (fix typing) Returns: List[dict]: A list of all files in the repository as dictionaries containing file details. @@ -116,12 +125,17 @@ def _fetch_repo_contents(folder: str) -> list[dict]: raise ValueError(f"Failed to fetch repository contents: {response.status_code}") @staticmethod +<<<<<<< HEAD def list_operators() -> list[str]: """Lists all available ONNX operators in the 'operators' folder of the Hugging Face repository. Returns: List[str]: A list of operator file paths. """ +======= + def list_operators() -> None: + """Lists all available ONNX operators in the 'operators' folder of the Hugging Face repository.""" +>>>>>>> 1b3c0dbb (fix typing) repo_contents = ONNXLoader._fetch_repo_contents("operators") # Filter for operators in the 'operators' directory @@ -130,12 +144,17 @@ def list_operators() -> list[str]: pprint.pp(operators) @staticmethod +<<<<<<< HEAD def list_models() -> list[str]: """Lists all available ONNX models in the 'models' folder of the Hugging Face repository. Returns: List[str]: A list of model file paths. """ +======= + def list_models() -> None: + """Lists all available ONNX models in the 'models' folder of the Hugging Face repository.""" +>>>>>>> 1b3c0dbb (fix typing) repo_contents = ONNXLoader._fetch_repo_contents("models") # Filter for models in the 'models' directory diff --git a/requirements/requirements-dev.txt b/requirements/requirements-dev.txt index 6d5dac3c63..60ee6d4559 100644 --- a/requirements/requirements-dev.txt +++ b/requirements/requirements-dev.txt @@ -2,7 +2,8 @@ accelerate coverage diffusers mypy -numpy<2 +numpy<3 +types-requests onnx pillow pre-commit>=2 From 97d816efdb803acf8d407ba426a9875aa95a01e9 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Sat, 7 Sep 2024 18:59:45 +0300 Subject: [PATCH 10/84] update --- kornia/onnx/utils.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/kornia/onnx/utils.py b/kornia/onnx/utils.py index fb55d352ac..3c13140b38 100644 --- a/kornia/onnx/utils.py +++ b/kornia/onnx/utils.py @@ -1,12 +1,8 @@ -from typing import Any, Optional import requests import pprint import os -import pprint import urllib.request -from typing import Optional - -import requests +from typing import Any, Optional from kornia.core.external import onnx From df2dfd20ec26893d30fb30b0f7ee47aeee0449fd Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 7 Sep 2024 16:00:17 +0000 Subject: [PATCH 11/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- requirements/requirements-dev.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/requirements-dev.txt b/requirements/requirements-dev.txt index 60ee6d4559..bd3222ebc9 100644 --- a/requirements/requirements-dev.txt +++ b/requirements/requirements-dev.txt @@ -3,7 +3,6 @@ coverage diffusers mypy numpy<3 -types-requests onnx pillow pre-commit>=2 @@ -11,3 +10,4 @@ pytest==8.3.2 pytest-timeout requests transformers +types-requests From d1411dd298e61ed9e3060adc4aa38ce68aaad669 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Sat, 7 Sep 2024 19:04:00 +0300 Subject: [PATCH 12/84] update --- kornia/onnx/utils.py | 23 ----------------------- 1 file changed, 23 deletions(-) diff --git a/kornia/onnx/utils.py b/kornia/onnx/utils.py index 3c13140b38..939f8dc172 100644 --- a/kornia/onnx/utils.py +++ b/kornia/onnx/utils.py @@ -100,14 +100,9 @@ def download( raise ValueError(f"Error in resolving `{url}`. {e}.") @staticmethod -<<<<<<< HEAD - def _fetch_repo_contents(folder: str) -> list[dict]: - """Fetches the contents of the Hugging Face repository using the Hugging Face API. -======= def _fetch_repo_contents(folder: str) -> list[dict[str, Any]]: """ Fetches the contents of the Hugging Face repository using the Hugging Face API. ->>>>>>> 1b3c0dbb (fix typing) Returns: List[dict]: A list of all files in the repository as dictionaries containing file details. @@ -121,17 +116,8 @@ def _fetch_repo_contents(folder: str) -> list[dict[str, Any]]: raise ValueError(f"Failed to fetch repository contents: {response.status_code}") @staticmethod -<<<<<<< HEAD - def list_operators() -> list[str]: - """Lists all available ONNX operators in the 'operators' folder of the Hugging Face repository. - - Returns: - List[str]: A list of operator file paths. - """ -======= def list_operators() -> None: """Lists all available ONNX operators in the 'operators' folder of the Hugging Face repository.""" ->>>>>>> 1b3c0dbb (fix typing) repo_contents = ONNXLoader._fetch_repo_contents("operators") # Filter for operators in the 'operators' directory @@ -140,17 +126,8 @@ def list_operators() -> None: pprint.pp(operators) @staticmethod -<<<<<<< HEAD - def list_models() -> list[str]: - """Lists all available ONNX models in the 'models' folder of the Hugging Face repository. - - Returns: - List[str]: A list of model file paths. - """ -======= def list_models() -> None: """Lists all available ONNX models in the 'models' folder of the Hugging Face repository.""" ->>>>>>> 1b3c0dbb (fix typing) repo_contents = ONNXLoader._fetch_repo_contents("models") # Filter for models in the 'models' directory From 89a3e73887abc378d38a23af53708d11013f8fdb Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 7 Sep 2024 16:04:53 +0000 Subject: [PATCH 13/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/onnx/utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/kornia/onnx/utils.py b/kornia/onnx/utils.py index 939f8dc172..8213de73f8 100644 --- a/kornia/onnx/utils.py +++ b/kornia/onnx/utils.py @@ -1,9 +1,10 @@ -import requests -import pprint import os +import pprint import urllib.request from typing import Any, Optional +import requests + from kornia.core.external import onnx __all__ = ["ONNXLoader"] @@ -101,8 +102,7 @@ def download( @staticmethod def _fetch_repo_contents(folder: str) -> list[dict[str, Any]]: - """ - Fetches the contents of the Hugging Face repository using the Hugging Face API. + """Fetches the contents of the Hugging Face repository using the Hugging Face API. Returns: List[dict]: A list of all files in the repository as dictionaries containing file details. From 88e4625f356b9d7e497dc2d6f6c6856f83ff2f04 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Sat, 7 Sep 2024 19:10:52 +0300 Subject: [PATCH 14/84] update --- kornia/color/gray.py | 12 ++++++------ kornia/color/hls.py | 8 ++++---- kornia/color/hsv.py | 8 ++++---- kornia/color/lab.py | 8 ++++---- kornia/color/luv.py | 8 ++++---- kornia/color/raw.py | 8 ++++---- kornia/color/rgb.py | 32 +++++++++++++++---------------- kornia/color/xyz.py | 8 ++++---- kornia/color/ycbcr.py | 8 ++++---- kornia/color/yuv.py | 8 ++++---- kornia/core/module.py | 6 +++--- kornia/enhance/adjust.py | 12 ++++++------ requirements/requirements-dev.txt | 2 ++ 13 files changed, 65 insertions(+), 63 deletions(-) diff --git a/kornia/color/gray.py b/kornia/color/gray.py index 96d2b63dac..45f4c958f0 100644 --- a/kornia/color/gray.py +++ b/kornia/color/gray.py @@ -125,8 +125,8 @@ class GrayscaleToRgb(Module): >>> output = rgb(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 1, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + 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) @@ -150,8 +150,8 @@ class RgbToGrayscale(Module): >>> output = gray(input) # 2x1x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 1, -1, -1] + 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__() @@ -181,8 +181,8 @@ class BgrToGrayscale(Module): >>> output = gray(input) # 2x1x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 1, -1, -1] + 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) diff --git a/kornia/color/hls.py b/kornia/color/hls.py index f06448fd38..7880999e53 100644 --- a/kornia/color/hls.py +++ b/kornia/color/hls.py @@ -142,8 +142,8 @@ class RgbToHls(Module): >>> output = hls(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return rgb_to_hls(image) @@ -170,8 +170,8 @@ class HlsToRgb(Module): >>> output = rgb(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return hls_to_rgb(image) diff --git a/kornia/color/hsv.py b/kornia/color/hsv.py index ce295beae8..37dfe4baf0 100644 --- a/kornia/color/hsv.py +++ b/kornia/color/hsv.py @@ -116,8 +116,8 @@ class RgbToHsv(Module): >>> output = hsv(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def __init__(self, eps: float = 1e-6) -> None: super().__init__() @@ -145,8 +145,8 @@ class HsvToRgb(Module): >>> output = rgb(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def forward(self, image: torch.Tensor) -> torch.Tensor: return hsv_to_rgb(image) diff --git a/kornia/color/lab.py b/kornia/color/lab.py index 3defcb2478..13b25b3be6 100644 --- a/kornia/color/lab.py +++ b/kornia/color/lab.py @@ -150,8 +150,8 @@ class RgbToLab(Module): [3] https://github.com/torch/image/blob/dc061b98fb7e946e00034a5fc73e883a299edc7f/generic/image.c#L1467 """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def forward(self, image: torch.Tensor) -> torch.Tensor: return rgb_to_lab(image) @@ -180,8 +180,8 @@ class LabToRgb(Module): [3] https://github.com/torch/image/blob/dc061b98fb7e946e00034a5fc73e883a299edc7f/generic/image.c#L1518 """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def forward(self, image: torch.Tensor, clip: bool = True) -> torch.Tensor: return lab_to_rgb(image, clip) diff --git a/kornia/color/luv.py b/kornia/color/luv.py index bfec8c41b3..0b336773e9 100644 --- a/kornia/color/luv.py +++ b/kornia/color/luv.py @@ -141,8 +141,8 @@ class RgbToLuv(Module): [3] http://www.poynton.com/ColorFAQ.html """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def forward(self, image: torch.Tensor) -> torch.Tensor: return rgb_to_luv(image) @@ -171,8 +171,8 @@ class LuvToRgb(Module): [3] http://www.poynton.com/ColorFAQ.html """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def forward(self, image: torch.Tensor) -> torch.Tensor: return luv_to_rgb(image) diff --git a/kornia/color/raw.py b/kornia/color/raw.py index aabaa03289..4e040ff052 100644 --- a/kornia/color/raw.py +++ b/kornia/color/raw.py @@ -288,8 +288,8 @@ class RawToRgb(Module): >>> output = rgb(rawinput) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 1, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 1, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def __init__(self, cfa: CFA) -> None: super().__init__() @@ -317,8 +317,8 @@ class RgbToRaw(Module): >>> output = raw(rgbinput) # 2x1x4x6 """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 1, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 1, -1, -1] def __init__(self, cfa: CFA) -> None: super().__init__() diff --git a/kornia/color/rgb.py b/kornia/color/rgb.py index 0217d98507..92f4310e8b 100644 --- a/kornia/color/rgb.py +++ b/kornia/color/rgb.py @@ -248,8 +248,8 @@ class BgrToRgb(Module): >>> output = rgb(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return bgr_to_rgb(image) @@ -273,8 +273,8 @@ class RgbToBgr(Module): >>> output = bgr(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return rgb_to_bgr(image) @@ -304,8 +304,8 @@ class RgbToRgba(Module): >>> output = rgba(input) # 2x4x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 4, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 4, -1, -1] def __init__(self, alpha_val: Union[float, Tensor]) -> None: super().__init__() @@ -339,8 +339,8 @@ class BgrToRgba(Module): >>> output = rgba(input) # 2x4x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 4, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 4, -1, -1] def __init__(self, alpha_val: Union[float, Tensor]) -> None: super().__init__() @@ -368,8 +368,8 @@ class RgbaToRgb(Module): >>> output = rgba(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 4, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 4, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return rgba_to_rgb(image) @@ -393,8 +393,8 @@ class RgbaToBgr(Module): >>> output = rgba(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 4, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 4, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return rgba_to_bgr(image) @@ -426,8 +426,8 @@ class RgbToLinearRgb(Module): [3] https://en.wikipedia.org/wiki/SRGB """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return rgb_to_linear_rgb(image) @@ -458,8 +458,8 @@ class LinearRgbToRgb(Module): [3] https://en.wikipedia.org/wiki/SRGB """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return linear_rgb_to_rgb(image) diff --git a/kornia/color/xyz.py b/kornia/color/xyz.py index c5493840db..cc6a06c656 100644 --- a/kornia/color/xyz.py +++ b/kornia/color/xyz.py @@ -91,8 +91,8 @@ class RgbToXyz(Module): [1] https://docs.opencv.org/4.0.1/de/d25/imgproc_color_conversions.html """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return rgb_to_xyz(image) @@ -117,8 +117,8 @@ class XyzToRgb(Module): [1] https://docs.opencv.org/4.0.1/de/d25/imgproc_color_conversions.html """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return xyz_to_rgb(image) diff --git a/kornia/color/ycbcr.py b/kornia/color/ycbcr.py index aac42300eb..1351dcfd8d 100644 --- a/kornia/color/ycbcr.py +++ b/kornia/color/ycbcr.py @@ -121,8 +121,8 @@ class RgbToYcbcr(Module): >>> output = ycbcr(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return rgb_to_ycbcr(image) @@ -146,8 +146,8 @@ class YcbcrToRgb(Module): >>> output = rgb(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return ycbcr_to_rgb(image) diff --git a/kornia/color/yuv.py b/kornia/color/yuv.py index 3e94b6615e..ab89ee2982 100644 --- a/kornia/color/yuv.py +++ b/kornia/color/yuv.py @@ -293,8 +293,8 @@ class RgbToYuv(Module): [1] https://es.wikipedia.org/wiki/YUV#RGB_a_Y'UV """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def forward(self, input: Tensor) -> Tensor: return rgb_to_yuv(input) @@ -391,8 +391,8 @@ class YuvToRgb(Module): >>> output = rgb(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def forward(self, input: Tensor) -> Tensor: return yuv_to_rgb(input) diff --git a/kornia/core/module.py b/kornia/core/module.py index 00e315561e..e222ccb4c1 100644 --- a/kornia/core/module.py +++ b/kornia/core/module.py @@ -2,7 +2,7 @@ import math import os from functools import wraps -from typing import Any, Callable, List, Optional, Tuple, Union +from typing import Any, Callable, ClassVar, List, Optional, Tuple, Union import torch @@ -15,8 +15,8 @@ class ONNXExportMixin: ONNX_EXPORTABLE: bool = True - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, -1, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, -1, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, -1, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, -1, -1, -1] def to_onnx( self, diff --git a/kornia/enhance/adjust.py b/kornia/enhance/adjust.py index 41c2ced013..ebac19a9ac 100644 --- a/kornia/enhance/adjust.py +++ b/kornia/enhance/adjust.py @@ -1042,8 +1042,8 @@ class AdjustSaturation(Module): tensor(0.) """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def __init__(self, saturation_factor: Union[float, Tensor]) -> None: super().__init__() @@ -1092,8 +1092,8 @@ class AdjustSaturationWithGraySubtraction(Module): tensor(0.) """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def __init__(self, saturation_factor: Union[float, Tensor]) -> None: super().__init__() @@ -1142,8 +1142,8 @@ class AdjustHue(Module): torch.Size([2, 3, 3, 3]) """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] def __init__(self, hue_factor: Union[float, Tensor]) -> None: super().__init__() diff --git a/requirements/requirements-dev.txt b/requirements/requirements-dev.txt index bd3222ebc9..4fc3699bb7 100644 --- a/requirements/requirements-dev.txt +++ b/requirements/requirements-dev.txt @@ -11,3 +11,5 @@ pytest-timeout requests transformers types-requests +onnx +onnxruntime From c059039ec49565f61c6b2ca0eb181d7d497a1a78 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 7 Sep 2024 16:12:25 +0000 Subject: [PATCH 15/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- requirements/requirements-dev.txt | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/requirements/requirements-dev.txt b/requirements/requirements-dev.txt index 4fc3699bb7..21866c5bfb 100644 --- a/requirements/requirements-dev.txt +++ b/requirements/requirements-dev.txt @@ -4,6 +4,7 @@ diffusers mypy numpy<3 onnx +onnxruntime pillow pre-commit>=2 pytest==8.3.2 @@ -11,5 +12,3 @@ pytest-timeout requests transformers types-requests -onnx -onnxruntime From fa4c98375ba54ddecac2bf0eb26676cf5492b4e6 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Sat, 7 Sep 2024 19:17:19 +0300 Subject: [PATCH 16/84] update --- kornia/color/gray.py | 2 +- kornia/color/hls.py | 2 +- kornia/color/hsv.py | 1 + kornia/color/lab.py | 2 +- kornia/color/luv.py | 2 +- kornia/color/raw.py | 1 + kornia/color/rgb.py | 2 +- kornia/color/xyz.py | 1 + kornia/color/ycbcr.py | 1 + kornia/color/yuv.py | 2 +- kornia/enhance/adjust.py | 2 +- kornia/geometry/transform/affwarp.py | 2 +- 12 files changed, 12 insertions(+), 8 deletions(-) diff --git a/kornia/color/gray.py b/kornia/color/gray.py index 45f4c958f0..822ca3acb5 100644 --- a/kornia/color/gray.py +++ b/kornia/color/gray.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Optional +from typing import Optional, ClassVar import torch diff --git a/kornia/color/hls.py b/kornia/color/hls.py index 7880999e53..41de37d7b5 100644 --- a/kornia/color/hls.py +++ b/kornia/color/hls.py @@ -1,5 +1,5 @@ import math -from typing import Tuple +from typing import ClassVar, Tuple import torch diff --git a/kornia/color/hsv.py b/kornia/color/hsv.py index 37dfe4baf0..5164e5ccf6 100644 --- a/kornia/color/hsv.py +++ b/kornia/color/hsv.py @@ -1,4 +1,5 @@ import math +from typing import ClassVar import torch diff --git a/kornia/color/lab.py b/kornia/color/lab.py index 13b25b3be6..609239fc2f 100644 --- a/kornia/color/lab.py +++ b/kornia/color/lab.py @@ -2,7 +2,7 @@ https://github.com/scikit-image/scikit-image/blob/a48bf6774718c64dade4548153ae16065b595ca9/skimage/color/colorconv.py """ - +from typing import ClassVar import torch from kornia.core import ImageModule as Module diff --git a/kornia/color/luv.py b/kornia/color/luv.py index 0b336773e9..8fba7bc06a 100644 --- a/kornia/color/luv.py +++ b/kornia/color/luv.py @@ -3,7 +3,7 @@ https://github.com/scikit-image/scikit-image/blob/a48bf6774718c64dade4548153ae16065b595ca9/skimage/color/colorconv.py """ -from typing import Tuple +from typing import ClassVar, Tuple import torch diff --git a/kornia/color/raw.py b/kornia/color/raw.py index 4e040ff052..c36938c43a 100644 --- a/kornia/color/raw.py +++ b/kornia/color/raw.py @@ -1,3 +1,4 @@ +from typing import ClassVar from enum import Enum import torch diff --git a/kornia/color/rgb.py b/kornia/color/rgb.py index 92f4310e8b..aebdca104c 100644 --- a/kornia/color/rgb.py +++ b/kornia/color/rgb.py @@ -1,4 +1,4 @@ -from typing import Union, cast +from typing import ClassVar, Union, cast import torch diff --git a/kornia/color/xyz.py b/kornia/color/xyz.py index cc6a06c656..c4bbcf6dd8 100644 --- a/kornia/color/xyz.py +++ b/kornia/color/xyz.py @@ -1,3 +1,4 @@ +from typing import ClassVar import torch from kornia.core import ImageModule as Module diff --git a/kornia/color/ycbcr.py b/kornia/color/ycbcr.py index 1351dcfd8d..ba73f2ab24 100644 --- a/kornia/color/ycbcr.py +++ b/kornia/color/ycbcr.py @@ -1,3 +1,4 @@ +from typing import ClassVar import torch from kornia.core import ImageModule as Module diff --git a/kornia/color/yuv.py b/kornia/color/yuv.py index ab89ee2982..5255321d82 100644 --- a/kornia/color/yuv.py +++ b/kornia/color/yuv.py @@ -1,4 +1,4 @@ -from typing import Tuple +from typing import ClassVar, Tuple import torch diff --git a/kornia/enhance/adjust.py b/kornia/enhance/adjust.py index ebac19a9ac..af179e9144 100644 --- a/kornia/enhance/adjust.py +++ b/kornia/enhance/adjust.py @@ -1,5 +1,5 @@ from math import pi -from typing import Optional, Union +from typing import Optional, ClassVar, Union import torch diff --git a/kornia/geometry/transform/affwarp.py b/kornia/geometry/transform/affwarp.py index 3abb418ec4..ccc18c0ec6 100644 --- a/kornia/geometry/transform/affwarp.py +++ b/kornia/geometry/transform/affwarp.py @@ -1,5 +1,5 @@ import warnings -from typing import Optional, Tuple, Union +from typing import ClassVar, Optional, Tuple, Union import torch From 95f065d871d339e2cee17597e287bf56a5a44cda Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 7 Sep 2024 16:18:04 +0000 Subject: [PATCH 17/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/color/gray.py | 2 +- kornia/color/lab.py | 2 ++ kornia/color/raw.py | 2 +- kornia/color/xyz.py | 1 + kornia/color/ycbcr.py | 1 + kornia/enhance/adjust.py | 2 +- 6 files changed, 7 insertions(+), 3 deletions(-) diff --git a/kornia/color/gray.py b/kornia/color/gray.py index 822ca3acb5..f647cfd337 100644 --- a/kornia/color/gray.py +++ b/kornia/color/gray.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Optional, ClassVar +from typing import ClassVar, Optional import torch diff --git a/kornia/color/lab.py b/kornia/color/lab.py index 609239fc2f..6062f57e66 100644 --- a/kornia/color/lab.py +++ b/kornia/color/lab.py @@ -2,7 +2,9 @@ https://github.com/scikit-image/scikit-image/blob/a48bf6774718c64dade4548153ae16065b595ca9/skimage/color/colorconv.py """ + from typing import ClassVar + import torch from kornia.core import ImageModule as Module diff --git a/kornia/color/raw.py b/kornia/color/raw.py index c36938c43a..fb6a7fea7a 100644 --- a/kornia/color/raw.py +++ b/kornia/color/raw.py @@ -1,5 +1,5 @@ -from typing import ClassVar from enum import Enum +from typing import ClassVar import torch diff --git a/kornia/color/xyz.py b/kornia/color/xyz.py index c4bbcf6dd8..7c5e89590e 100644 --- a/kornia/color/xyz.py +++ b/kornia/color/xyz.py @@ -1,4 +1,5 @@ from typing import ClassVar + import torch from kornia.core import ImageModule as Module diff --git a/kornia/color/ycbcr.py b/kornia/color/ycbcr.py index ba73f2ab24..a43394b139 100644 --- a/kornia/color/ycbcr.py +++ b/kornia/color/ycbcr.py @@ -1,4 +1,5 @@ from typing import ClassVar + import torch from kornia.core import ImageModule as Module diff --git a/kornia/enhance/adjust.py b/kornia/enhance/adjust.py index af179e9144..f625b48e36 100644 --- a/kornia/enhance/adjust.py +++ b/kornia/enhance/adjust.py @@ -1,5 +1,5 @@ from math import pi -from typing import Optional, ClassVar, Union +from typing import ClassVar, Optional, Union import torch From 31452821f0857738fed03aab47eeff339eb9ec61 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Sat, 7 Sep 2024 19:26:51 +0300 Subject: [PATCH 18/84] update --- kornia/filters/motion.py | 6 ++++-- kornia/filters/sobel.py | 10 ++++++---- kornia/onnx/sequential.py | 18 +++++++++--------- kornia/onnx/utils.py | 7 +++++-- 4 files changed, 24 insertions(+), 17 deletions(-) diff --git a/kornia/filters/motion.py b/kornia/filters/motion.py index d6eb4d1781..7442f1167f 100644 --- a/kornia/filters/motion.py +++ b/kornia/filters/motion.py @@ -1,5 +1,7 @@ from __future__ import annotations +from typing import ClassVar + from kornia.core import ImageModule as Module from kornia.core import Tensor from kornia.core.check import KORNIA_CHECK @@ -84,8 +86,8 @@ class MotionBlur3D(Module): >>> output = motion_blur(input) # 2x4x5x7x9 """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, -1, -1, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, -1, -1, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, -1, -1, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, -1, -1, -1, -1] def __init__( self, diff --git a/kornia/filters/sobel.py b/kornia/filters/sobel.py index 897cb3d9a6..6b7a048795 100644 --- a/kornia/filters/sobel.py +++ b/kornia/filters/sobel.py @@ -1,5 +1,7 @@ from __future__ import annotations +from typing import ClassVar + import torch import torch.nn.functional as F @@ -172,8 +174,8 @@ class SpatialGradient(Module): >>> output = SpatialGradient()(input) # 1x3x2x4x4 """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, -1, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, -1, 2, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, -1, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, -1, 2, -1, -1] def __init__(self, mode: str = "sobel", order: int = 1, normalized: bool = True) -> None: super().__init__() @@ -209,8 +211,8 @@ class SpatialGradient3d(Module): torch.Size([1, 4, 3, 2, 4, 4]) """ - ONNX_DEFAULT_INPUTSHAPE: list[int] = [-1, -1, -1, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: list[int] = [-1, -1, -1, -1, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, -1, -1, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, -1, -1, -1, -1, -1] def __init__(self, mode: str = "diff", order: int = 1) -> None: super().__init__() diff --git a/kornia/onnx/sequential.py b/kornia/onnx/sequential.py index 98b9573e8a..a630668c1c 100644 --- a/kornia/onnx/sequential.py +++ b/kornia/onnx/sequential.py @@ -39,9 +39,9 @@ class ONNXSequential: def __init__( self, - *args: Union[onnx.ModelProto, str], # type:ignore + *args: Union["onnx.ModelProto", str], # type:ignore providers: Optional[list[str]] = None, - session_options: Optional[ort.SessionOptions] = None, # type:ignore + session_options: Optional["ort.SessionOptions"] = None, # type:ignore io_maps: Optional[list[tuple[str, str]]] = None, cache_dir: Optional[str] = None, ) -> None: @@ -50,7 +50,7 @@ def __init__( self._combined_op = self._combine(io_maps) self._session = self.create_session() - def _load_op(self, arg: Union[onnx.ModelProto, str]) -> onnx.ModelProto: # type:ignore + def _load_op(self, arg: Union["onnx.ModelProto", str]) -> "onnx.ModelProto": # type:ignore """Loads an ONNX model, either from a file path or use the provided ONNX ModelProto. Args: @@ -63,7 +63,7 @@ def _load_op(self, arg: Union[onnx.ModelProto, str]) -> onnx.ModelProto: # type return self.onnx_loader.load_model(arg) return arg - def _combine(self, io_maps: Optional[list[tuple[str, str]]] = None) -> onnx.ModelProto: # type:ignore + def _combine(self, io_maps: Optional[list[tuple[str, str]]] = None) -> "onnx.ModelProto": # type:ignore """Combine the provided ONNX models into a single ONNX graph. Optionally, map inputs and outputs between operators using the `io_map`. @@ -106,8 +106,8 @@ def export(self, file_path: str) -> None: def create_session( self, providers: Optional[list[str]] = None, - session_options: Optional[ort.SessionOptions] = None, # type:ignore - ) -> ort.InferenceSession: # type:ignore + session_options: Optional["ort.SessionOptions"] = None, # type:ignore + ) -> "ort.InferenceSession": # type:ignore """Create an optimized ONNXRuntime InferenceSession for the combined model. Args: @@ -129,7 +129,7 @@ def create_session( ) return session - def set_session(self, session: ort.InferenceSession) -> None: # type: ignore + def set_session(self, session: "ort.InferenceSession") -> None: # type: ignore """Set a custom ONNXRuntime InferenceSession. Args: @@ -138,7 +138,7 @@ def set_session(self, session: ort.InferenceSession) -> None: # type: ignore """ self._session = session - def get_session(self) -> ort.InferenceSession: # type: ignore + def get_session(self) -> "ort.InferenceSession": # type: ignore """Get the current ONNXRuntime InferenceSession. Returns: @@ -146,7 +146,7 @@ def get_session(self) -> ort.InferenceSession: # type: ignore """ return self._session - def __call__(self, *inputs: np.ndarray) -> list[np.ndarray]: # type:ignore + def __call__(self, *inputs: "np.ndarray") -> list["np.ndarray"]: # type:ignore """Perform inference using the combined ONNX model. Args: diff --git a/kornia/onnx/utils.py b/kornia/onnx/utils.py index 8213de73f8..14d14ea498 100644 --- a/kornia/onnx/utils.py +++ b/kornia/onnx/utils.py @@ -50,7 +50,7 @@ def _get_file_path(self, model_name: str, cache_dir: Optional[str]) -> str: file_path = os.path.join(cache_dir, "/".join(model_name.split("/")[:-1]), file_name) return file_path - def load_model(self, model_name: str, download: bool = False, **kwargs) -> onnx.ModelProto: # type:ignore + def load_model(self, model_name: str, download: bool = False, **kwargs) -> "onnx.ModelProto": # type:ignore """Loads an ONNX model from the local cache or downloads it from Hugging Face if necessary. Args: @@ -108,7 +108,10 @@ def _fetch_repo_contents(folder: str) -> list[dict[str, Any]]: List[dict]: A list of all files in the repository as dictionaries containing file details. """ url = f"https://huggingface.co/api/models/kornia/ONNX_models/tree/main/{folder}" - response = requests.get(url) + + if not url.startswith(("http:", "https:")): + raise ValueError("URL must start with 'http:' or 'https:'") + response = requests.get(url, timeout=10) if response.status_code == 200: return response.json() # Returns the JSON content of the repo From 84ff540450649e99fa45f98f0898ab21651e1bf3 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Sat, 7 Sep 2024 19:28:33 +0300 Subject: [PATCH 19/84] update --- kornia/onnx/utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/kornia/onnx/utils.py b/kornia/onnx/utils.py index 14d14ea498..10d83bb120 100644 --- a/kornia/onnx/utils.py +++ b/kornia/onnx/utils.py @@ -94,6 +94,9 @@ def download( os.makedirs(os.path.dirname(file_path), exist_ok=True) # Create the cache directory if it doesn't exist + if not url.startswith(("http:", "https:")): + raise ValueError("URL must start with 'http:' or 'https:'") + # Download the file and save it try: urllib.request.urlretrieve(url, file_path) @@ -109,8 +112,6 @@ def _fetch_repo_contents(folder: str) -> list[dict[str, Any]]: """ url = f"https://huggingface.co/api/models/kornia/ONNX_models/tree/main/{folder}" - if not url.startswith(("http:", "https:")): - raise ValueError("URL must start with 'http:' or 'https:'") response = requests.get(url, timeout=10) if response.status_code == 200: From 01a261355f75cba884271df2bbd87e3b0da479fb Mon Sep 17 00:00:00 2001 From: Jian S Date: Mon, 9 Sep 2024 02:20:07 +0800 Subject: [PATCH 20/84] update --- kornia/core/module.py | 48 +++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 46 insertions(+), 2 deletions(-) diff --git a/kornia/core/module.py b/kornia/core/module.py index e222ccb4c1..a434441430 100644 --- a/kornia/core/module.py +++ b/kornia/core/module.py @@ -11,9 +11,26 @@ from ._backend import Module, Tensor, from_numpy from .external import PILImage as Image from .external import numpy as np +from .external import onnx class ONNXExportMixin: + """Mixin class that provides ONNX export functionality for objects that support it. + + Attributes: + ONNX_EXPORTABLE: + A flag indicating whether the object can be exported to ONNX. Default is True. + ONNX_DEFAULT_INPUTSHAPE: + Default input shape for the ONNX export. A list of integers where `-1` indicates + dynamic dimensions. Default is [-1, -1, -1, -1]. + ONNX_DEFAULT_OUTPUTSHAP: + Default output shape for the ONNX export. A list of integers where `-1` indicates + dynamic dimensions. Default is [-1, -1, -1, -1]. + + Note: + - If `ONNX_EXPORTABLE` is False, indicating that the object cannot be exported to ONNX. + """ + ONNX_EXPORTABLE: bool = True ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, -1, -1, -1] ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, -1, -1, -1] @@ -24,6 +41,24 @@ def to_onnx( input_shape: Optional[list[int]] = None, output_shape: Optional[list[int]] = None, ) -> None: + """Exports the current object to an ONNX model file. + + Args: + onnx_name (Optional[str]): + The name of the output ONNX file. If not provided, a default name in the + format "Kornia-.onnx" will be used. + input_shape (Optional[list[int]]): + The input shape for the model as a list of integers. If None, + `ONNX_DEFAULT_INPUTSHAPE` will be used. Dynamic dimensions can be indicated by `-1`. + output_shape (Optional[list[int]]): + The output shape for the model as a list of integers. If None, + `ONNX_DEFAULT_OUTPUTSHAPE` will be used. Dynamic dimensions can be indicated by `-1`. + + Notes: + - A dummy input tensor is created based on the provided or default input shape. + - Dynamic axes for input and output tensors are configured where dimensions are marked `-1`. + - The model is exported with `torch.onnx.export`, with constant folding enabled and opset version set to 17. + """ if not self.ONNX_EXPORTABLE: raise RuntimeError("This object cannot be exported to ONNX.") @@ -36,8 +71,8 @@ def to_onnx( onnx_name = f"Kornia-{self.__class__.__name__}.onnx" # Creating a dummy input with the given shape - psuedo_shape = (1, 3, 256, 256) - dummy_input = torch.randn(*[(psuedo_shape[i] if dim == -1 else dim) for i, dim in enumerate(input_shape)]) + pseudo_shape = (1, 3, 256, 256) + dummy_input = torch.randn(*[(pseudo_shape[i] if dim == -1 else dim) for i, dim in enumerate(input_shape)]) # Dynamic axis configuration for input and output dynamic_axes = { @@ -56,6 +91,15 @@ def to_onnx( output_names=["output"], dynamic_axes=dynamic_axes, ) + + onnx_model = onnx.load(onnx_name) # type: ignore + + for key, value in [("source", "kornia"), ("version", kornia.__version__), ("class", self.__class__.__name__)]: + metadata_props = onnx_model.metadata_props.add() # type: ignore + metadata_props.key = key # type: ignore + metadata_props.value = str(value) # type: ignore + + onnx.save(onnx_model, onnx_name) class ImageModuleMixIn: From 92ea7cfe09290d71a7ee632f10dbbb8a58265a44 Mon Sep 17 00:00:00 2001 From: Jian S Date: Mon, 9 Sep 2024 06:19:26 +0800 Subject: [PATCH 21/84] update --- docs/source/onnx.rst | 6 +++--- kornia/core/module.py | 12 ++++++------ 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/docs/source/onnx.rst b/docs/source/onnx.rst index 093bd694d1..c08e70e24c 100644 --- a/docs/source/onnx.rst +++ b/docs/source/onnx.rst @@ -1,5 +1,5 @@ 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. @@ -91,7 +91,7 @@ Here's how you can quickly get started with `ONNXSequential`: Frequently Asked Questions (FAQ) -------------------------------- +-------------------------------- **1. Can I chain models from different sources?** @@ -106,7 +106,7 @@ You can use the `io_map` parameter to control how outputs of one model are mappe 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). diff --git a/kornia/core/module.py b/kornia/core/module.py index a434441430..9fdc3587de 100644 --- a/kornia/core/module.py +++ b/kornia/core/module.py @@ -32,8 +32,8 @@ class ONNXExportMixin: """ ONNX_EXPORTABLE: bool = True - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, -1, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, -1, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, -1, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, -1, -1, -1] def to_onnx( self, @@ -95,11 +95,11 @@ def to_onnx( onnx_model = onnx.load(onnx_name) # type: ignore for key, value in [("source", "kornia"), ("version", kornia.__version__), ("class", self.__class__.__name__)]: - metadata_props = onnx_model.metadata_props.add() # type: ignore - metadata_props.key = key # type: ignore - metadata_props.value = str(value) # type: ignore + metadata_props = onnx_model.metadata_props.add() + metadata_props.key = key + metadata_props.value = str(value) - onnx.save(onnx_model, onnx_name) + onnx.save(onnx_model, onnx_name) # type: ignore class ImageModuleMixIn: From b71fa4d59890454704065d146d545a72a390eb32 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 8 Sep 2024 18:20:23 +0000 Subject: [PATCH 22/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/core/module.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/kornia/core/module.py b/kornia/core/module.py index 9fdc3587de..b2a27a91f1 100644 --- a/kornia/core/module.py +++ b/kornia/core/module.py @@ -18,13 +18,13 @@ class ONNXExportMixin: """Mixin class that provides ONNX export functionality for objects that support it. Attributes: - ONNX_EXPORTABLE: + ONNX_EXPORTABLE: A flag indicating whether the object can be exported to ONNX. Default is True. - ONNX_DEFAULT_INPUTSHAPE: - Default input shape for the ONNX export. A list of integers where `-1` indicates + ONNX_DEFAULT_INPUTSHAPE: + Default input shape for the ONNX export. A list of integers where `-1` indicates dynamic dimensions. Default is [-1, -1, -1, -1]. - ONNX_DEFAULT_OUTPUTSHAP: - Default output shape for the ONNX export. A list of integers where `-1` indicates + ONNX_DEFAULT_OUTPUTSHAP: + Default output shape for the ONNX export. A list of integers where `-1` indicates dynamic dimensions. Default is [-1, -1, -1, -1]. Note: @@ -44,14 +44,14 @@ def to_onnx( """Exports the current object to an ONNX model file. Args: - onnx_name (Optional[str]): - The name of the output ONNX file. If not provided, a default name in the + onnx_name (Optional[str]): + The name of the output ONNX file. If not provided, a default name in the format "Kornia-.onnx" will be used. - input_shape (Optional[list[int]]): - The input shape for the model as a list of integers. If None, + input_shape (Optional[list[int]]): + The input shape for the model as a list of integers. If None, `ONNX_DEFAULT_INPUTSHAPE` will be used. Dynamic dimensions can be indicated by `-1`. - output_shape (Optional[list[int]]): - The output shape for the model as a list of integers. If None, + output_shape (Optional[list[int]]): + The output shape for the model as a list of integers. If None, `ONNX_DEFAULT_OUTPUTSHAPE` will be used. Dynamic dimensions can be indicated by `-1`. Notes: @@ -91,7 +91,7 @@ def to_onnx( output_names=["output"], dynamic_axes=dynamic_axes, ) - + onnx_model = onnx.load(onnx_name) # type: ignore for key, value in [("source", "kornia"), ("version", kornia.__version__), ("class", self.__class__.__name__)]: From 7ceaba678e1fa157f359a033d9388fc3234e9e29 Mon Sep 17 00:00:00 2001 From: Jian S Date: Mon, 9 Sep 2024 06:30:57 +0800 Subject: [PATCH 23/84] update --- kornia/color/gray.py | 14 +++++++------- kornia/color/hls.py | 10 +++++----- kornia/color/hsv.py | 10 +++++----- kornia/color/lab.py | 10 +++++----- kornia/color/luv.py | 10 +++++----- kornia/color/raw.py | 10 +++++----- kornia/color/rgb.py | 34 +++++++++++++++++----------------- kornia/color/xyz.py | 10 +++++----- kornia/color/ycbcr.py | 10 +++++----- kornia/color/yuv.py | 10 +++++----- kornia/core/module.py | 10 +++++----- kornia/enhance/adjust.py | 14 +++++++------- kornia/onnx/utils.py | 10 +++++----- 13 files changed, 81 insertions(+), 81 deletions(-) diff --git a/kornia/color/gray.py b/kornia/color/gray.py index f647cfd337..714673d862 100644 --- a/kornia/color/gray.py +++ b/kornia/color/gray.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import ClassVar, Optional +from typing import ClassVar, List, Optional import torch @@ -125,8 +125,8 @@ 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] + 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) @@ -150,8 +150,8 @@ 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] + 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__() @@ -181,8 +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] + 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) diff --git a/kornia/color/hls.py b/kornia/color/hls.py index 41de37d7b5..0506762c8c 100644 --- a/kornia/color/hls.py +++ b/kornia/color/hls.py @@ -1,5 +1,5 @@ import math -from typing import ClassVar, Tuple +from typing import ClassVar, List, Tuple import torch @@ -142,8 +142,8 @@ class RgbToHls(Module): >>> output = hls(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return rgb_to_hls(image) @@ -170,8 +170,8 @@ class HlsToRgb(Module): >>> output = rgb(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return hls_to_rgb(image) diff --git a/kornia/color/hsv.py b/kornia/color/hsv.py index 5164e5ccf6..4b82447306 100644 --- a/kornia/color/hsv.py +++ b/kornia/color/hsv.py @@ -1,5 +1,5 @@ import math -from typing import ClassVar +from typing import ClassVar, List import torch @@ -117,8 +117,8 @@ class RgbToHsv(Module): >>> output = hsv(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def __init__(self, eps: float = 1e-6) -> None: super().__init__() @@ -146,8 +146,8 @@ class HsvToRgb(Module): >>> output = rgb(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def forward(self, image: torch.Tensor) -> torch.Tensor: return hsv_to_rgb(image) diff --git a/kornia/color/lab.py b/kornia/color/lab.py index 6062f57e66..3dedc03e2c 100644 --- a/kornia/color/lab.py +++ b/kornia/color/lab.py @@ -3,7 +3,7 @@ https://github.com/scikit-image/scikit-image/blob/a48bf6774718c64dade4548153ae16065b595ca9/skimage/color/colorconv.py """ -from typing import ClassVar +from typing import ClassVar, List import torch @@ -152,8 +152,8 @@ class RgbToLab(Module): [3] https://github.com/torch/image/blob/dc061b98fb7e946e00034a5fc73e883a299edc7f/generic/image.c#L1467 """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def forward(self, image: torch.Tensor) -> torch.Tensor: return rgb_to_lab(image) @@ -182,8 +182,8 @@ class LabToRgb(Module): [3] https://github.com/torch/image/blob/dc061b98fb7e946e00034a5fc73e883a299edc7f/generic/image.c#L1518 """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def forward(self, image: torch.Tensor, clip: bool = True) -> torch.Tensor: return lab_to_rgb(image, clip) diff --git a/kornia/color/luv.py b/kornia/color/luv.py index 8fba7bc06a..34c00bb746 100644 --- a/kornia/color/luv.py +++ b/kornia/color/luv.py @@ -3,7 +3,7 @@ https://github.com/scikit-image/scikit-image/blob/a48bf6774718c64dade4548153ae16065b595ca9/skimage/color/colorconv.py """ -from typing import ClassVar, Tuple +from typing import ClassVar, List, Tuple import torch @@ -141,8 +141,8 @@ class RgbToLuv(Module): [3] http://www.poynton.com/ColorFAQ.html """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def forward(self, image: torch.Tensor) -> torch.Tensor: return rgb_to_luv(image) @@ -171,8 +171,8 @@ class LuvToRgb(Module): [3] http://www.poynton.com/ColorFAQ.html """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def forward(self, image: torch.Tensor) -> torch.Tensor: return luv_to_rgb(image) diff --git a/kornia/color/raw.py b/kornia/color/raw.py index fb6a7fea7a..b13459fb3e 100644 --- a/kornia/color/raw.py +++ b/kornia/color/raw.py @@ -1,5 +1,5 @@ from enum import Enum -from typing import ClassVar +from typing import ClassVar, List import torch @@ -289,8 +289,8 @@ class RawToRgb(Module): >>> output = rgb(rawinput) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 1, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 1, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def __init__(self, cfa: CFA) -> None: super().__init__() @@ -318,8 +318,8 @@ class RgbToRaw(Module): >>> output = raw(rgbinput) # 2x1x4x6 """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 1, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 1, -1, -1] def __init__(self, cfa: CFA) -> None: super().__init__() diff --git a/kornia/color/rgb.py b/kornia/color/rgb.py index aebdca104c..068b19f5eb 100644 --- a/kornia/color/rgb.py +++ b/kornia/color/rgb.py @@ -1,4 +1,4 @@ -from typing import ClassVar, Union, cast +from typing import ClassVar, List, Union, cast import torch @@ -248,8 +248,8 @@ class BgrToRgb(Module): >>> output = rgb(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return bgr_to_rgb(image) @@ -273,8 +273,8 @@ class RgbToBgr(Module): >>> output = bgr(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return rgb_to_bgr(image) @@ -304,8 +304,8 @@ class RgbToRgba(Module): >>> output = rgba(input) # 2x4x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 4, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 4, -1, -1] def __init__(self, alpha_val: Union[float, Tensor]) -> None: super().__init__() @@ -339,8 +339,8 @@ class BgrToRgba(Module): >>> output = rgba(input) # 2x4x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 4, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 4, -1, -1] def __init__(self, alpha_val: Union[float, Tensor]) -> None: super().__init__() @@ -368,8 +368,8 @@ class RgbaToRgb(Module): >>> output = rgba(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 4, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 4, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return rgba_to_rgb(image) @@ -393,8 +393,8 @@ class RgbaToBgr(Module): >>> output = rgba(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 4, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 4, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return rgba_to_bgr(image) @@ -426,8 +426,8 @@ class RgbToLinearRgb(Module): [3] https://en.wikipedia.org/wiki/SRGB """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return rgb_to_linear_rgb(image) @@ -458,8 +458,8 @@ class LinearRgbToRgb(Module): [3] https://en.wikipedia.org/wiki/SRGB """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return linear_rgb_to_rgb(image) diff --git a/kornia/color/xyz.py b/kornia/color/xyz.py index 7c5e89590e..456f1e8e55 100644 --- a/kornia/color/xyz.py +++ b/kornia/color/xyz.py @@ -1,4 +1,4 @@ -from typing import ClassVar +from typing import ClassVar, List import torch @@ -93,8 +93,8 @@ class RgbToXyz(Module): [1] https://docs.opencv.org/4.0.1/de/d25/imgproc_color_conversions.html """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return rgb_to_xyz(image) @@ -119,8 +119,8 @@ class XyzToRgb(Module): [1] https://docs.opencv.org/4.0.1/de/d25/imgproc_color_conversions.html """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return xyz_to_rgb(image) diff --git a/kornia/color/ycbcr.py b/kornia/color/ycbcr.py index a43394b139..014d1fcaba 100644 --- a/kornia/color/ycbcr.py +++ b/kornia/color/ycbcr.py @@ -1,4 +1,4 @@ -from typing import ClassVar +from typing import ClassVar, List import torch @@ -123,8 +123,8 @@ class RgbToYcbcr(Module): >>> output = ycbcr(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return rgb_to_ycbcr(image) @@ -148,8 +148,8 @@ class YcbcrToRgb(Module): >>> output = rgb(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def forward(self, image: Tensor) -> Tensor: return ycbcr_to_rgb(image) diff --git a/kornia/color/yuv.py b/kornia/color/yuv.py index 5255321d82..88f6f835fe 100644 --- a/kornia/color/yuv.py +++ b/kornia/color/yuv.py @@ -1,4 +1,4 @@ -from typing import ClassVar, Tuple +from typing import ClassVar, List, Tuple import torch @@ -293,8 +293,8 @@ class RgbToYuv(Module): [1] https://es.wikipedia.org/wiki/YUV#RGB_a_Y'UV """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def forward(self, input: Tensor) -> Tensor: return rgb_to_yuv(input) @@ -391,8 +391,8 @@ class YuvToRgb(Module): >>> output = rgb(input) # 2x3x4x5 """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def forward(self, input: Tensor) -> Tensor: return yuv_to_rgb(input) diff --git a/kornia/core/module.py b/kornia/core/module.py index b2a27a91f1..2c6c5cc10d 100644 --- a/kornia/core/module.py +++ b/kornia/core/module.py @@ -38,19 +38,19 @@ class ONNXExportMixin: def to_onnx( self, onnx_name: Optional[str] = None, - input_shape: Optional[list[int]] = None, - output_shape: Optional[list[int]] = None, + input_shape: Optional[List[int]] = None, + output_shape: Optional[List[int]] = None, ) -> None: """Exports the current object to an ONNX model file. Args: - onnx_name (Optional[str]): + onnx_name: The name of the output ONNX file. If not provided, a default name in the format "Kornia-.onnx" will be used. - input_shape (Optional[list[int]]): + input_shape: The input shape for the model as a list of integers. If None, `ONNX_DEFAULT_INPUTSHAPE` will be used. Dynamic dimensions can be indicated by `-1`. - output_shape (Optional[list[int]]): + output_shape: The output shape for the model as a list of integers. If None, `ONNX_DEFAULT_OUTPUTSHAPE` will be used. Dynamic dimensions can be indicated by `-1`. diff --git a/kornia/enhance/adjust.py b/kornia/enhance/adjust.py index f625b48e36..18e79b6eb1 100644 --- a/kornia/enhance/adjust.py +++ b/kornia/enhance/adjust.py @@ -1,5 +1,5 @@ from math import pi -from typing import ClassVar, Optional, Union +from typing import ClassVar, List, Optional, Union import torch @@ -1042,8 +1042,8 @@ class AdjustSaturation(Module): tensor(0.) """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def __init__(self, saturation_factor: Union[float, Tensor]) -> None: super().__init__() @@ -1092,8 +1092,8 @@ class AdjustSaturationWithGraySubtraction(Module): tensor(0.) """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def __init__(self, saturation_factor: Union[float, Tensor]) -> None: super().__init__() @@ -1142,8 +1142,8 @@ class AdjustHue(Module): torch.Size([2, 3, 3, 3]) """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] def __init__(self, hue_factor: Union[float, Tensor]) -> None: super().__init__() diff --git a/kornia/onnx/utils.py b/kornia/onnx/utils.py index 10d83bb120..264b1ad1cf 100644 --- a/kornia/onnx/utils.py +++ b/kornia/onnx/utils.py @@ -95,13 +95,13 @@ def download( os.makedirs(os.path.dirname(file_path), exist_ok=True) # Create the cache directory if it doesn't exist if not url.startswith(("http:", "https:")): + try: + urllib.request.urlretrieve(url, file_path) + except urllib.error.HTTPError as e: + raise ValueError(f"Error in resolving `{url}`. {e}.") + else: raise ValueError("URL must start with 'http:' or 'https:'") - # Download the file and save it - try: - urllib.request.urlretrieve(url, file_path) - except urllib.error.HTTPError as e: - raise ValueError(f"Error in resolving `{url}`. {e}.") @staticmethod def _fetch_repo_contents(folder: str) -> list[dict[str, Any]]: From a552c1d1d5c7eb33d95650d5cf3bc650c507d3fe Mon Sep 17 00:00:00 2001 From: Jian S Date: Mon, 9 Sep 2024 16:17:11 +0800 Subject: [PATCH 24/84] update --- kornia/constants.py | 3 +++ kornia/onnx/utils.py | 23 +++++++++++++++-------- 2 files changed, 18 insertions(+), 8 deletions(-) diff --git a/kornia/constants.py b/kornia/constants.py index 72ac09863f..19fd9de6dd 100644 --- a/kornia/constants.py +++ b/kornia/constants.py @@ -1,5 +1,6 @@ from enum import Enum, EnumMeta from typing import Iterator, Type, TypeVar, Union +import logging import torch @@ -7,6 +8,8 @@ __all__ = ["pi", "DType", "Resample", "BorderType", "SamplePadding", "TKEnum"] +logging.basicConfig(level = logging.INFO) + pi = torch.tensor(3.14159265358979323846) diff --git a/kornia/onnx/utils.py b/kornia/onnx/utils.py index 264b1ad1cf..a9a37c68d2 100644 --- a/kornia/onnx/utils.py +++ b/kornia/onnx/utils.py @@ -1,6 +1,7 @@ import os import pprint import urllib.request +import logging from typing import Any, Optional import requests @@ -9,13 +10,15 @@ __all__ = ["ONNXLoader"] +logger = logging.getLogger(__name__) + class ONNXLoader: """Manages ONNX models, handling local caching, downloading from Hugging Face, and loading models. Attributes: cache_dir: The directory where ONNX models are cached locally. - Defaults to None, which will use a default `.kornia_onnx_models` directory. + Defaults to None, which will use a default `.kornia_hub/onnx_models` directory. .. code-block:: python onnx_loader = ONNXLoader() @@ -43,14 +46,14 @@ def _get_file_path(self, model_name: str, cache_dir: Optional[str]) -> str: if self.cache_dir is not None: cache_dir = self.cache_dir else: - cache_dir = ".kornia_onnx_models" + cache_dir = ".kornia_hub/onnx_models" # The filename is the model name (without directory path) file_name = f"{model_name.split('/')[-1]}.onnx" file_path = os.path.join(cache_dir, "/".join(model_name.split("/")[:-1]), file_name) return file_path - def load_model(self, model_name: str, download: bool = False, **kwargs) -> "onnx.ModelProto": # type:ignore + def load_model(self, model_name: str, download: bool = True, **kwargs) -> "onnx.ModelProto": # type:ignore """Loads an ONNX model from the local cache or downloads it from Hugging Face if necessary. Args: @@ -65,11 +68,15 @@ def load_model(self, model_name: str, download: bool = False, **kwargs) -> "onnx """ if model_name.startswith("hf://"): model_name = model_name[len("hf://") :] - file_path = self._get_file_path(model_name, self.cache_dir) + cache_dir = kwargs.get("cache_dir", None) or self.cache_dir + file_path = self._get_file_path(model_name, cache_dir) if not os.path.exists(file_path): # Construct the raw URL for the ONNX file - url = f"https://huggingface.co/kornia/ONNX_models/resolve/main/{model_name}.onnx" - self.download(url, file_path, **kwargs) + if download: + url = f"https://huggingface.co/kornia/ONNX_models/resolve/main/{model_name}.onnx" + self.download(url, file_path) + else: + raise ValueError(f"`{model_name}` is not found in `{file_path}`. You may set `download=True`.") return onnx.load(file_path) # type:ignore if os.path.exists(model_name): @@ -81,7 +88,6 @@ def download( self, url: str, file_path: str, - cache_dir: Optional[str] = None, ) -> None: """Downloads an ONNX model from the specified URL and saves it to the specified file path. @@ -94,8 +100,9 @@ def download( os.makedirs(os.path.dirname(file_path), exist_ok=True) # Create the cache directory if it doesn't exist - if not url.startswith(("http:", "https:")): + if url.startswith(("http:", "https:")): try: + logger.info(f"Downloading `{url}` to `{file_path}`.") urllib.request.urlretrieve(url, file_path) except urllib.error.HTTPError as e: raise ValueError(f"Error in resolving `{url}`. {e}.") From a4c188cc7f3935a2dd9c9316e22adf0b6e18066f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 8 Sep 2024 22:31:12 +0000 Subject: [PATCH 25/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/onnx/utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/kornia/onnx/utils.py b/kornia/onnx/utils.py index a9a37c68d2..a91a396a36 100644 --- a/kornia/onnx/utils.py +++ b/kornia/onnx/utils.py @@ -109,7 +109,6 @@ def download( else: raise ValueError("URL must start with 'http:' or 'https:'") - @staticmethod def _fetch_repo_contents(folder: str) -> list[dict[str, Any]]: """Fetches the contents of the Hugging Face repository using the Hugging Face API. From 62c429532511f89a598aa8573639837640d7ca47 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 9 Sep 2024 08:17:43 +0000 Subject: [PATCH 26/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/constants.py | 4 ++-- kornia/onnx/utils.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/kornia/constants.py b/kornia/constants.py index 19fd9de6dd..c4bc6f3d84 100644 --- a/kornia/constants.py +++ b/kornia/constants.py @@ -1,6 +1,6 @@ +import logging from enum import Enum, EnumMeta from typing import Iterator, Type, TypeVar, Union -import logging import torch @@ -8,7 +8,7 @@ __all__ = ["pi", "DType", "Resample", "BorderType", "SamplePadding", "TKEnum"] -logging.basicConfig(level = logging.INFO) +logging.basicConfig(level=logging.INFO) pi = torch.tensor(3.14159265358979323846) diff --git a/kornia/onnx/utils.py b/kornia/onnx/utils.py index a91a396a36..fbcf513066 100644 --- a/kornia/onnx/utils.py +++ b/kornia/onnx/utils.py @@ -1,7 +1,7 @@ +import logging import os import pprint import urllib.request -import logging from typing import Any, Optional import requests From 0a15677785308300dfdb4b6a2e22881c8e9ba74e Mon Sep 17 00:00:00 2001 From: shijianjian Date: Mon, 9 Sep 2024 14:41:53 +0300 Subject: [PATCH 27/84] update --- docs/source/onnx.rst | 16 ++++++++++++++-- kornia/color/gray.py | 12 ++++++------ kornia/onnx/sequential.py | 8 -------- kornia/onnx/utils.py | 9 +-------- 4 files changed, 21 insertions(+), 24 deletions(-) diff --git a/docs/source/onnx.rst b/docs/source/onnx.rst index c08e70e24c..ff3381c562 100644 --- a/docs/source/onnx.rst +++ b/docs/source/onnx.rst @@ -81,8 +81,9 @@ Here's how you can quickly get started with `ONNXSequential`: # Initialize with CUDA execution provider onnx_seq = ONNXSequential( - "hf://operators/kornia.color.gray.RgbToGrayscale", - "hf://operators/kornia.geometry.transform.affwarp.Resize_512x512", + "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'] ) @@ -122,3 +123,14 @@ 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") + # Load a local converted/downloaded operator + onnx_loader.load_model("operators/kornia.color.gray.GrayscaleToRgb") + + :members: \ No newline at end of file diff --git a/kornia/color/gray.py b/kornia/color/gray.py index 714673d862..c3420f60d9 100644 --- a/kornia/color/gray.py +++ b/kornia/color/gray.py @@ -125,8 +125,8 @@ 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] + 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) @@ -150,8 +150,8 @@ 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] + 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__() @@ -181,8 +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] + 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) diff --git a/kornia/onnx/sequential.py b/kornia/onnx/sequential.py index a630668c1c..1961d441aa 100644 --- a/kornia/onnx/sequential.py +++ b/kornia/onnx/sequential.py @@ -27,14 +27,6 @@ class ONNXSequential: cache_dir: cache_dir: The directory where ONNX models are cached locally (only for downloading from HuggingFace). Defaults to None, which will use a default `.kornia_onnx_models` directory. - - .. code-block:: python - # Load ops from HuggingFace repos then chain to your own model! - model = kornia.onnx.ONNXSequential( - "hf://operators/kornia.color.gray.RgbToGrayscale", - "hf://operators/kornia.geometry.transform.affwarp.Resize_512x512", - "MY_OTHER_MODEL.onnx" - ) """ def __init__( diff --git a/kornia/onnx/utils.py b/kornia/onnx/utils.py index fbcf513066..9a39dc84c6 100644 --- a/kornia/onnx/utils.py +++ b/kornia/onnx/utils.py @@ -19,13 +19,6 @@ class ONNXLoader: Attributes: cache_dir: The directory where ONNX models are cached locally. Defaults to None, which will use a default `.kornia_hub/onnx_models` directory. - - .. code-block:: python - onnx_loader = ONNXLoader() - # Load a HuggingFace operator - onnx_loader.load_model("hf://operators/kornia.color.gray.GrayscaleToRgb") - # Load a local converted/downloaded operator - onnx_loader.load_model("operators/kornia.color.gray.GrayscaleToRgb") """ def __init__(self, cache_dir: Optional[str] = None): @@ -103,7 +96,7 @@ def download( if url.startswith(("http:", "https:")): try: logger.info(f"Downloading `{url}` to `{file_path}`.") - urllib.request.urlretrieve(url, file_path) + urllib.request.urlretrieve(url, file_path) # noqa: S310 except urllib.error.HTTPError as e: raise ValueError(f"Error in resolving `{url}`. {e}.") else: From 88979410bdd3c3cf08237b9c09413ebf7aa9d384 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Tue, 10 Sep 2024 22:57:58 +0300 Subject: [PATCH 28/84] update --- kornia/enhance/normalize.py | 101 ++++++++++++++++++------------------ kornia/utils/__init__.py | 2 + kornia/utils/sample.py | 81 +++++++++++++++++++++++++++++ 3 files changed, 134 insertions(+), 50 deletions(-) create mode 100644 kornia/utils/sample.py diff --git a/kornia/enhance/normalize.py b/kornia/enhance/normalize.py index 423d6ee3bd..6e758527f1 100644 --- a/kornia/enhance/normalize.py +++ b/kornia/enhance/normalize.py @@ -100,28 +100,34 @@ def normalize(data: Tensor, mean: Tensor, std: Tensor) -> Tensor: torch.Size([1, 4, 3, 3]) """ shape = data.shape - if len(mean.shape) == 0 or mean.shape[0] == 1: - mean = mean.expand(shape[1]) - if len(std.shape) == 0 or std.shape[0] == 1: - std = std.expand(shape[1]) - # Allow broadcast on channel dimension - if mean.shape and mean.shape[0] != 1: - if mean.shape[0] != data.shape[1] and mean.shape[:2] != data.shape[:2]: - raise ValueError(f"mean length and number of channels do not match. Got {mean.shape} and {data.shape}.") + if torch.onnx.is_in_onnx_export(): + if not isinstance(mean, Tensor) or not isinstance(std, Tensor): + raise ValueError("Only tensor is accepted when converting to ONNX.") + if mean.shape[0] != 1 or std.shape[0] != 1: + raise ValueError("Batch dimension must be one for broadcasting when converting to ONNX.") + else: + if isinstance(mean, float): + mean = torch.tensor([mean] * shape[1], device=data.device, dtype=data.dtype) - # Allow broadcast on channel dimension - if std.shape and std.shape[0] != 1: - if std.shape[0] != data.shape[1] and std.shape[:2] != data.shape[:2]: - raise ValueError(f"std length and number of channels do not match. Got {std.shape} and {data.shape}.") + if isinstance(std, float): + std = torch.tensor([std] * shape[1], device=data.device, dtype=data.dtype) - mean = torch.as_tensor(mean, device=data.device, dtype=data.dtype) - std = torch.as_tensor(std, device=data.device, dtype=data.dtype) + # Allow broadcast on channel dimension + if mean.shape and mean.shape[0] != 1: + if mean.shape[0] != data.shape[1] and mean.shape[:2] != data.shape[:2]: + raise ValueError(f"mean length and number of channels do not match. Got {mean.shape} and {data.shape}.") - if mean.shape: - mean = mean[..., :, None] - if std.shape: - std = std[..., :, None] + # Allow broadcast on channel dimension + if std.shape and std.shape[0] != 1: + if std.shape[0] != data.shape[1] and std.shape[:2] != data.shape[:2]: + raise ValueError(f"std length and number of channels do not match. Got {std.shape} and {data.shape}.") + + mean = torch.as_tensor(mean, device=data.device, dtype=data.dtype) + std = torch.as_tensor(std, device=data.device, dtype=data.dtype) + + mean = mean[..., :, None] + std = std[..., :, None] out: Tensor = (data.view(shape[0], shape[1], -1) - mean) / std @@ -203,38 +209,33 @@ def denormalize(data: Tensor, mean: Union[Tensor, float], std: Union[Tensor, flo """ shape = data.shape - if isinstance(mean, float): - mean = torch.tensor([mean] * shape[1], device=data.device, dtype=data.dtype) - - if isinstance(std, float): - std = torch.tensor([std] * shape[1], device=data.device, dtype=data.dtype) - - if not isinstance(data, Tensor): - raise TypeError(f"data should be a tensor. Got {type(data)}") - - if not isinstance(mean, Tensor): - raise TypeError(f"mean should be a tensor or a float. Got {type(mean)}") - - if not isinstance(std, Tensor): - raise TypeError(f"std should be a tensor or float. Got {type(std)}") - - # Allow broadcast on channel dimension - if mean.shape and mean.shape[0] != 1: - if mean.shape[0] != data.shape[-3] and mean.shape[:2] != data.shape[:2]: - raise ValueError(f"mean length and number of channels do not match. Got {mean.shape} and {data.shape}.") - - # Allow broadcast on channel dimension - if std.shape and std.shape[0] != 1: - if std.shape[0] != data.shape[-3] and std.shape[:2] != data.shape[:2]: - raise ValueError(f"std length and number of channels do not match. Got {std.shape} and {data.shape}.") - - mean = torch.as_tensor(mean, device=data.device, dtype=data.dtype) - std = torch.as_tensor(std, device=data.device, dtype=data.dtype) - - if mean.shape: - mean = mean[..., :, None] - if std.shape: - std = std[..., :, None] + if torch.onnx.is_in_onnx_export(): + if not isinstance(mean, Tensor) or not isinstance(std, Tensor): + raise ValueError("Only tensor is accepted when converting to ONNX.") + if mean.shape[0] != 1 or std.shape[0] != 1: + raise ValueError("Batch dimension must be one for broadcasting when converting to ONNX.") + else: + if isinstance(mean, float): + mean = torch.tensor([mean] * shape[1], device=data.device, dtype=data.dtype) + + if isinstance(std, float): + std = torch.tensor([std] * shape[1], device=data.device, dtype=data.dtype) + + # Allow broadcast on channel dimension + if mean.shape and mean.shape[0] != 1: + if mean.shape[0] != data.shape[-3] and mean.shape[:2] != data.shape[:2]: + raise ValueError(f"mean length and number of channels do not match. Got {mean.shape} and {data.shape}.") + + # Allow broadcast on channel dimension + if std.shape and std.shape[0] != 1: + if std.shape[0] != data.shape[-3] and std.shape[:2] != data.shape[:2]: + raise ValueError(f"std length and number of channels do not match. Got {std.shape} and {data.shape}.") + + mean = torch.as_tensor(mean, device=data.device, dtype=data.dtype) + std = torch.as_tensor(std, device=data.device, dtype=data.dtype) + + mean = mean[..., :, None] + std = std[..., :, None] out: Tensor = (data.view(shape[0], shape[1], -1) * std) + mean diff --git a/kornia/utils/__init__.py b/kornia/utils/__init__.py index 8fc674f6fb..95cfe1016e 100644 --- a/kornia/utils/__init__.py +++ b/kornia/utils/__init__.py @@ -28,6 +28,7 @@ ) from .one_hot import one_hot from .pointcloud_io import load_pointcloud_ply, save_pointcloud_ply +from .sample import get_sample_images __all__ = [ "batched_forward", @@ -62,4 +63,5 @@ "is_mps_tensor_safe", "dataclass_to_dict", "dict_to_dataclass", + "get_sample_images" ] diff --git a/kornia/utils/sample.py b/kornia/utils/sample.py new file mode 100644 index 0000000000..ee0d2d7351 --- /dev/null +++ b/kornia/utils/sample.py @@ -0,0 +1,81 @@ +import os +import logging +import torch +import kornia +import requests +from io import BytesIO +from typing import Optional, Union + +from kornia.io import load_image +from kornia.core import Tensor, stack +from kornia.core.external import PILImage as Image + +__all__ = [ + "get_sample_images", +] + +IMAGE_URLS: list[str] = [ + "https://raw.githubusercontent.com/kornia/data/main/panda.jpg", + "https://raw.githubusercontent.com/kornia/data/main/simba.png", + "https://raw.githubusercontent.com/kornia/data/main/girona.png", + "https://raw.githubusercontent.com/kornia/data/main/baby_giraffe.png", + "https://raw.githubusercontent.com/kornia/data/main/persistencia_memoria.jpg", + "https://raw.githubusercontent.com/kornia/data/main/delorean.png", +] + + +def download_image(url: str, save_to: str) -> None: + """Download an image from a given URL and save it to a specified file path. + + Args: + url: The URL of the image to download. + save_to: The file path where the downloaded image will be saved. + """ + im = Image.open(requests.get(url, stream=True).raw) # type:ignore + im.save(save_to) # type: ignore + + +def get_sample_images( + resize: tuple[int, int] = None, paths: list[str] = IMAGE_URLS, download: bool = True, + cache_dir: Optional[str] = None +) -> Union[Tensor, list[Tensor]]: + """Loads multiple images from the given URLs. + + Optionally download them, resize them if specified, and return them as a batch of tensors or a list of tensors. + + Args: + paths: A list of path or URL from which to load or download images. + Defaults to a pre-defined constant `IMAGE_URLS` if not provided. + resize: Optional target size for resizing all images as a tuple (height, width). + If not provided, the images will not be resized, and their original sizes will be retained. + download (bool): Whether to download the images if they are not already cached. Defaults to True. + cache_dir (Optional[str]): The directory where the downloaded images will be cached. + Defaults to ".kornia_hub/images". + + Returns: + torch.Tensor | list[torch.Tensor]: + If `resize` is provided, returns a single stacked tensor with shape (B, C, H, W). + Otherwise, returns a list of tensors, each with its original shape (C, H, W). + """ + if cache_dir is None: + cache_dir = ".kornia_hub/images" + os.makedirs(cache_dir, exist_ok=True) + tensors = [] + for path in paths: + if path.startswith("http"): + name = os.path.basename(path) + fname = os.path.join(cache_dir, name) + if not os.path.exists(fname): + logging.info(f"Downloading `{path}` to `{fname}`.") + download_image(path, fname) + else: + fname = path + img_tensor = load_image(fname) + if resize is not None: + img_tensor = kornia.geometry.resize(img_tensor, resize) + tensors.append(img_tensor) + + if resize is not None: + return stack(tensors) + else: + return tensors From 1f0ae6e298349f033a738be559bc52666e001a82 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 10 Sep 2024 19:59:08 +0000 Subject: [PATCH 29/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- docs/source/onnx.rst | 2 +- kornia/color/gray.py | 2 +- kornia/geometry/transform/affwarp.py | 2 +- kornia/utils/__init__.py | 2 +- kornia/utils/sample.py | 19 ++++++++++--------- 5 files changed, 14 insertions(+), 13 deletions(-) diff --git a/docs/source/onnx.rst b/docs/source/onnx.rst index ff3381c562..626313db2b 100644 --- a/docs/source/onnx.rst +++ b/docs/source/onnx.rst @@ -133,4 +133,4 @@ API Documentation # Load a local converted/downloaded operator onnx_loader.load_model("operators/kornia.color.gray.GrayscaleToRgb") - :members: \ No newline at end of file + :members: diff --git a/kornia/color/gray.py b/kornia/color/gray.py index c3420f60d9..f647cfd337 100644 --- a/kornia/color/gray.py +++ b/kornia/color/gray.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import ClassVar, List, Optional +from typing import ClassVar, Optional import torch diff --git a/kornia/geometry/transform/affwarp.py b/kornia/geometry/transform/affwarp.py index ccc18c0ec6..3abb418ec4 100644 --- a/kornia/geometry/transform/affwarp.py +++ b/kornia/geometry/transform/affwarp.py @@ -1,5 +1,5 @@ import warnings -from typing import ClassVar, Optional, Tuple, Union +from typing import Optional, Tuple, Union import torch diff --git a/kornia/utils/__init__.py b/kornia/utils/__init__.py index 95cfe1016e..61806c28f8 100644 --- a/kornia/utils/__init__.py +++ b/kornia/utils/__init__.py @@ -63,5 +63,5 @@ "is_mps_tensor_safe", "dataclass_to_dict", "dict_to_dataclass", - "get_sample_images" + "get_sample_images", ] diff --git a/kornia/utils/sample.py b/kornia/utils/sample.py index ee0d2d7351..feed129b6b 100644 --- a/kornia/utils/sample.py +++ b/kornia/utils/sample.py @@ -1,14 +1,13 @@ -import os import logging -import torch -import kornia -import requests -from io import BytesIO +import os from typing import Optional, Union -from kornia.io import load_image +import requests + +import kornia from kornia.core import Tensor, stack from kornia.core.external import PILImage as Image +from kornia.io import load_image __all__ = [ "get_sample_images", @@ -36,11 +35,13 @@ def download_image(url: str, save_to: str) -> None: def get_sample_images( - resize: tuple[int, int] = None, paths: list[str] = IMAGE_URLS, download: bool = True, - cache_dir: Optional[str] = None + resize: tuple[int, int] = None, + paths: list[str] = IMAGE_URLS, + download: bool = True, + cache_dir: Optional[str] = None, ) -> Union[Tensor, list[Tensor]]: """Loads multiple images from the given URLs. - + Optionally download them, resize them if specified, and return them as a batch of tensors or a list of tensors. Args: From 37032b6d98c323e010836cc7832e3d051aa8bb7e Mon Sep 17 00:00:00 2001 From: shijianjian Date: Tue, 10 Sep 2024 23:11:05 +0300 Subject: [PATCH 30/84] update --- kornia/models/detector/utils.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) create mode 100644 kornia/models/detector/utils.py diff --git a/kornia/models/detector/utils.py b/kornia/models/detector/utils.py new file mode 100644 index 0000000000..f0a3d212be --- /dev/null +++ b/kornia/models/detector/utils.py @@ -0,0 +1,21 @@ +from typing import Optional + +from kornia.core import Module + + +class BoxFiltering(Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, boxes: Tensor, confidence_threshold: Tensor) -> Tensor: + """Filter boxes according to the desired threshold. + + Args: + boxes: [B, D, 6], where B is the batchsize, D is the number of detections in the image, + 6 represent (class_id, confidence_score, x, y, w, h). + confidence_threshold: an 0-d scalar that represents the desired threshold. + """ + + return boxes[(boxes[:, :, 1] > confidence_threshold).unsqueeze(-1).expand_as(boxes)].view( + boxes.shape[0], -1, boxes.shape[-1] + ) From 88983c7dfc7de8d83f72a79f58f50df4d0cc4ab4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 11 Sep 2024 08:52:11 +0000 Subject: [PATCH 31/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/models/detector/utils.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/kornia/models/detector/utils.py b/kornia/models/detector/utils.py index f0a3d212be..f9f8c01a38 100644 --- a/kornia/models/detector/utils.py +++ b/kornia/models/detector/utils.py @@ -1,12 +1,10 @@ -from typing import Optional - from kornia.core import Module class BoxFiltering(Module): def __init__(self) -> None: super().__init__() - + def forward(self, boxes: Tensor, confidence_threshold: Tensor) -> Tensor: """Filter boxes according to the desired threshold. From b8dc7838ba926b1cd6c7f15cc6c647fb59eef601 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Wed, 11 Sep 2024 12:04:19 +0300 Subject: [PATCH 32/84] update --- kornia/models/detector/utils.py | 2 +- kornia/utils/sample.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/kornia/models/detector/utils.py b/kornia/models/detector/utils.py index f9f8c01a38..8472220110 100644 --- a/kornia/models/detector/utils.py +++ b/kornia/models/detector/utils.py @@ -1,4 +1,4 @@ -from kornia.core import Module +from kornia.core import Module, Tensor class BoxFiltering(Module): diff --git a/kornia/utils/sample.py b/kornia/utils/sample.py index feed129b6b..6e1165f315 100644 --- a/kornia/utils/sample.py +++ b/kornia/utils/sample.py @@ -1,6 +1,6 @@ import logging import os -from typing import Optional, Union +from typing import List, Optional, Union import requests @@ -13,7 +13,7 @@ "get_sample_images", ] -IMAGE_URLS: list[str] = [ +IMAGE_URLS: List[str] = [ "https://raw.githubusercontent.com/kornia/data/main/panda.jpg", "https://raw.githubusercontent.com/kornia/data/main/simba.png", "https://raw.githubusercontent.com/kornia/data/main/girona.png", @@ -30,12 +30,12 @@ def download_image(url: str, save_to: str) -> None: url: The URL of the image to download. save_to: The file path where the downloaded image will be saved. """ - im = Image.open(requests.get(url, stream=True).raw) # type:ignore + im = Image.open(requests.get(url, stream=True, timeout=30).raw) # type:ignore im.save(save_to) # type: ignore def get_sample_images( - resize: tuple[int, int] = None, + resize: Optional[tuple[int, int]] = None, paths: list[str] = IMAGE_URLS, download: bool = True, cache_dir: Optional[str] = None, From 1625fd352b3c7e223bdc87c6bc40613436b1e4ad Mon Sep 17 00:00:00 2001 From: shijianjian Date: Wed, 11 Sep 2024 12:12:17 +0300 Subject: [PATCH 33/84] update --- kornia/utils/sample.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kornia/utils/sample.py b/kornia/utils/sample.py index 6e1165f315..17177efcd9 100644 --- a/kornia/utils/sample.py +++ b/kornia/utils/sample.py @@ -31,7 +31,7 @@ def download_image(url: str, save_to: str) -> None: save_to: The file path where the downloaded image will be saved. """ im = Image.open(requests.get(url, stream=True, timeout=30).raw) # type:ignore - im.save(save_to) # type: ignore + im.save(save_to) def get_sample_images( From 6835aee96466b2a8fc458098c8df948a0d38ef98 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Wed, 11 Sep 2024 15:14:20 +0300 Subject: [PATCH 34/84] update --- kornia/utils/sample.py | 4 ++-- requirements/requirements-docs.txt | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/kornia/utils/sample.py b/kornia/utils/sample.py index 17177efcd9..aea98cbcbb 100644 --- a/kornia/utils/sample.py +++ b/kornia/utils/sample.py @@ -36,10 +36,10 @@ def download_image(url: str, save_to: str) -> None: def get_sample_images( resize: Optional[tuple[int, int]] = None, - paths: list[str] = IMAGE_URLS, + paths: List[str] = IMAGE_URLS, download: bool = True, cache_dir: Optional[str] = None, -) -> Union[Tensor, list[Tensor]]: +) -> Union[Tensor, List[Tensor]]: """Loads multiple images from the given URLs. Optionally download them, resize them if specified, and return them as a batch of tensors or a list of tensors. diff --git a/requirements/requirements-docs.txt b/requirements/requirements-docs.txt index 770406e842..f8fd423bcf 100644 --- a/requirements/requirements-docs.txt +++ b/requirements/requirements-docs.txt @@ -1,6 +1,7 @@ furo kornia_moons matplotlib +onnx opencv-python PyYAML>=5.1 sphinx From 4ebe1afb8d26b7c1051b1aa33bde8e2eb9ba0e3d Mon Sep 17 00:00:00 2001 From: shijianjian Date: Fri, 13 Sep 2024 01:16:45 +0300 Subject: [PATCH 35/84] update --- requirements/requirements-docs.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/requirements/requirements-docs.txt b/requirements/requirements-docs.txt index f8fd423bcf..87d2805a71 100644 --- a/requirements/requirements-docs.txt +++ b/requirements/requirements-docs.txt @@ -2,6 +2,7 @@ furo kornia_moons matplotlib onnx +onnxruntime opencv-python PyYAML>=5.1 sphinx From 9459fda76dda3fd4a27d20c2967d1701e8020885 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Fri, 13 Sep 2024 09:01:16 +0300 Subject: [PATCH 36/84] update --- kornia/onnx/sequential.py | 23 ++++++++--------------- 1 file changed, 8 insertions(+), 15 deletions(-) diff --git a/kornia/onnx/sequential.py b/kornia/onnx/sequential.py index 1961d441aa..672870129d 100644 --- a/kornia/onnx/sequential.py +++ b/kornia/onnx/sequential.py @@ -13,20 +13,16 @@ class ONNXSequential: """ONNXSequential to chain multiple ONNX operators together. Args: - *args: - A variable number of ONNX models (either ONNX ModelProto objects or file paths). - providers: - A list of execution providers for ONNXRuntime (e.g., ['CUDAExecutionProvider', 'CPUExecutionProvider']). - session_options: - Optional ONNXRuntime session options for optimizing the session. - io_maps: - An optional list of list of tuples specifying input-output mappings for combining models. + *args: A variable number of ONNX models (either ONNX ModelProto objects or file paths). + providers: A list of execution providers for ONNXRuntime + (e.g., ['CUDAExecutionProvider', 'CPUExecutionProvider']). + session_options: Optional ONNXRuntime session options for optimizing the session. + io_maps: An optional list of list of tuples specifying input-output mappings for combining models. If None, we assume the default input name and output name are "input" and "output" accordingly, and only one input and output node for each graph. If not None, `io_maps[0]` shall represent the `io_map` for combining the first and second ONNX models. - cache_dir: - cache_dir: The directory where ONNX models are cached locally (only for downloading from HuggingFace). - Defaults to None, which will use a default `.kornia_onnx_models` directory. + cache_dir: The directory where ONNX models are cached locally (only for downloading from HuggingFace). + Defaults to None, which will use a default `.kornia_hub/onnx_models` directory. """ def __init__( @@ -66,9 +62,6 @@ def _combine(self, io_maps: Optional[list[tuple[str, str]]] = None) -> "onnx.Mod Returns: onnx.ModelProto: The combined ONNX model as a single ONNX graph. - - Raises: - ValueError: If no operators are provided for combination. """ if len(self.operators) == 0: raise ValueError("No operators found.") @@ -90,7 +83,7 @@ def export(self, file_path: str) -> None: """Export the combined ONNX model to a file. Args: - file_path: str + file_path: The file path to export the combined ONNX model. """ onnx.save(self._combined_op, file_path) # type:ignore From 6c172ee30876428080540e4cd28f446c401df40b Mon Sep 17 00:00:00 2001 From: shijianjian Date: Fri, 13 Sep 2024 16:01:19 +0300 Subject: [PATCH 37/84] update --- kornia/__init__.py | 1 + .../rt_detr/architecture/rtdetr_head.py | 1 + .../contrib/models/rt_detr/post_processor.py | 11 ++--- kornia/core/__init__.py | 4 +- kornia/core/module.py | 40 ++++++++++++------ kornia/filters/motion.py | 1 + kornia/models/__init__.py | 1 + kornia/models/detector/__init__.py | 1 + kornia/models/detector/rtdetr.py | 17 ++++---- kornia/models/detector/utils.py | 42 +++++++++++++++---- 10 files changed, 85 insertions(+), 34 deletions(-) diff --git a/kornia/__init__.py b/kornia/__init__.py index 5896d56f52..9bdc9c1d66 100644 --- a/kornia/__init__.py +++ b/kornia/__init__.py @@ -15,6 +15,7 @@ io, losses, metrics, + models, morphology, onnx, tracking, diff --git a/kornia/contrib/models/rt_detr/architecture/rtdetr_head.py b/kornia/contrib/models/rt_detr/architecture/rtdetr_head.py index 6345577e5a..3c731238b5 100644 --- a/kornia/contrib/models/rt_detr/architecture/rtdetr_head.py +++ b/kornia/contrib/models/rt_detr/architecture/rtdetr_head.py @@ -274,6 +274,7 @@ def __init__( num_denoising: int = 100, ) -> None: super().__init__() + self.num_classes = num_classes self.num_queries = num_queries # TODO: verify this is correct if len(in_channels) > num_levels: diff --git a/kornia/contrib/models/rt_detr/post_processor.py b/kornia/contrib/models/rt_detr/post_processor.py index 95b35623df..d62e5be7f0 100644 --- a/kornia/contrib/models/rt_detr/post_processor.py +++ b/kornia/contrib/models/rt_detr/post_processor.py @@ -2,11 +2,12 @@ from __future__ import annotations -from typing import Optional +from typing import Optional, Union, List import torch from kornia.core import Module, Tensor, concatenate +from kornia.models.detector.utils import BoxFiltering def mod(a: Tensor, b: int) -> Tensor: @@ -37,14 +38,16 @@ def __init__( num_classes: int = 80, num_top_queries: int = 300, confidence_filtering: bool = True, + filter_as_zero: bool = False, ) -> None: super().__init__() self.confidence_threshold = confidence_threshold self.num_classes = num_classes self.confidence_filtering = confidence_filtering self.num_top_queries = num_top_queries + self.box_filtering = BoxFiltering(confidence_threshold, filter_as_zero=filter_as_zero) - def forward(self, logits: Tensor, boxes: Tensor, original_sizes: Tensor) -> Tensor: + def forward(self, logits: Tensor, boxes: Tensor, original_sizes: Tensor) -> Union[Tensor, List[Tensor]]: """Post-process outputs from DETR. Args: @@ -88,6 +91,4 @@ def forward(self, logits: Tensor, boxes: Tensor, original_sizes: Tensor) -> Tens if not self.confidence_filtering or self.confidence_threshold == 0: return all_boxes - return all_boxes[(all_boxes[:, :, 1] > self.confidence_threshold).unsqueeze(-1).expand_as(all_boxes)].view( - all_boxes.shape[0], -1, all_boxes.shape[-1] - ) + return self.box_filtering(all_boxes, self.confidence_threshold) diff --git a/kornia/core/__init__.py b/kornia/core/__init__.py index 2e1ebc2f7c..c715601fd4 100644 --- a/kornia/core/__init__.py +++ b/kornia/core/__init__.py @@ -32,7 +32,7 @@ zeros, zeros_like, ) -from .module import ImageModule +from .module import ImageModule, ONNXExportMixin, ImageModuleMixIn from .tensor_wrapper import TensorWrapper # type: ignore __all__ = [ @@ -70,4 +70,6 @@ "TensorWrapper", "map_coordinates", "ImageModule", + "ONNXExportMixin", + "ImageModuleMixIn", ] diff --git a/kornia/core/module.py b/kornia/core/module.py index 2c6c5cc10d..140cd924e0 100644 --- a/kornia/core/module.py +++ b/kornia/core/module.py @@ -2,13 +2,13 @@ import math import os from functools import wraps -from typing import Any, Callable, ClassVar, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, ClassVar, List, Optional, Tuple, Union import torch import kornia -from ._backend import Module, Tensor, from_numpy +from ._backend import Module, Tensor, from_numpy, rand from .external import PILImage as Image from .external import numpy as np from .external import onnx @@ -26,6 +26,10 @@ class ONNXExportMixin: ONNX_DEFAULT_OUTPUTSHAP: Default output shape for the ONNX export. A list of integers where `-1` indicates dynamic dimensions. Default is [-1, -1, -1, -1]. + ONNX_EXPORT_PSEUDO_SHAPE: + This is used to create a dummy input tensor for the ONNX export. Default is [1, 3, 256, 256]. + It dimension shall match the ONNX_DEFAULT_INPUTSHAPE and ONNX_DEFAULT_OUTPUTSHAPE. + Non-image dimensions are allowed. Note: - If `ONNX_EXPORTABLE` is False, indicating that the object cannot be exported to ONNX. @@ -34,6 +38,8 @@ class ONNXExportMixin: ONNX_EXPORTABLE: bool = True ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, -1, -1, -1] ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, -1, -1, -1] + ONNX_EXPORT_PSEUDO_SHAPE: ClassVar[List[int]] = [1, 3, 256, 256] + ADDITIONAL_METADATA: ClassVar[List[Tuple[str, str]]] = [] def to_onnx( self, @@ -70,15 +76,8 @@ def to_onnx( if onnx_name is None: onnx_name = f"Kornia-{self.__class__.__name__}.onnx" - # Creating a dummy input with the given shape - pseudo_shape = (1, 3, 256, 256) - dummy_input = torch.randn(*[(pseudo_shape[i] if dim == -1 else dim) for i, dim in enumerate(input_shape)]) - - # Dynamic axis configuration for input and output - dynamic_axes = { - "input": {i: "dim_" + str(i) for i, dim in enumerate(input_shape) if dim == -1}, - "output": {i: "dim_" + str(i) for i, dim in enumerate(output_shape) if dim == -1}, - } + dummy_input = self._create_dummy_input(input_shape) + dynamic_axes = self._create_dynamic_axes(input_shape, output_shape) torch.onnx.export( self, # type: ignore @@ -92,9 +91,26 @@ def to_onnx( dynamic_axes=dynamic_axes, ) + self._add_metadata(onnx_name) + + def _create_dummy_input(self, input_shape: List[int]) -> Union[Tuple[Any, ...], Tensor]: + return rand(*[ + (self.ONNX_EXPORT_PSEUDO_SHAPE[i] if dim == -1 else dim) for i, dim in enumerate(input_shape)]) + + def _create_dynamic_axes(self, input_shape: List[int], output_shape: List[int]) -> Dict[str, Dict[int, str]]: + return { + "input": {i: "dim_" + str(i) for i, dim in enumerate(input_shape) if dim == -1}, + "output": {i: "dim_" + str(i) for i, dim in enumerate(output_shape) if dim == -1}, + } + + def _add_metadata(self, onnx_name: str) -> None: onnx_model = onnx.load(onnx_name) # type: ignore - for key, value in [("source", "kornia"), ("version", kornia.__version__), ("class", self.__class__.__name__)]: + for key, value in [ + ("source", "kornia"), + ("version", kornia.__version__), + ("class", self.__class__.__name__), + ] + self.ADDITIONAL_METADATA: metadata_props = onnx_model.metadata_props.add() metadata_props.key = key metadata_props.value = str(value) diff --git a/kornia/filters/motion.py b/kornia/filters/motion.py index 7442f1167f..f084036eb8 100644 --- a/kornia/filters/motion.py +++ b/kornia/filters/motion.py @@ -88,6 +88,7 @@ class MotionBlur3D(Module): ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, -1, -1, -1, -1] ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, -1, -1, -1, -1] + ONNX_EXPORT_PSEUDO_SHAPE: ClassVar[list[int]] = [1, 3, 80, 80, 80] def __init__( self, diff --git a/kornia/models/__init__.py b/kornia/models/__init__.py index e69de29bb2..f7a598a13c 100644 --- a/kornia/models/__init__.py +++ b/kornia/models/__init__.py @@ -0,0 +1 @@ +from . import detector diff --git a/kornia/models/detector/__init__.py b/kornia/models/detector/__init__.py index 55a62efc22..c51e7ed66f 100644 --- a/kornia/models/detector/__init__.py +++ b/kornia/models/detector/__init__.py @@ -1 +1,2 @@ from .rtdetr import * +from .utils import * diff --git a/kornia/models/detector/rtdetr.py b/kornia/models/detector/rtdetr.py index 4bc1c6bd2d..ff78ff2289 100644 --- a/kornia/models/detector/rtdetr.py +++ b/kornia/models/detector/rtdetr.py @@ -26,7 +26,7 @@ def build( config: Optional[RTDETRConfig] = None, pretrained: bool = True, image_size: Optional[int] = 640, - confidence_threshold: float = 0.5, + confidence_threshold: Optional[float] = None, confidence_filtering: Optional[bool] = None, ) -> ObjectDetector: """Builds and returns an RT-DETR object detector model. @@ -46,11 +46,6 @@ def build( The size to which input images will be resized during preprocessing. If None, no resizing will be performed before passing to the model. Recommended scales include [480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800]. - confidence_threshold: - The confidence threshold used during post-processing to filter detections. - confidence_filtering: - If to perform filtering on resulting boxes. If None, the filtering will be blocked when exporting - to ONNX, while it would perform as per confidence_threshold when build the model. Returns: ObjectDetector @@ -70,13 +65,17 @@ def build( warnings.warn("No `model_name` or `config` found. Will build pretrained `rtdetr_r18vd`.") model = RTDETR.from_pretrained("rtdetr_r18vd") + if confidence_threshold is None: + confidence_threshold = config.confidence_threshold if config is not None else 0.3 + return ObjectDetector( model, ResizePreProcessor((image_size, image_size)) if image_size is not None else nn.Identity(), DETRPostProcessor( - confidence_threshold, - num_classes=config.num_classes if config is not None else 80, - confidence_filtering=confidence_filtering or not torch.onnx.is_in_onnx_export, + confidence_threshold=confidence_threshold, + confidence_filtering=confidence_filtering or not torch.onnx.is_in_onnx_export(), + num_classes=model.decoder.num_classes, + num_top_queries=model.decoder.num_queries, ), ) diff --git a/kornia/models/detector/utils.py b/kornia/models/detector/utils.py index 8472220110..938db67f0c 100644 --- a/kornia/models/detector/utils.py +++ b/kornia/models/detector/utils.py @@ -1,11 +1,21 @@ -from kornia.core import Module, Tensor +from kornia.core import Module, ONNXExportMixin, Tensor, rand +from typing import ClassVar, List, Any, Union, Tuple, Optional +__all__ = ["BoxFiltering"] -class BoxFiltering(Module): - def __init__(self) -> None: + +class BoxFiltering(Module, ONNXExportMixin): + + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, -1, 6] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, -1, 6] + ONNX_EXPORT_PSEUDO_SHAPE: ClassVar[List[int]] = [5, 20, 6] + + def __init__(self, confidence_threshold: Optional[Tensor] = None, filter_as_zero: bool = False) -> None: super().__init__() + self.filter_as_zero = filter_as_zero + self.confidence_threshold = confidence_threshold - def forward(self, boxes: Tensor, confidence_threshold: Tensor) -> Tensor: + def forward(self, boxes: Tensor, confidence_threshold: Optional[Tensor] = None) -> Union[Tensor, List[Tensor]]: """Filter boxes according to the desired threshold. Args: @@ -13,7 +23,25 @@ def forward(self, boxes: Tensor, confidence_threshold: Tensor) -> Tensor: 6 represent (class_id, confidence_score, x, y, w, h). confidence_threshold: an 0-d scalar that represents the desired threshold. """ + if confidence_threshold is None: + confidence_threshold = self.confidence_threshold + if confidence_threshold is None: + raise ValueError("`confidence_threshold` must be provided if not set in the constructor.") + + if self.filter_as_zero: + box_mask = (boxes[:, :, 1] > confidence_threshold).unsqueeze(-1).expand_as(boxes) + filtered_boxes = boxes * box_mask.float() + else: + filtered_boxes = [] + for i in range(boxes.shape[0]): + box = boxes[i:i + 1][(boxes[i:i + 1, :, 1] > confidence_threshold).unsqueeze(-1).expand_as(boxes[i:i + 1])] + filtered_boxes.append(box.view(1, -1, boxes.shape[-1])) + + return filtered_boxes - return boxes[(boxes[:, :, 1] > confidence_threshold).unsqueeze(-1).expand_as(boxes)].view( - boxes.shape[0], -1, boxes.shape[-1] - ) + def _create_dummy_input(self, input_shape: List[int]) -> Union[Tuple[Any, ...], Tensor]: + pseudo_input = rand(*[ + (self.ONNX_EXPORT_PSEUDO_SHAPE[i] if dim == -1 else dim) for i, dim in enumerate(input_shape)]) + if self.confidence_threshold is None: + return pseudo_input, 0.1 + return pseudo_input From d883dd348a672fa4369f995fc7cada4ffe63fd8c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 13 Sep 2024 13:03:24 +0000 Subject: [PATCH 38/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/contrib/models/rt_detr/post_processor.py | 2 +- kornia/core/__init__.py | 2 +- kornia/core/module.py | 5 ++--- kornia/models/detector/utils.py | 13 ++++++++----- 4 files changed, 12 insertions(+), 10 deletions(-) diff --git a/kornia/contrib/models/rt_detr/post_processor.py b/kornia/contrib/models/rt_detr/post_processor.py index d62e5be7f0..56049eb6f9 100644 --- a/kornia/contrib/models/rt_detr/post_processor.py +++ b/kornia/contrib/models/rt_detr/post_processor.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Optional, Union, List +from typing import List, Optional, Union import torch diff --git a/kornia/core/__init__.py b/kornia/core/__init__.py index c715601fd4..58ace6afa2 100644 --- a/kornia/core/__init__.py +++ b/kornia/core/__init__.py @@ -32,7 +32,7 @@ zeros, zeros_like, ) -from .module import ImageModule, ONNXExportMixin, ImageModuleMixIn +from .module import ImageModule, ImageModuleMixIn, ONNXExportMixin from .tensor_wrapper import TensorWrapper # type: ignore __all__ = [ diff --git a/kornia/core/module.py b/kornia/core/module.py index 140cd924e0..90973c2f78 100644 --- a/kornia/core/module.py +++ b/kornia/core/module.py @@ -2,7 +2,7 @@ import math import os from functools import wraps -from typing import Any, Callable, Dict, ClassVar, List, Optional, Tuple, Union +from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Union import torch @@ -94,8 +94,7 @@ def to_onnx( self._add_metadata(onnx_name) def _create_dummy_input(self, input_shape: List[int]) -> Union[Tuple[Any, ...], Tensor]: - return rand(*[ - (self.ONNX_EXPORT_PSEUDO_SHAPE[i] if dim == -1 else dim) for i, dim in enumerate(input_shape)]) + return rand(*[(self.ONNX_EXPORT_PSEUDO_SHAPE[i] if dim == -1 else dim) for i, dim in enumerate(input_shape)]) def _create_dynamic_axes(self, input_shape: List[int], output_shape: List[int]) -> Dict[str, Dict[int, str]]: return { diff --git a/kornia/models/detector/utils.py b/kornia/models/detector/utils.py index 938db67f0c..dcb50d0a77 100644 --- a/kornia/models/detector/utils.py +++ b/kornia/models/detector/utils.py @@ -1,11 +1,11 @@ +from typing import Any, ClassVar, List, Optional, Tuple, Union + from kornia.core import Module, ONNXExportMixin, Tensor, rand -from typing import ClassVar, List, Any, Union, Tuple, Optional __all__ = ["BoxFiltering"] class BoxFiltering(Module, ONNXExportMixin): - ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, -1, 6] ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, -1, 6] ONNX_EXPORT_PSEUDO_SHAPE: ClassVar[List[int]] = [5, 20, 6] @@ -34,14 +34,17 @@ def forward(self, boxes: Tensor, confidence_threshold: Optional[Tensor] = None) else: filtered_boxes = [] for i in range(boxes.shape[0]): - box = boxes[i:i + 1][(boxes[i:i + 1, :, 1] > confidence_threshold).unsqueeze(-1).expand_as(boxes[i:i + 1])] + box = boxes[i : i + 1][ + (boxes[i : i + 1, :, 1] > confidence_threshold).unsqueeze(-1).expand_as(boxes[i : i + 1]) + ] filtered_boxes.append(box.view(1, -1, boxes.shape[-1])) return filtered_boxes def _create_dummy_input(self, input_shape: List[int]) -> Union[Tuple[Any, ...], Tensor]: - pseudo_input = rand(*[ - (self.ONNX_EXPORT_PSEUDO_SHAPE[i] if dim == -1 else dim) for i, dim in enumerate(input_shape)]) + pseudo_input = rand( + *[(self.ONNX_EXPORT_PSEUDO_SHAPE[i] if dim == -1 else dim) for i, dim in enumerate(input_shape)] + ) if self.confidence_threshold is None: return pseudo_input, 0.1 return pseudo_input From 478e413040d9681fecf90033fa2561aa3a679dcd Mon Sep 17 00:00:00 2001 From: shijianjian Date: Fri, 13 Sep 2024 18:08:04 +0300 Subject: [PATCH 39/84] update --- kornia/onnx/sequential.py | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/kornia/onnx/sequential.py b/kornia/onnx/sequential.py index 672870129d..c46b56c206 100644 --- a/kornia/onnx/sequential.py +++ b/kornia/onnx/sequential.py @@ -6,7 +6,7 @@ from .utils import ONNXLoader -__all__ = ["ONNXSequential"] +__all__ = ["ONNXSequential", "load"] class ONNXSequential: @@ -14,6 +14,8 @@ class ONNXSequential: Args: *args: A variable number of ONNX models (either ONNX ModelProto objects or file paths). + For Hugging Face-hosted models, use the format 'hf://model_name'. Valid `model_name` can be found on + https://huggingface.co/kornia/ONNX_models. providers: A list of execution providers for ONNXRuntime (e.g., ['CUDAExecutionProvider', 'CPUExecutionProvider']). session_options: Optional ONNXRuntime session options for optimizing the session. @@ -36,7 +38,7 @@ def __init__( self.onnx_loader = ONNXLoader(cache_dir) self.operators = args self._combined_op = self._combine(io_maps) - self._session = self.create_session() + self._session = self.create_session(providers=providers, session_options=session_options) def _load_op(self, arg: Union["onnx.ModelProto", str]) -> "onnx.ModelProto": # type:ignore """Loads an ONNX model, either from a file path or use the provided ONNX ModelProto. @@ -148,3 +150,17 @@ def __call__(self, *inputs: "np.ndarray") -> list["np.ndarray"]: # type:ignore outputs = self._session.run(None, ort_input_values) return outputs + + +def load(model_name: str) -> "ONNXSequential": + """Load an ONNX model from either a file path or HuggingFace. + + The loaded model is an ONNXSequential object, of which you may run the model with + the `__call__` method, with less boilerplate. + + Args: + model_name: The name of the model to load. For Hugging Face-hosted models, + use the format 'hf://model_name'. Valid `model_name` can be found on + https://huggingface.co/kornia/ONNX_models. + """ + return ONNXSequential(model_name) \ No newline at end of file From ef413d95a93582b6a3ec67d335c4f5f631880f23 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 13 Sep 2024 15:09:21 +0000 Subject: [PATCH 40/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/onnx/sequential.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kornia/onnx/sequential.py b/kornia/onnx/sequential.py index c46b56c206..a4f7b68c77 100644 --- a/kornia/onnx/sequential.py +++ b/kornia/onnx/sequential.py @@ -163,4 +163,4 @@ def load(model_name: str) -> "ONNXSequential": use the format 'hf://model_name'. Valid `model_name` can be found on https://huggingface.co/kornia/ONNX_models. """ - return ONNXSequential(model_name) \ No newline at end of file + return ONNXSequential(model_name) From 68843c49dde19c458012da03c42cc22bae6f0575 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Fri, 13 Sep 2024 18:57:08 +0300 Subject: [PATCH 41/84] update --- docs/source/onnx.rst | 4 ++-- kornia/contrib/models/rt_detr/post_processor.py | 7 +++++-- kornia/models/detector/utils.py | 2 +- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/docs/source/onnx.rst b/docs/source/onnx.rst index 626313db2b..c6e50626f8 100644 --- a/docs/source/onnx.rst +++ b/docs/source/onnx.rst @@ -129,8 +129,8 @@ API Documentation .. code-block:: python onnx_loader = ONNXLoader() # Load a HuggingFace operator - onnx_loader.load_model("hf://operators/kornia.color.gray.GrayscaleToRgb") + 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") + onnx_loader.load_model("operators/kornia.color.gray.GrayscaleToRgb") # doctest: +SKIP :members: diff --git a/kornia/contrib/models/rt_detr/post_processor.py b/kornia/contrib/models/rt_detr/post_processor.py index 56049eb6f9..d3be4dce83 100644 --- a/kornia/contrib/models/rt_detr/post_processor.py +++ b/kornia/contrib/models/rt_detr/post_processor.py @@ -6,7 +6,7 @@ import torch -from kornia.core import Module, Tensor, concatenate +from kornia.core import Module, Tensor, concatenate, tensor from kornia.models.detector.utils import BoxFiltering @@ -45,7 +45,10 @@ def __init__( self.num_classes = num_classes self.confidence_filtering = confidence_filtering self.num_top_queries = num_top_queries - self.box_filtering = BoxFiltering(confidence_threshold, filter_as_zero=filter_as_zero) + self.box_filtering = BoxFiltering( + tensor(confidence_threshold) if confidence_threshold is not None else None, + filter_as_zero=filter_as_zero + ) def forward(self, logits: Tensor, boxes: Tensor, original_sizes: Tensor) -> Union[Tensor, List[Tensor]]: """Post-process outputs from DETR. diff --git a/kornia/models/detector/utils.py b/kornia/models/detector/utils.py index dcb50d0a77..dc73a44be6 100644 --- a/kornia/models/detector/utils.py +++ b/kornia/models/detector/utils.py @@ -32,7 +32,7 @@ def forward(self, boxes: Tensor, confidence_threshold: Optional[Tensor] = None) box_mask = (boxes[:, :, 1] > confidence_threshold).unsqueeze(-1).expand_as(boxes) filtered_boxes = boxes * box_mask.float() else: - filtered_boxes = [] + filtered_boxes: list[Tensor] = [] for i in range(boxes.shape[0]): box = boxes[i : i + 1][ (boxes[i : i + 1, :, 1] > confidence_threshold).unsqueeze(-1).expand_as(boxes[i : i + 1]) From fa39dd63cfdeeec607c5da670026576edabc9560 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 13 Sep 2024 15:57:36 +0000 Subject: [PATCH 42/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/contrib/models/rt_detr/post_processor.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/kornia/contrib/models/rt_detr/post_processor.py b/kornia/contrib/models/rt_detr/post_processor.py index d3be4dce83..ade891778a 100644 --- a/kornia/contrib/models/rt_detr/post_processor.py +++ b/kornia/contrib/models/rt_detr/post_processor.py @@ -46,8 +46,7 @@ def __init__( self.confidence_filtering = confidence_filtering self.num_top_queries = num_top_queries self.box_filtering = BoxFiltering( - tensor(confidence_threshold) if confidence_threshold is not None else None, - filter_as_zero=filter_as_zero + tensor(confidence_threshold) if confidence_threshold is not None else None, filter_as_zero=filter_as_zero ) def forward(self, logits: Tensor, boxes: Tensor, original_sizes: Tensor) -> Union[Tensor, List[Tensor]]: From 125b4ebcd56b71ed4d6fb6e1f5c0de8087e4286b Mon Sep 17 00:00:00 2001 From: shijianjian Date: Fri, 13 Sep 2024 19:26:34 +0300 Subject: [PATCH 43/84] update --- docs/source/onnx.rst | 1 + kornia/contrib/models/rt_detr/post_processor.py | 2 +- kornia/core/module.py | 3 ++- kornia/models/detector/utils.py | 6 ++++-- 4 files changed, 8 insertions(+), 4 deletions(-) diff --git a/docs/source/onnx.rst b/docs/source/onnx.rst index c6e50626f8..5ae0ba8787 100644 --- a/docs/source/onnx.rst +++ b/docs/source/onnx.rst @@ -127,6 +127,7 @@ API Documentation .. 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 diff --git a/kornia/contrib/models/rt_detr/post_processor.py b/kornia/contrib/models/rt_detr/post_processor.py index ade891778a..c21951bc6c 100644 --- a/kornia/contrib/models/rt_detr/post_processor.py +++ b/kornia/contrib/models/rt_detr/post_processor.py @@ -49,7 +49,7 @@ def __init__( tensor(confidence_threshold) if confidence_threshold is not None else None, filter_as_zero=filter_as_zero ) - def forward(self, logits: Tensor, boxes: Tensor, original_sizes: Tensor) -> Union[Tensor, List[Tensor]]: + def forward(self, logits: Tensor, boxes: Tensor, original_sizes: Tensor) -> Union[Tensor, list[Tensor]]: """Post-process outputs from DETR. Args: diff --git a/kornia/core/module.py b/kornia/core/module.py index 90973c2f78..76c85c9640 100644 --- a/kornia/core/module.py +++ b/kornia/core/module.py @@ -109,7 +109,8 @@ def _add_metadata(self, onnx_name: str) -> None: ("source", "kornia"), ("version", kornia.__version__), ("class", self.__class__.__name__), - ] + self.ADDITIONAL_METADATA: + *self.ADDITIONAL_METADATA + ]: metadata_props = onnx_model.metadata_props.add() metadata_props.key = key metadata_props.value = str(value) diff --git a/kornia/models/detector/utils.py b/kornia/models/detector/utils.py index dc73a44be6..b150bf87a5 100644 --- a/kornia/models/detector/utils.py +++ b/kornia/models/detector/utils.py @@ -1,4 +1,4 @@ -from typing import Any, ClassVar, List, Optional, Tuple, Union +from typing import Any, ClassVar, List, Optional, Tuple, Union, cast from kornia.core import Module, ONNXExportMixin, Tensor, rand @@ -28,11 +28,13 @@ def forward(self, boxes: Tensor, confidence_threshold: Optional[Tensor] = None) if confidence_threshold is None: raise ValueError("`confidence_threshold` must be provided if not set in the constructor.") + filtered_boxes: Union[Tensor, List[Tensor]] if self.filter_as_zero: box_mask = (boxes[:, :, 1] > confidence_threshold).unsqueeze(-1).expand_as(boxes) filtered_boxes = boxes * box_mask.float() else: - filtered_boxes: list[Tensor] = [] + filtered_boxes = cast(list[Tensor], filtered_boxes) + filtered_boxes = [] for i in range(boxes.shape[0]): box = boxes[i : i + 1][ (boxes[i : i + 1, :, 1] > confidence_threshold).unsqueeze(-1).expand_as(boxes[i : i + 1]) From beed3e8e2b138fed7a213b2897f17614a75e31fd Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 13 Sep 2024 16:27:39 +0000 Subject: [PATCH 44/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/contrib/models/rt_detr/post_processor.py | 2 +- kornia/core/module.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/kornia/contrib/models/rt_detr/post_processor.py b/kornia/contrib/models/rt_detr/post_processor.py index c21951bc6c..ae9c4ac791 100644 --- a/kornia/contrib/models/rt_detr/post_processor.py +++ b/kornia/contrib/models/rt_detr/post_processor.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import List, Optional, Union +from typing import Optional, Union import torch diff --git a/kornia/core/module.py b/kornia/core/module.py index 76c85c9640..12a875df6a 100644 --- a/kornia/core/module.py +++ b/kornia/core/module.py @@ -109,7 +109,7 @@ def _add_metadata(self, onnx_name: str) -> None: ("source", "kornia"), ("version", kornia.__version__), ("class", self.__class__.__name__), - *self.ADDITIONAL_METADATA + *self.ADDITIONAL_METADATA, ]: metadata_props = onnx_model.metadata_props.add() metadata_props.key = key From 9a7be766c26db582bea7909c563103acc9bc3000 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Fri, 13 Sep 2024 19:33:30 +0300 Subject: [PATCH 45/84] update --- kornia/utils/sample.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/kornia/utils/sample.py b/kornia/utils/sample.py index aea98cbcbb..198a47b1d1 100644 --- a/kornia/utils/sample.py +++ b/kornia/utils/sample.py @@ -1,6 +1,6 @@ import logging import os -from typing import List, Optional, Union +from typing import List, Optional, Tuple, Union import requests @@ -35,7 +35,7 @@ def download_image(url: str, save_to: str) -> None: def get_sample_images( - resize: Optional[tuple[int, int]] = None, + resize: Optional[Tuple[int, int]] = None, paths: List[str] = IMAGE_URLS, download: bool = True, cache_dir: Optional[str] = None, @@ -66,9 +66,12 @@ def get_sample_images( if path.startswith("http"): name = os.path.basename(path) fname = os.path.join(cache_dir, name) - if not os.path.exists(fname): + if not os.path.exists(fname) and download: logging.info(f"Downloading `{path}` to `{fname}`.") download_image(path, fname) + elif not os.path.exists(fname) and not download: + logging.error( + f"Image `{path}` not found at `{fname}`. You may want to set `download=True` to download it.") else: fname = path img_tensor = load_image(fname) From e5cc24df40dcf1123155cde7da0bc312945f6d15 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 13 Sep 2024 16:34:41 +0000 Subject: [PATCH 46/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/utils/sample.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/kornia/utils/sample.py b/kornia/utils/sample.py index 198a47b1d1..4fedfa239b 100644 --- a/kornia/utils/sample.py +++ b/kornia/utils/sample.py @@ -71,7 +71,8 @@ def get_sample_images( download_image(path, fname) elif not os.path.exists(fname) and not download: logging.error( - f"Image `{path}` not found at `{fname}`. You may want to set `download=True` to download it.") + f"Image `{path}` not found at `{fname}`. You may want to set `download=True` to download it." + ) else: fname = path img_tensor = load_image(fname) From 6a929e2fb91914f74c660d33e4a8cd5aca60a594 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Fri, 13 Sep 2024 19:52:23 +0300 Subject: [PATCH 47/84] update --- kornia/enhance/normalize.py | 8 ++++---- kornia/models/detector/rtdetr.py | 4 ++-- kornia/models/detector/utils.py | 8 ++++---- tests/contrib/test_object_detector.py | 2 +- 4 files changed, 11 insertions(+), 11 deletions(-) diff --git a/kornia/enhance/normalize.py b/kornia/enhance/normalize.py index 6e758527f1..3f4fd3f447 100644 --- a/kornia/enhance/normalize.py +++ b/kornia/enhance/normalize.py @@ -126,8 +126,8 @@ def normalize(data: Tensor, mean: Tensor, std: Tensor) -> Tensor: mean = torch.as_tensor(mean, device=data.device, dtype=data.dtype) std = torch.as_tensor(std, device=data.device, dtype=data.dtype) - mean = mean[..., :, None] - std = std[..., :, None] + mean = mean[..., None] + std = std[..., None] out: Tensor = (data.view(shape[0], shape[1], -1) - mean) / std @@ -234,8 +234,8 @@ def denormalize(data: Tensor, mean: Union[Tensor, float], std: Union[Tensor, flo mean = torch.as_tensor(mean, device=data.device, dtype=data.dtype) std = torch.as_tensor(std, device=data.device, dtype=data.dtype) - mean = mean[..., :, None] - std = std[..., :, None] + mean = mean[..., None] + std = std[..., None] out: Tensor = (data.view(shape[0], shape[1], -1) * std) + mean diff --git a/kornia/models/detector/rtdetr.py b/kornia/models/detector/rtdetr.py index ff78ff2289..818abc10ee 100644 --- a/kornia/models/detector/rtdetr.py +++ b/kornia/models/detector/rtdetr.py @@ -1,5 +1,5 @@ import warnings -from typing import Optional +from typing import Optional, Tuple import torch from torch import nn @@ -88,7 +88,7 @@ def to_onnx( image_size: Optional[int] = 640, confidence_threshold: float = 0.5, confidence_filtering: Optional[bool] = None, - ) -> tuple[str, ObjectDetector]: + ) -> Tuple[str, ObjectDetector]: """Exports an RT-DETR object detection model to ONNX format. Either `model_name` or `config` must be provided. If neither is provided, diff --git a/kornia/models/detector/utils.py b/kornia/models/detector/utils.py index b150bf87a5..ff33014df9 100644 --- a/kornia/models/detector/utils.py +++ b/kornia/models/detector/utils.py @@ -33,13 +33,13 @@ def forward(self, boxes: Tensor, confidence_threshold: Optional[Tensor] = None) box_mask = (boxes[:, :, 1] > confidence_threshold).unsqueeze(-1).expand_as(boxes) filtered_boxes = boxes * box_mask.float() else: - filtered_boxes = cast(list[Tensor], filtered_boxes) filtered_boxes = [] + filtered_boxes = cast(list[Tensor], filtered_boxes) for i in range(boxes.shape[0]): - box = boxes[i : i + 1][ - (boxes[i : i + 1, :, 1] > confidence_threshold).unsqueeze(-1).expand_as(boxes[i : i + 1]) + box = boxes[i][ + (boxes[i, :, 1] > confidence_threshold).unsqueeze(-1).expand_as(boxes[i]) ] - filtered_boxes.append(box.view(1, -1, boxes.shape[-1])) + filtered_boxes.append(box.view(-1, boxes.shape[-1])) return filtered_boxes diff --git a/tests/contrib/test_object_detector.py b/tests/contrib/test_object_detector.py index 00eee98b81..3a3b71b500 100644 --- a/tests/contrib/test_object_detector.py +++ b/tests/contrib/test_object_detector.py @@ -29,7 +29,7 @@ def test_smoke(self, device, dtype): assert pre_processor_out[0].shape[-2] == 32 assert len(detections) == batch_size for dets in detections: - assert dets.shape[1] == 6 + assert dets.shape[1] == 6, dets.shape assert torch.all(dets[:, 0].int() == dets[:, 0]) assert torch.all(dets[:, 1] >= 0.3) From 15150472cb61fecbb73385ef8c1107c22916048c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 13 Sep 2024 16:54:33 +0000 Subject: [PATCH 48/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/models/detector/utils.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/kornia/models/detector/utils.py b/kornia/models/detector/utils.py index ff33014df9..aa0b0fab66 100644 --- a/kornia/models/detector/utils.py +++ b/kornia/models/detector/utils.py @@ -36,9 +36,7 @@ def forward(self, boxes: Tensor, confidence_threshold: Optional[Tensor] = None) filtered_boxes = [] filtered_boxes = cast(list[Tensor], filtered_boxes) for i in range(boxes.shape[0]): - box = boxes[i][ - (boxes[i, :, 1] > confidence_threshold).unsqueeze(-1).expand_as(boxes[i]) - ] + box = boxes[i][(boxes[i, :, 1] > confidence_threshold).unsqueeze(-1).expand_as(boxes[i])] filtered_boxes.append(box.view(-1, boxes.shape[-1])) return filtered_boxes From b21d55543b2a5b93c4e9607b5cae88948c2713ec Mon Sep 17 00:00:00 2001 From: shijianjian Date: Fri, 13 Sep 2024 20:41:02 +0300 Subject: [PATCH 49/84] update --- kornia/models/detector/utils.py | 1 - kornia/onnx/utils.py | 4 ++-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/kornia/models/detector/utils.py b/kornia/models/detector/utils.py index aa0b0fab66..7895227528 100644 --- a/kornia/models/detector/utils.py +++ b/kornia/models/detector/utils.py @@ -34,7 +34,6 @@ def forward(self, boxes: Tensor, confidence_threshold: Optional[Tensor] = None) filtered_boxes = boxes * box_mask.float() else: filtered_boxes = [] - filtered_boxes = cast(list[Tensor], filtered_boxes) for i in range(boxes.shape[0]): box = boxes[i][(boxes[i, :, 1] > confidence_threshold).unsqueeze(-1).expand_as(boxes[i])] filtered_boxes.append(box.view(-1, boxes.shape[-1])) diff --git a/kornia/onnx/utils.py b/kornia/onnx/utils.py index 9a39dc84c6..98cd774f30 100644 --- a/kornia/onnx/utils.py +++ b/kornia/onnx/utils.py @@ -2,7 +2,7 @@ import os import pprint import urllib.request -from typing import Any, Optional +from typing import Any, Dict, List, Optional import requests @@ -103,7 +103,7 @@ def download( raise ValueError("URL must start with 'http:' or 'https:'") @staticmethod - def _fetch_repo_contents(folder: str) -> list[dict[str, Any]]: + def _fetch_repo_contents(folder: str) -> List[Dict[str, Any]]: """Fetches the contents of the Hugging Face repository using the Hugging Face API. Returns: From 56f81907ebbc6f8ad8719bc4d70c5476b22fdf29 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 13 Sep 2024 17:43:25 +0000 Subject: [PATCH 50/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/models/detector/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kornia/models/detector/utils.py b/kornia/models/detector/utils.py index 7895227528..fc2f9db34b 100644 --- a/kornia/models/detector/utils.py +++ b/kornia/models/detector/utils.py @@ -1,4 +1,4 @@ -from typing import Any, ClassVar, List, Optional, Tuple, Union, cast +from typing import Any, ClassVar, List, Optional, Tuple, Union from kornia.core import Module, ONNXExportMixin, Tensor, rand From 14c277c7481540ba942853747fe11d68d68fbeca Mon Sep 17 00:00:00 2001 From: shijianjian Date: Fri, 13 Sep 2024 20:59:12 +0300 Subject: [PATCH 51/84] update --- kornia/onnx/sequential.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/kornia/onnx/sequential.py b/kornia/onnx/sequential.py index a4f7b68c77..0fee754484 100644 --- a/kornia/onnx/sequential.py +++ b/kornia/onnx/sequential.py @@ -1,4 +1,4 @@ -from typing import Optional, Union +from typing import List, Optional, Tuple, Union from kornia.core.external import numpy as np from kornia.core.external import onnx @@ -30,9 +30,9 @@ class ONNXSequential: def __init__( self, *args: Union["onnx.ModelProto", str], # type:ignore - providers: Optional[list[str]] = None, + providers: Optional[List[str]] = None, session_options: Optional["ort.SessionOptions"] = None, # type:ignore - io_maps: Optional[list[tuple[str, str]]] = None, + io_maps: Optional[List[Tuple[str, str]]] = None, cache_dir: Optional[str] = None, ) -> None: self.onnx_loader = ONNXLoader(cache_dir) @@ -53,7 +53,7 @@ def _load_op(self, arg: Union["onnx.ModelProto", str]) -> "onnx.ModelProto": # return self.onnx_loader.load_model(arg) return arg - def _combine(self, io_maps: Optional[list[tuple[str, str]]] = None) -> "onnx.ModelProto": # type:ignore + def _combine(self, io_maps: Optional[List[Tuple[str, str]]] = None) -> "onnx.ModelProto": # type:ignore """Combine the provided ONNX models into a single ONNX graph. Optionally, map inputs and outputs between operators using the `io_map`. @@ -92,7 +92,7 @@ def export(self, file_path: str) -> None: def create_session( self, - providers: Optional[list[str]] = None, + providers: Optional[List[str]] = None, session_options: Optional["ort.SessionOptions"] = None, # type:ignore ) -> "ort.InferenceSession": # type:ignore """Create an optimized ONNXRuntime InferenceSession for the combined model. @@ -133,7 +133,7 @@ def get_session(self) -> "ort.InferenceSession": # type: ignore """ return self._session - def __call__(self, *inputs: "np.ndarray") -> list["np.ndarray"]: # type:ignore + def __call__(self, *inputs: "np.ndarray") -> List["np.ndarray"]: # type:ignore """Perform inference using the combined ONNX model. Args: From c110f3fbe44b44b3a4c161b7bbf42071e6310913 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Sat, 14 Sep 2024 14:27:17 +0300 Subject: [PATCH 52/84] update --- kornia/constants.py | 1 + kornia/contrib/models/sam/model.py | 19 ++++ kornia/core/external.py | 1 + kornia/models/detector/rtdetr.py | 2 +- kornia/models/tracker/__init__.py | 0 kornia/models/tracker/boxmot_tracker.py | 121 ++++++++++++++++++++++++ requirements/requirements-dev.txt | 2 +- tests/onnx/test_sequential.py | 65 +++++++++++++ tests/onnx/test_utils.py | 117 +++++++++++++++++++++++ 9 files changed, 326 insertions(+), 2 deletions(-) create mode 100644 kornia/models/tracker/__init__.py create mode 100644 kornia/models/tracker/boxmot_tracker.py create mode 100644 tests/onnx/test_sequential.py create mode 100644 tests/onnx/test_utils.py diff --git a/kornia/constants.py b/kornia/constants.py index c4bc6f3d84..9bf340d329 100644 --- a/kornia/constants.py +++ b/kornia/constants.py @@ -8,6 +8,7 @@ __all__ = ["pi", "DType", "Resample", "BorderType", "SamplePadding", "TKEnum"] +# NOTE: to remove later logging.basicConfig(level=logging.INFO) pi = torch.tensor(3.14159265358979323846) diff --git a/kornia/contrib/models/sam/model.py b/kornia/contrib/models/sam/model.py index f9cbbfa140..c528baf365 100644 --- a/kornia/contrib/models/sam/model.py +++ b/kornia/contrib/models/sam/model.py @@ -86,6 +86,25 @@ def __init__( self.prompt_encoder = prompt_encoder self.mask_decoder = mask_decoder + @staticmethod + def from_name(name: str) -> Sam: + """Build/load the SAM model based on it's name. + + Args: + name: The name of the SAM model. Valid names are: + - 'vit_b' + - 'vit_l' + - 'vit_h' + - 'mobile_sam' + + Returns: + The respective SAM model + """ + if name in ["vit_b", "vit_l", "vit_h", "mobile_sam"]: + return Sam.from_config(SamConfig(name)) + else: + raise ValueError(f"Invalid SAM model name: {name}") + @staticmethod def from_config(config: SamConfig) -> Sam: """Build/load the SAM model based on it's config. diff --git a/kornia/core/external.py b/kornia/core/external.py index 96de033f91..86b64e227d 100644 --- a/kornia/core/external.py +++ b/kornia/core/external.py @@ -96,3 +96,4 @@ def __dir__(self) -> List[str]: diffusers = LazyLoader("diffusers") onnx = LazyLoader("onnx") onnxruntime = LazyLoader("onnxruntime") +boxmot = LazyLoader("boxmot") diff --git a/kornia/models/detector/rtdetr.py b/kornia/models/detector/rtdetr.py index 818abc10ee..e4117127ed 100644 --- a/kornia/models/detector/rtdetr.py +++ b/kornia/models/detector/rtdetr.py @@ -86,7 +86,7 @@ def to_onnx( config: Optional[RTDETRConfig] = None, pretrained: bool = True, image_size: Optional[int] = 640, - confidence_threshold: float = 0.5, + confidence_threshold: Optional[float] = None, confidence_filtering: Optional[bool] = None, ) -> Tuple[str, ObjectDetector]: """Exports an RT-DETR object detection model to ONNX format. diff --git a/kornia/models/tracker/__init__.py b/kornia/models/tracker/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/kornia/models/tracker/boxmot_tracker.py b/kornia/models/tracker/boxmot_tracker.py new file mode 100644 index 0000000000..692abf6b69 --- /dev/null +++ b/kornia/models/tracker/boxmot_tracker.py @@ -0,0 +1,121 @@ +import os +from pathlib import Path +from typing import Union + +from kornia.core.external import boxmot +from kornia.core.external import numpy as np +from kornia.core import Tensor +from kornia.contrib.object_detection import ObjectDetector +from kornia.models.detector.rtdetr import RTDETRDetectorBuilder +from kornia.utils.image import tensor_to_image + + +class BoxMotTracker: + """BoxMotTracker is a module that wraps a detector and a tracker model. + + This module uses BoxMot library for tracking. + + Args: + detector: ObjectDetector: The detector model. + tracker_model_name: The name of the tracker model. Valid options are: + - "BoTSORT" + - "DeepOCSORT" + - "OCSORT" + - "HybridSORT" + - "ByteTrack" + - "StrongSORT" + - "ImprAssoc" + tracker_model_weights: Path to the model weights for ReID (Re-Identification). + device: Device on which to run the model (e.g., 'cpu' or 'cuda'). + fp16: Whether to use half-precision (fp16) for faster inference on compatible devices. + per_class: Whether to perform per-class tracking + track_high_thresh: High threshold for detection confidence. + Detections above this threshold are used in the first association round. + track_low_thresh: Low threshold for detection confidence. + Detections below this threshold are ignored. + new_track_thresh: Threshold for creating a new track. + Detections above this threshold will be considered as potential new tracks. + track_buffer: Number of frames to keep a track alive after it was last detected. + match_thresh: Threshold for the matching step in data association. + proximity_thresh: Threshold for IoU (Intersection over Union) distance in first-round association. + appearance_thresh: Threshold for appearance embedding distance in the ReID module. + cmc_method: Method for correcting camera motion. Options include "sof" (simple optical flow). + frame_rate: Frame rate of the video being processed. Used to scale the track buffer size. + fuse_first_associate: Whether to fuse appearance and motion information during the first association step. + with_reid: Whether to use ReID (Re-Identification) features for association. + """ + + def __init__( + self, detector: Union[ObjectDetector, str] = "rtdetr_r18vd", tracker_model_name: str = "BoTSORT", + tracker_model_weights: str = "osnet_x0_25_msmt17.pt", device: str = "cpu", fp16: bool = False, **kwargs + ) -> None: + super().__init__() + if isinstance(detector, str): + if detector.startswith("rtdetr"): + detector = RTDETRDetectorBuilder.build(model_name=detector) + else: + raise ValueError(f"Detector `{detector}` not available. You may pass an ObjectDetector instance instead.") + self.detector = detector + os.makedirs(".kornia_hub/models/boxmot", exist_ok=True) + self.tracker = getattr(boxmot, tracker_model_name)( + model_weights=Path(os.path.join(".kornia_hub/models/boxmot", tracker_model_weights)), + device=device, + fp16=fp16, + **kwargs + ) + + def update(self, image: Tensor) -> None: + """Update the tracker with a new image. + + Args: + image: The input image. + """ + + if not (image.ndim == 4 and image.shape[0] == 1) and not image.ndim == 3: + raise ValueError( + f"Input tensor must be of shape (1, 3, H, W) or (3, H, W). Got {image.shape}" + ) + + if image.ndim == 3: + image = image.unsqueeze(0) + + detections: Union[Tensor, list[Tensor]] = self.detector(image) + + detections = detections[0].cpu().numpy() # Batch size is 1 + + detections = np.array( # type: ignore + [ + detections[:, 2], + detections[:, 3], + detections[:, 2] + detections[:, 4], + detections[:, 3] + detections[:, 5], + detections[:, 1], + detections[:, 0], + ] + ).T + + if detections.shape[0] == 0: + # empty N X (x, y, x, y, conf, cls) + detections = np.empty((0, 6)) # type: ignore + + frame_raw = (tensor_to_image(image) * 255).astype(np.uint8) # type: ignore + + return self.tracker.update( # type: ignore + detections, frame_raw + ) # --> M X (x, y, x, y, id, conf, cls, ind) + + def visualize(self, image: Tensor, show_trajectories: bool =True) -> np.ndarray: # type: ignore + """Visualize the results of the tracker. + + Args: + image: The input image. + show_trajectories: Whether to show the trajectories. + + Returns: + The image with the results of the tracker. + """ + frame_raw = (tensor_to_image(image) * 255).astype(np.uint8) # type: ignore + self.update(image) + self.tracker.plot_results(frame_raw, show_trajectories=show_trajectories) # type: ignore + + return frame_raw diff --git a/requirements/requirements-dev.txt b/requirements/requirements-dev.txt index 21866c5bfb..632dcc17da 100644 --- a/requirements/requirements-dev.txt +++ b/requirements/requirements-dev.txt @@ -2,7 +2,7 @@ accelerate coverage diffusers mypy -numpy<3 +numpy<2 onnx onnxruntime pillow diff --git a/tests/onnx/test_sequential.py b/tests/onnx/test_sequential.py new file mode 100644 index 0000000000..fd8bcb6c9d --- /dev/null +++ b/tests/onnx/test_sequential.py @@ -0,0 +1,65 @@ +import pytest +from unittest.mock import MagicMock, patch +import numpy as np +import onnx +import onnxruntime as ort +from onnx.helper import make_tensor_value_info, make_graph, make_model, make_node + +from kornia.onnx.sequential import ONNXSequential + + +class TestONNXSequential: + @pytest.fixture + def mock_model_proto(self): + # Create a minimal ONNX model with an input and output + input_info = make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1, 2]) + output_info = make_tensor_value_info("output", onnx.TensorProto.FLOAT, [1, 2]) + node = make_node("Identity", ["input"], ["output"]) + graph = make_graph([node], "test_graph", [input_info], [output_info]) + model = make_model(graph) + return model + + @pytest.fixture + def onnx_sequential(self, mock_model_proto): + # Return an ONNXSequential instance with mocked models + return ONNXSequential(mock_model_proto, mock_model_proto) + + def test_load_op_from_proto(self, mock_model_proto, onnx_sequential): + # Test loading a model from an ONNX ModelProto object + model = onnx_sequential._load_op(mock_model_proto) + assert model == mock_model_proto + + @patch("onnx.compose.merge_models") + def test_combine_models(self, mock_merge_models, mock_model_proto): + # Create a small ONNX model as the return value of merge_models + input_info = make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1, 2]) + output_info = make_tensor_value_info("output", onnx.TensorProto.FLOAT, [1, 2]) + node = make_node("Identity", ["input"], ["output"]) + graph = make_graph([node], "combined_graph", [input_info], [output_info]) + combined_model = make_model(graph) + + mock_merge_models.return_value = combined_model + + # Test combining multiple ONNX models with io_maps + onnx_sequential = ONNXSequential(mock_model_proto, mock_model_proto) + combined_op = onnx_sequential._combine([("output1", "input2")]) + + assert isinstance(combined_op, onnx.ModelProto) + + @patch("onnx.save") + def test_export_combined_model(self, mock_save, onnx_sequential): + # Test exporting the combined ONNX model + onnx_sequential.export("exported_model.onnx") + mock_save.assert_called_once_with(onnx_sequential._combined_op, "exported_model.onnx") + + @patch("onnxruntime.InferenceSession") + def test_create_session(self, mock_inference_session, onnx_sequential): + # Test creating an ONNXRuntime session + session = onnx_sequential.create_session() + assert session == mock_inference_session() + + def test_set_get_session(self, onnx_sequential): + # Test setting and getting a custom session + mock_session = MagicMock(spec=ort.InferenceSession) + onnx_sequential.set_session(mock_session) + assert onnx_sequential.get_session() == mock_session diff --git a/tests/onnx/test_utils.py b/tests/onnx/test_utils.py new file mode 100644 index 0000000000..42bbbaa8e7 --- /dev/null +++ b/tests/onnx/test_utils.py @@ -0,0 +1,117 @@ +import os +import pytest +from unittest import mock +from onnx import ModelProto # Assuming `onnx` is installed and ModelProto is part of the library +from kornia.onnx.utils import ONNXLoader +import urllib + + +class TestONNXLoader: + @pytest.fixture + def loader(self): + return ONNXLoader(cache_dir=".test_cache") + + def test_get_file_path_with_custom_cache_dir(self, loader): + model_name = "operators/some_model" + expected_path = ".test_cache/operators/some_model.onnx" + assert loader._get_file_path(model_name, loader.cache_dir) == expected_path + + def test_get_file_path_with_default_cache_dir(self): + loader = ONNXLoader() + model_name = "operators/some_model" + expected_path = ".kornia_hub/onnx_models/operators/some_model.onnx" + assert loader._get_file_path(model_name, None) == expected_path + + @mock.patch("onnx.load") + @mock.patch("os.path.exists") + def test_load_model_local(self, mock_exists, mock_onnx_load, loader): + model_name = "local_model.onnx" + mock_exists.return_value = True + + # Simulate onnx.load returning a dummy ModelProto + mock_model = mock.Mock(spec=ModelProto) + mock_onnx_load.return_value = mock_model + + model = loader.load_model(model_name) + assert model == mock_model + mock_onnx_load.assert_called_once_with(model_name) + + @mock.patch("urllib.request.urlretrieve") + @mock.patch("os.path.exists") + def test_load_model_download(self, mock_exists, mock_urlretrieve, loader): + model_name = "hf://operators/some_model" + mock_exists.return_value = False + mock_urlretrieve.return_value = None # Simulating successful download + + with mock.patch("onnx.load") as mock_onnx_load: + mock_model = mock.Mock(spec=ModelProto) + mock_onnx_load.return_value = mock_model + + model = loader.load_model(model_name) + assert model == mock_model + mock_urlretrieve.assert_called_once_with( + "https://huggingface.co/kornia/ONNX_models/resolve/main/operators/some_model.onnx", + ".test_cache/operators/some_model.onnx" + ) + + def test_load_model_file_not_found(self, loader): + model_name = "non_existent_model.onnx" + + with pytest.raises(ValueError, match=f"File {model_name} not found"): + loader.load_model(model_name) + + @mock.patch("urllib.request.urlretrieve") + @mock.patch("os.makedirs") + def test_download_success(self, mock_makedirs, mock_urlretrieve, loader): + url = "https://huggingface.co/some_model.onnx" + file_path = ".test_cache/some_model.onnx" + + loader.download(url, file_path) + + mock_makedirs.assert_called_once_with(os.path.dirname(file_path), exist_ok=True) + mock_urlretrieve.assert_called_once_with(url, file_path) + + @mock.patch("urllib.request.urlretrieve", side_effect=urllib.error.HTTPError(url=None, code=404, msg="Not Found", hdrs=None, fp=None)) + def test_download_failure(self, mock_urlretrieve, loader): + url = "https://huggingface.co/non_existent_model.onnx" + file_path = ".test_cache/non_existent_model.onnx" + + with pytest.raises(ValueError, match="Error in resolving"): + loader.download(url, file_path) + + @mock.patch("requests.get") + def test_fetch_repo_contents_success(self, mock_get): + mock_response = mock.Mock() + mock_response.status_code = 200 + mock_response.json.return_value = [{"path": "operators/model.onnx"}] + mock_get.return_value = mock_response + + contents = ONNXLoader._fetch_repo_contents("operators") + assert contents == [{"path": "operators/model.onnx"}] + + @mock.patch("requests.get") + def test_fetch_repo_contents_failure(self, mock_get): + mock_response = mock.Mock() + mock_response.status_code = 404 + mock_get.return_value = mock_response + + with pytest.raises(ValueError, match="Failed to fetch repository contents"): + ONNXLoader._fetch_repo_contents("operators") + + @mock.patch("kornia.onnx.utils.ONNXLoader._fetch_repo_contents") + def test_list_operators(self, mock_fetch_repo_contents, capsys): + mock_fetch_repo_contents.return_value = [{"path": "operators/some_model.onnx"}] + + ONNXLoader.list_operators() + + captured = capsys.readouterr() + assert "operators/some_model.onnx" in captured.out + + @mock.patch("kornia.onnx.utils.ONNXLoader._fetch_repo_contents") + def test_list_models(self, mock_fetch_repo_contents, capsys): + mock_fetch_repo_contents.return_value = [{"path": "models/some_model.onnx"}] + + ONNXLoader.list_models() + + captured = capsys.readouterr() + assert "models/some_model.onnx" in captured.out From 93c4968380a121067ab58fd45196acd542bf63aa Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 14 Sep 2024 11:27:38 +0000 Subject: [PATCH 53/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/models/tracker/boxmot_tracker.py | 25 +++++++++++++++---------- tests/onnx/test_sequential.py | 6 +++--- tests/onnx/test_utils.py | 13 +++++++++---- 3 files changed, 27 insertions(+), 17 deletions(-) diff --git a/kornia/models/tracker/boxmot_tracker.py b/kornia/models/tracker/boxmot_tracker.py index 692abf6b69..f591709538 100644 --- a/kornia/models/tracker/boxmot_tracker.py +++ b/kornia/models/tracker/boxmot_tracker.py @@ -2,10 +2,10 @@ from pathlib import Path from typing import Union +from kornia.contrib.object_detection import ObjectDetector +from kornia.core import Tensor from kornia.core.external import boxmot from kornia.core.external import numpy as np -from kornia.core import Tensor -from kornia.contrib.object_detection import ObjectDetector from kornia.models.detector.rtdetr import RTDETRDetectorBuilder from kornia.utils.image import tensor_to_image @@ -46,22 +46,29 @@ class BoxMotTracker: """ def __init__( - self, detector: Union[ObjectDetector, str] = "rtdetr_r18vd", tracker_model_name: str = "BoTSORT", - tracker_model_weights: str = "osnet_x0_25_msmt17.pt", device: str = "cpu", fp16: bool = False, **kwargs + self, + detector: Union[ObjectDetector, str] = "rtdetr_r18vd", + tracker_model_name: str = "BoTSORT", + tracker_model_weights: str = "osnet_x0_25_msmt17.pt", + device: str = "cpu", + fp16: bool = False, + **kwargs, ) -> None: super().__init__() if isinstance(detector, str): if detector.startswith("rtdetr"): detector = RTDETRDetectorBuilder.build(model_name=detector) else: - raise ValueError(f"Detector `{detector}` not available. You may pass an ObjectDetector instance instead.") + raise ValueError( + f"Detector `{detector}` not available. You may pass an ObjectDetector instance instead." + ) self.detector = detector os.makedirs(".kornia_hub/models/boxmot", exist_ok=True) self.tracker = getattr(boxmot, tracker_model_name)( model_weights=Path(os.path.join(".kornia_hub/models/boxmot", tracker_model_weights)), device=device, fp16=fp16, - **kwargs + **kwargs, ) def update(self, image: Tensor) -> None: @@ -72,9 +79,7 @@ def update(self, image: Tensor) -> None: """ if not (image.ndim == 4 and image.shape[0] == 1) and not image.ndim == 3: - raise ValueError( - f"Input tensor must be of shape (1, 3, H, W) or (3, H, W). Got {image.shape}" - ) + raise ValueError(f"Input tensor must be of shape (1, 3, H, W) or (3, H, W). Got {image.shape}") if image.ndim == 3: image = image.unsqueeze(0) @@ -104,7 +109,7 @@ def update(self, image: Tensor) -> None: detections, frame_raw ) # --> M X (x, y, x, y, id, conf, cls, ind) - def visualize(self, image: Tensor, show_trajectories: bool =True) -> np.ndarray: # type: ignore + def visualize(self, image: Tensor, show_trajectories: bool = True) -> np.ndarray: # type: ignore """Visualize the results of the tracker. Args: diff --git a/tests/onnx/test_sequential.py b/tests/onnx/test_sequential.py index fd8bcb6c9d..630d7fb70e 100644 --- a/tests/onnx/test_sequential.py +++ b/tests/onnx/test_sequential.py @@ -1,9 +1,9 @@ -import pytest from unittest.mock import MagicMock, patch -import numpy as np + import onnx import onnxruntime as ort -from onnx.helper import make_tensor_value_info, make_graph, make_model, make_node +import pytest +from onnx.helper import make_graph, make_model, make_node, make_tensor_value_info from kornia.onnx.sequential import ONNXSequential diff --git a/tests/onnx/test_utils.py b/tests/onnx/test_utils.py index 42bbbaa8e7..43e70f2799 100644 --- a/tests/onnx/test_utils.py +++ b/tests/onnx/test_utils.py @@ -1,9 +1,11 @@ import os -import pytest +import urllib from unittest import mock + +import pytest from onnx import ModelProto # Assuming `onnx` is installed and ModelProto is part of the library + from kornia.onnx.utils import ONNXLoader -import urllib class TestONNXLoader: @@ -51,7 +53,7 @@ def test_load_model_download(self, mock_exists, mock_urlretrieve, loader): assert model == mock_model mock_urlretrieve.assert_called_once_with( "https://huggingface.co/kornia/ONNX_models/resolve/main/operators/some_model.onnx", - ".test_cache/operators/some_model.onnx" + ".test_cache/operators/some_model.onnx", ) def test_load_model_file_not_found(self, loader): @@ -71,7 +73,10 @@ def test_download_success(self, mock_makedirs, mock_urlretrieve, loader): mock_makedirs.assert_called_once_with(os.path.dirname(file_path), exist_ok=True) mock_urlretrieve.assert_called_once_with(url, file_path) - @mock.patch("urllib.request.urlretrieve", side_effect=urllib.error.HTTPError(url=None, code=404, msg="Not Found", hdrs=None, fp=None)) + @mock.patch( + "urllib.request.urlretrieve", + side_effect=urllib.error.HTTPError(url=None, code=404, msg="Not Found", hdrs=None, fp=None), + ) def test_download_failure(self, mock_urlretrieve, loader): url = "https://huggingface.co/non_existent_model.onnx" file_path = ".test_cache/non_existent_model.onnx" From fc5ecb4280c014273b9e13d92e95a1f2b1ca682a Mon Sep 17 00:00:00 2001 From: shijianjian Date: Sat, 14 Sep 2024 17:03:25 +0300 Subject: [PATCH 54/84] update --- kornia/models/tracker/boxmot_tracker.py | 23 +++++++++++------------ requirements/requirements-dev.txt | 1 + 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/kornia/models/tracker/boxmot_tracker.py b/kornia/models/tracker/boxmot_tracker.py index f591709538..d0fdcfb3bb 100644 --- a/kornia/models/tracker/boxmot_tracker.py +++ b/kornia/models/tracker/boxmot_tracker.py @@ -90,24 +90,23 @@ def update(self, image: Tensor) -> None: detections = np.array( # type: ignore [ - detections[:, 2], - detections[:, 3], - detections[:, 2] + detections[:, 4], - detections[:, 3] + detections[:, 5], - detections[:, 1], - detections[:, 0], + detections[:, 2], # type: ignore + detections[:, 3], # type: ignore + detections[:, 2] + detections[:, 4], # type: ignore + detections[:, 3] + detections[:, 5], # type: ignore + detections[:, 1], # type: ignore + detections[:, 0], # type: ignore ] ).T - if detections.shape[0] == 0: + if detections.shape[0] == 0: # type: ignore # empty N X (x, y, x, y, conf, cls) detections = np.empty((0, 6)) # type: ignore frame_raw = (tensor_to_image(image) * 255).astype(np.uint8) # type: ignore - return self.tracker.update( # type: ignore - detections, frame_raw - ) # --> M X (x, y, x, y, id, conf, cls, ind) + # --> M X (x, y, x, y, id, conf, cls, ind) + return self.tracker.update(detections, frame_raw) # type: ignore def visualize(self, image: Tensor, show_trajectories: bool = True) -> np.ndarray: # type: ignore """Visualize the results of the tracker. @@ -119,8 +118,8 @@ def visualize(self, image: Tensor, show_trajectories: bool = True) -> np.ndarray Returns: The image with the results of the tracker. """ - frame_raw = (tensor_to_image(image) * 255).astype(np.uint8) # type: ignore + frame_raw = (tensor_to_image(image) * 255).astype(np.uint8) self.update(image) - self.tracker.plot_results(frame_raw, show_trajectories=show_trajectories) # type: ignore + self.tracker.plot_results(frame_raw, show_trajectories=show_trajectories) return frame_raw diff --git a/requirements/requirements-dev.txt b/requirements/requirements-dev.txt index 632dcc17da..f52cc1aa94 100644 --- a/requirements/requirements-dev.txt +++ b/requirements/requirements-dev.txt @@ -1,4 +1,5 @@ accelerate +boxmot coverage diffusers mypy From 898364c889a9074243d68de60d748f35d26d90bb Mon Sep 17 00:00:00 2001 From: shijianjian Date: Sat, 14 Sep 2024 18:39:08 +0300 Subject: [PATCH 55/84] update --- kornia/core/external.py | 1 + kornia/models/segmentor/__init__.py | 0 .../models/segmentor/segmentation_models.py | 74 +++++++++++++++++++ kornia/models/tracker/__init__.py | 1 + kornia/models/tracker/boxmot_tracker.py | 2 + 5 files changed, 78 insertions(+) create mode 100644 kornia/models/segmentor/__init__.py create mode 100644 kornia/models/segmentor/segmentation_models.py diff --git a/kornia/core/external.py b/kornia/core/external.py index 86b64e227d..de1544f7ca 100644 --- a/kornia/core/external.py +++ b/kornia/core/external.py @@ -97,3 +97,4 @@ def __dir__(self) -> List[str]: onnx = LazyLoader("onnx") onnxruntime = LazyLoader("onnxruntime") boxmot = LazyLoader("boxmot") +segmentation_models_pytorch = LazyLoader("segmentation_models_pytorch") diff --git a/kornia/models/segmentor/__init__.py b/kornia/models/segmentor/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/kornia/models/segmentor/segmentation_models.py b/kornia/models/segmentor/segmentation_models.py new file mode 100644 index 0000000000..40bc6c9b54 --- /dev/null +++ b/kornia/models/segmentor/segmentation_models.py @@ -0,0 +1,74 @@ +from typing import Optional + +import kornia +from kornia.core import Module, Tensor, tensor +from kornia.core.external import segmentation_models_pytorch as smp + + +class SegmentationModels(Module): + """SegmentationModel is a module that wraps a segmentation model. + + This module uses SegmentationModel library for segmentation. + + Args: + model_name: Name of the model to use. Valid options are: + "Unet", "UnetPlusPlus", "MAnet", "LinkNet", "FPN", "PSPNet", "PAN", "DeepLabV3", "DeepLabV3Plus". + encoder_name: Name of the encoder to use. + encoder_depth: Depth of the encoder. + encoder_weights: Weights of the encoder. + decoder_channels: Number of channels in the decoder. + in_channels: Number of channels in the input. + classes: Number of classes to predict. + **kwargs: Additional arguments to pass to the model. Detailed arguments can be found at: + https://github.com/qubvel-org/segmentation_models.pytorch/tree/main/segmentation_models_pytorch/decoders + + Note: + Only encoder weights are available. + Pretrained weights for the whole model are not available. + """ + def __init__( + self, + model_name: str = "Unet", + encoder_name: str = "resnet34", + encoder_weights: Optional[str] = "imagenet", + in_channels: int = 3, + classes: int = 1, + **kwargs + ) -> None: + super().__init__() + self.preproc_params = smp.encoders.get_preprocessing_params(encoder_name) # type: ignore + self.segmentation_model = getattr(smp, model_name)( + encoder_name=encoder_name, + encoder_weights=encoder_weights, + in_channels=in_channels, + classes=classes, + **kwargs + ) + + def preprocessing(self, input: Tensor) -> Tensor: + if self.preproc_params["input_space"] == "RGB": + pass + elif self.preproc_params["input_space"] == "BGR": + input = kornia.color.rgb_to_bgr(input) + else: + raise ValueError(f"Unsupported input space: {self.preproc_params['input_space']}") + + if self.preproc_params["input_range"] is not None: + if input.max() > 1 and self.preproc_params["input_range"][1] == 1: + input = input / 255.0 + + if self.preproc_params["mean"] is None: + mean = tensor(self.preproc_params["mean"]).to(input.device) + else: + mean = tensor(self.preproc_params["mean"]).to(input.device) + + if self.preproc_params["std"] is None: + std = tensor(self.preproc_params["std"]).to(input.device) + else: + std = tensor(self.preproc_params["std"]).to(input.device) + + return kornia.enhance.normalize(input, mean, std) + + def forward(self, input: Tensor) -> Tensor: + input = self.preprocessing(input) + return self.segmentation_model(input) diff --git a/kornia/models/tracker/__init__.py b/kornia/models/tracker/__init__.py index e69de29bb2..a68c9dce90 100644 --- a/kornia/models/tracker/__init__.py +++ b/kornia/models/tracker/__init__.py @@ -0,0 +1 @@ +from .boxmot_tracker import * diff --git a/kornia/models/tracker/boxmot_tracker.py b/kornia/models/tracker/boxmot_tracker.py index d0fdcfb3bb..97a09bf87b 100644 --- a/kornia/models/tracker/boxmot_tracker.py +++ b/kornia/models/tracker/boxmot_tracker.py @@ -9,6 +9,8 @@ from kornia.models.detector.rtdetr import RTDETRDetectorBuilder from kornia.utils.image import tensor_to_image +__all__ = ["BoxMotTracker"] + class BoxMotTracker: """BoxMotTracker is a module that wraps a detector and a tracker model. From 35af5f1663254c5ee2bda95d65ca202eb0009d16 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 14 Sep 2024 15:39:29 +0000 Subject: [PATCH 56/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/models/segmentor/segmentation_models.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/kornia/models/segmentor/segmentation_models.py b/kornia/models/segmentor/segmentation_models.py index 40bc6c9b54..8e619bc97f 100644 --- a/kornia/models/segmentor/segmentation_models.py +++ b/kornia/models/segmentor/segmentation_models.py @@ -21,11 +21,12 @@ class SegmentationModels(Module): classes: Number of classes to predict. **kwargs: Additional arguments to pass to the model. Detailed arguments can be found at: https://github.com/qubvel-org/segmentation_models.pytorch/tree/main/segmentation_models_pytorch/decoders - + Note: Only encoder weights are available. Pretrained weights for the whole model are not available. """ + def __init__( self, model_name: str = "Unet", @@ -33,7 +34,7 @@ def __init__( encoder_weights: Optional[str] = "imagenet", in_channels: int = 3, classes: int = 1, - **kwargs + **kwargs, ) -> None: super().__init__() self.preproc_params = smp.encoders.get_preprocessing_params(encoder_name) # type: ignore @@ -42,7 +43,7 @@ def __init__( encoder_weights=encoder_weights, in_channels=in_channels, classes=classes, - **kwargs + **kwargs, ) def preprocessing(self, input: Tensor) -> Tensor: @@ -52,7 +53,7 @@ def preprocessing(self, input: Tensor) -> Tensor: input = kornia.color.rgb_to_bgr(input) else: raise ValueError(f"Unsupported input space: {self.preproc_params['input_space']}") - + if self.preproc_params["input_range"] is not None: if input.max() > 1 and self.preproc_params["input_range"][1] == 1: input = input / 255.0 @@ -61,7 +62,7 @@ def preprocessing(self, input: Tensor) -> Tensor: mean = tensor(self.preproc_params["mean"]).to(input.device) else: mean = tensor(self.preproc_params["mean"]).to(input.device) - + if self.preproc_params["std"] is None: std = tensor(self.preproc_params["std"]).to(input.device) else: From ba18fcd2bb3d4ed087aa09c132e6c90334b56fa8 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Sun, 15 Sep 2024 00:25:55 +0300 Subject: [PATCH 57/84] update --- kornia/enhance/normalize.py | 5 ++- .../models/segmentor/segmentation_models.py | 41 +++++++++++-------- kornia/onnx/sequential.py | 2 +- 3 files changed, 29 insertions(+), 19 deletions(-) diff --git a/kornia/enhance/normalize.py b/kornia/enhance/normalize.py index 3f4fd3f447..a291fa58cf 100644 --- a/kornia/enhance/normalize.py +++ b/kornia/enhance/normalize.py @@ -105,7 +105,10 @@ def normalize(data: Tensor, mean: Tensor, std: Tensor) -> Tensor: if not isinstance(mean, Tensor) or not isinstance(std, Tensor): raise ValueError("Only tensor is accepted when converting to ONNX.") if mean.shape[0] != 1 or std.shape[0] != 1: - raise ValueError("Batch dimension must be one for broadcasting when converting to ONNX.") + raise ValueError( + "Batch dimension must be one for broadcasting when converting to ONNX." + f"Try changing mean shape and std shape from ({mean.shape}, {std.shape}) to (1, C) or (1, C, 1, 1)." + ) else: if isinstance(mean, float): mean = torch.tensor([mean] * shape[1], device=data.device, dtype=data.dtype) diff --git a/kornia/models/segmentor/segmentation_models.py b/kornia/models/segmentor/segmentation_models.py index 8e619bc97f..09093aac11 100644 --- a/kornia/models/segmentor/segmentation_models.py +++ b/kornia/models/segmentor/segmentation_models.py @@ -1,11 +1,12 @@ from typing import Optional import kornia -from kornia.core import Module, Tensor, tensor +from kornia.core import Module, Tensor, tensor, zeros_like, ones_like +from kornia.core.module import ONNXExportMixin from kornia.core.external import segmentation_models_pytorch as smp -class SegmentationModels(Module): +class SegmentationModels(Module, ONNXExportMixin): """SegmentationModel is a module that wraps a segmentation model. This module uses SegmentationModel library for segmentation. @@ -27,6 +28,9 @@ class SegmentationModels(Module): Pretrained weights for the whole model are not available. """ + ONNX_DEFAULT_INPUTSHAPE = (-1, 3, -1, -1) + ONNX_DEFAULT_OUTPUTSHAPE = (-1, -1, -1, -1) + def __init__( self, model_name: str = "Unet", @@ -47,27 +51,30 @@ def __init__( ) def preprocessing(self, input: Tensor) -> Tensor: - if self.preproc_params["input_space"] == "RGB": + # Ensure the color space transformation is ONNX-friendly + input_space = self.preproc_params["input_space"] + input = kornia.color.rgb_to_bgr(input) if input_space == "BGR" else input # Assume input is already RGB if not BGR + + # Normalize input range if needed + input_range = self.preproc_params["input_range"] + if input_range[1] == 255: + input = input * 255.0 + elif input_range[1] == 1: pass - elif self.preproc_params["input_space"] == "BGR": - input = kornia.color.rgb_to_bgr(input) else: - raise ValueError(f"Unsupported input space: {self.preproc_params['input_space']}") + raise ValueError(f"Unsupported input range: {input_range}") - if self.preproc_params["input_range"] is not None: - if input.max() > 1 and self.preproc_params["input_range"][1] == 1: - input = input / 255.0 - - if self.preproc_params["mean"] is None: - mean = tensor(self.preproc_params["mean"]).to(input.device) + # Handle mean and std normalization + if self.preproc_params["mean"] is not None: + mean = tensor([self.preproc_params["mean"]]).to(input.device) else: - mean = tensor(self.preproc_params["mean"]).to(input.device) + mean = zeros_like(input) - if self.preproc_params["std"] is None: - std = tensor(self.preproc_params["std"]).to(input.device) + if self.preproc_params["std"] is not None: + std = tensor([self.preproc_params["std"]]).to(input.device) else: - std = tensor(self.preproc_params["std"]).to(input.device) - + std = ones_like(input) + return kornia.enhance.normalize(input, mean, std) def forward(self, input: Tensor) -> Tensor: diff --git a/kornia/onnx/sequential.py b/kornia/onnx/sequential.py index 0fee754484..f992d57eae 100644 --- a/kornia/onnx/sequential.py +++ b/kornia/onnx/sequential.py @@ -112,7 +112,7 @@ def create_session( session = ort.InferenceSession( # type:ignore self._combined_op.SerializeToString(), sess_options=sess_options, - providers=providers or ["CPUExecutionProvider"], + providers=providers or ['CUDAExecutionProvider', 'CPUExecutionProvider'], ) return session From 1c508f53a45287a1b50234f9c391050f17aadf5a Mon Sep 17 00:00:00 2001 From: shijianjian Date: Sun, 15 Sep 2024 00:38:48 +0300 Subject: [PATCH 58/84] update --- kornia/onnx/sequential.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/kornia/onnx/sequential.py b/kornia/onnx/sequential.py index f992d57eae..d4dea8c26e 100644 --- a/kornia/onnx/sequential.py +++ b/kornia/onnx/sequential.py @@ -112,7 +112,7 @@ def create_session( session = ort.InferenceSession( # type:ignore self._combined_op.SerializeToString(), sess_options=sess_options, - providers=providers or ['CUDAExecutionProvider', 'CPUExecutionProvider'], + providers=providers or ['CPUExecutionProvider'], ) return session @@ -133,6 +133,14 @@ def get_session(self) -> "ort.InferenceSession": # type: ignore """ return self._session + def as_cpu(self) -> None: + """Set the session to run on CPU.""" + self._session = self.create_session(providers=['CPUExecutionProvider']) + + def as_cuda(self, device_id: int = 0) -> None: + """Set the session to run on CUDA.""" + self._session = self.create_session(providers=[('CUDAExecutionProvider', {'device_id': device_id})]) + def __call__(self, *inputs: "np.ndarray") -> List["np.ndarray"]: # type:ignore """Perform inference using the combined ONNX model. From e0cabded8eacedff3f2c27a9b3c415a0e69c028f Mon Sep 17 00:00:00 2001 From: shijianjian Date: Sun, 15 Sep 2024 08:51:52 +0300 Subject: [PATCH 59/84] update --- kornia/contrib/object_detection.py | 244 ++---------------- kornia/filters/dexined.py | 17 +- kornia/models/__init__.py | 2 + kornia/models/detector/__init__.py | 1 + kornia/models/detector/base.py | 204 +++++++++++++++ kornia/models/detector/rtdetr.py | 3 +- kornia/models/detector/utils.py | 7 + kornia/models/edge_detector/__init__.py | 2 + kornia/models/edge_detector/base.py | 67 +++++ kornia/models/edge_detector/dexined.py | 55 ++++ kornia/models/segmentor/__init__.py | 1 + kornia/models/segmentor/base.py | 63 +++++ .../models/segmentor/segmentation_models.py | 2 +- kornia/models/tracker/boxmot_tracker.py | 2 +- kornia/models/utils.py | 78 ++++++ kornia/onnx/sequential.py | 4 +- 16 files changed, 515 insertions(+), 237 deletions(-) create mode 100644 kornia/models/detector/base.py create mode 100644 kornia/models/edge_detector/__init__.py create mode 100644 kornia/models/edge_detector/base.py create mode 100644 kornia/models/edge_detector/dexined.py create mode 100644 kornia/models/segmentor/base.py create mode 100644 kornia/models/utils.py diff --git a/kornia/contrib/object_detection.py b/kornia/contrib/object_detection.py index 4a4b965dc2..d6a7de5ad7 100644 --- a/kornia/contrib/object_detection.py +++ b/kornia/contrib/object_detection.py @@ -1,21 +1,13 @@ -from __future__ import annotations +import warnings -import datetime -import logging -import os -from dataclasses import dataclass -from enum import Enum -from typing import Optional, Union - -import torch - -from kornia.core import Module, Tensor, concatenate -from kornia.core.check import KORNIA_CHECK_SHAPE -from kornia.core.external import PILImage as Image -from kornia.core.external import numpy as np -from kornia.geometry.transform import resize -from kornia.io import write_image -from kornia.utils.draw import draw_rectangle +from kornia.models.utils import ResizePreProcessor as ResizePreProcessorBase +from kornia.models.detector.base import ( + BoundingBox as BoundingBoxBase, + BoundingBoxDataFormat, + results_from_detections as results_from_detections_base, + ObjectDetector as ObjectDetectorBase, + ObjectDetectorResult as ObjectDetectorResultBase, +) __all__ = [ "BoundingBoxDataFormat", @@ -26,218 +18,22 @@ "ObjectDetectorResult", ] -logger = logging.getLogger(__name__) - - -class BoundingBoxDataFormat(Enum): - """Enum class that maps bounding box data format.""" - - XYWH = 0 - XYXY = 1 - CXCYWH = 2 - CENTER_XYWH = 2 - - -# NOTE: probably we should use a more generic name like BoundingBox2D -# and add a BoundingBox3D class for 3D bounding boxes. Also for serialization -# we should have an explicit class for each format to make it more production ready -# specially to serialize to protobuf and not saturate at a high rates. - - -@dataclass(frozen=True) -class BoundingBox: - """Bounding box data class. - - Useful for representing bounding boxes in different formats for object detection. - - Args: - data: tuple of bounding box data. The length of the tuple depends on the data format. - data_format: bounding box data format. - """ - - data: tuple[float, float, float, float] - data_format: BoundingBoxDataFormat - - -@dataclass(frozen=True) -class ObjectDetectorResult: - """Object detection result. - - Args: - class_id: class id of the detected object. - confidence: confidence score of the detected object. - bbox: bounding box of the detected object in xywh format. - """ - - class_id: int - confidence: float - bbox: BoundingBox - - -def results_from_detections(detections: Tensor, format: str | BoundingBoxDataFormat) -> list[ObjectDetectorResult]: - """Convert a detection tensor to a list of :py:class:`ObjectDetectorResult`. - - Args: - detections: tensor with shape :math:`(D, 6)`, where :math:`D` is the number of detections in the given image, - :math:`6` represents class id, score, and `xywh` bounding box. - - Returns: - list of :py:class:`ObjectDetectorResult`. - """ - KORNIA_CHECK_SHAPE(detections, ["D", "6"]) - - if isinstance(format, str): - format = BoundingBoxDataFormat[format.upper()] - - results: list[ObjectDetectorResult] = [] - for det in detections: - det = det.squeeze().tolist() - if len(det) != 6: - continue - results.append( - ObjectDetectorResult( - class_id=int(det[0]), - confidence=det[1], - bbox=BoundingBox(data=(det[2], det[3], det[4], det[5]), data_format=format), - ) - ) - return results - - -class ResizePreProcessor(Module): - """This module resizes a list of image tensors to the given size. - - Additionally, also returns the original image sizes for further post-processing. - """ - - def __init__(self, size: tuple[int, int], interpolation_mode: str = "bilinear") -> None: - """ - Args: - size: images will be resized to this value. If a 2-integer tuple is given, it is interpreted as - (height, width). - interpolation_mode: interpolation mode for image resizing. Supported values: ``nearest``, ``bilinear``, - ``bicubic``, ``area``, and ``nearest-exact``. - """ - super().__init__() - self.size = size - self.interpolation_mode = interpolation_mode - - def forward(self, imgs: Union[Tensor, list[Tensor]]) -> tuple[Tensor, Tensor]: - """ - Returns: - resized_imgs: resized images in a batch. - original_sizes: the original image sizes of (height, width). - """ - # TODO: support other input formats e.g. file path, numpy - resized_imgs: list[Tensor] = [] - - iters = len(imgs) if isinstance(imgs, list) else imgs.shape[0] - original_sizes = imgs[0].new_zeros((iters, 2)) - for i in range(iters): - img = imgs[i] - original_sizes[i, 0] = img.shape[-2] # Height - original_sizes[i, 1] = img.shape[-1] # Width - resized_imgs.append(resize(img[None], size=self.size, interpolation=self.interpolation_mode)) - return concatenate(resized_imgs), original_sizes - - -# TODO: move this to kornia.models as AlgorithmicModel api -class ObjectDetector(Module): - """This class wraps an object detection model and performs pre-processing and post-processing.""" - - def __init__(self, model: Module, pre_processor: Module, post_processor: Module) -> None: - """Construct an Object Detector object. - - Args: - model: an object detection model. - pre_processor: a pre-processing module - post_processor: a post-processing module. - """ - super().__init__() - self.model = model.eval() - self.pre_processor = pre_processor.eval() - self.post_processor = post_processor.eval() +class BoundingBox(BoundingBoxBase): + warnings.warn("BoundingBox is deprecated and will be removed in v0.8.0. Use kornia.models.detector.BoundingBox instead.", DeprecationWarning) - @torch.inference_mode() - def forward(self, images: Union[Tensor, list[Tensor]]) -> Tensor: - """Detect objects in a given list of images. - Args: - images: If list of RGB images. Each image is a Tensor with shape :math:`(3, H, W)`. - If Tensor, a Tensor with shape :math:`(B, 3, H, W)`. +def results_from_detections(*args, **kwargs): + warnings.warn("results_from_detections is deprecated and will be removed in v0.8.0. Use kornia.models.detector.results_from_detections instead.", DeprecationWarning) + return results_from_detections_base(*args, **kwargs) - Returns: - list of detections found in each image. For item in a batch, shape is :math:`(D, 6)`, where :math:`D` is the - number of detections in the given image, :math:`6` represents class id, score, and `xywh` bounding box. - """ - images, images_sizes = self.pre_processor(images) - logits, boxes = self.model(images) - detections = self.post_processor(logits, boxes, images_sizes) - return detections - def draw( - self, images: Union[Tensor, list[Tensor]], detections: Optional[Tensor] = None, output_type: str = "torch" - ) -> Union[Tensor, list[Tensor], list[Image.Image]]: # type: ignore - """Very simple drawing. +class ResizePreProcessor(ResizePreProcessorBase): + warnings.warn("ResizePreProcessor is deprecated and will be removed in v0.8.0. Use kornia.models.utils.ResizePreProcessor instead.", DeprecationWarning) - Needs to be more fancy later. - """ - if detections is None: - detections = self.forward(images) - output = [] - for image, detection in zip(images, detections): - out_img = image[None].clone() - for out in detection: - out_img = draw_rectangle( - out_img, - torch.Tensor([[[out[-4], out[-3], out[-4] + out[-2], out[-3] + out[-1]]]]), - ) - if output_type == "torch": - output.append(out_img[0]) - elif output_type == "pil": - output.append(Image.fromarray((out_img[0] * 255).permute(1, 2, 0).numpy().astype(np.uint8))) # type: ignore - else: - raise RuntimeError(f"Unsupported output type `{output_type}`.") - return output - def save( - self, images: Union[Tensor, list[Tensor]], detections: Optional[Tensor] = None, directory: Optional[str] = None - ) -> None: - """Saves the output image(s) to a directory. +class ObjectDetector(ObjectDetectorBase): + warnings.warn("ObjectDetector is deprecated and will be removed in v0.8.0. Use kornia.models.detector.ObjectDetector instead.", DeprecationWarning) - Args: - name: Directory to save the images. - n_row: Number of images displayed in each row of the grid. - """ - if directory is None: - name = f"detection-{datetime.datetime.now(tz=datetime.timezone.utc).strftime('%Y%m%d%H%M%S')!s}" - directory = os.path.join("Kornia_outputs", name) - outputs = self.draw(images, detections) - os.makedirs(directory, exist_ok=True) - for i, out_image in enumerate(outputs): - write_image( - os.path.join(directory, f"{str(i).zfill(6)}.jpg"), - out_image.mul(255.0).byte(), - ) - logger.info(f"Outputs are saved in {directory}") - def compile( - self, - *, - fullgraph: bool = False, - dynamic: bool = False, - backend: str = "inductor", - mode: Optional[str] = None, - options: Optional[dict[str, str | int | bool]] = None, - disable: bool = False, - ) -> None: - """Compile the internal object detection model with :py:func:`torch.compile()`.""" - self.model = torch.compile( # type: ignore - self.model, - fullgraph=fullgraph, - dynamic=dynamic, - backend=backend, - mode=mode, - options=options, - disable=disable, - ) +class ObjectDetectorResult(ObjectDetectorResultBase): + warnings.warn("ObjectDetectorResult is deprecated and will be removed in v0.8.0. Use kornia.models.detector.ObjectDetectorResult instead.", DeprecationWarning) diff --git a/kornia/filters/dexined.py b/kornia/filters/dexined.py index 7c34fa9cf7..12bfac591c 100644 --- a/kornia/filters/dexined.py +++ b/kornia/filters/dexined.py @@ -125,8 +125,7 @@ def compute_out_features(self, idx: int, up_scale: int) -> int: def forward(self, x: Tensor, out_shape: list[int]) -> Tensor: out = self.features(x) - if out.shape[-2:] != out_shape: - out = F.interpolate(out, out_shape, mode="bilinear") + out = F.interpolate(out, out_shape, mode="bilinear") return out @@ -176,7 +175,7 @@ class DexiNed(Module): >>> img = torch.rand(1, 3, 320, 320) >>> net = DexiNed(pretrained=False) >>> out = net(img) - >>> out[-1].shape + >>> out.shape torch.Size([1, 1, 320, 320]) """ @@ -228,7 +227,7 @@ def load_from_file(self, path_file: str) -> None: self.load_state_dict(pretrained_dict, strict=True) self.eval() - def forward(self, x: Tensor) -> list[Tensor]: + def get_features(self, x: Tensor) -> list[Tensor]: # Block 1 block_1 = self.block_1(x) block_1_side = self.side_1(block_1) @@ -272,11 +271,13 @@ def forward(self, x: Tensor) -> list[Tensor]: out_5 = self.up_block_5(block_5, out_shape) out_6 = self.up_block_6(block_6, out_shape) results = [out_1, out_2, out_3, out_4, out_5, out_6] + return results + + def forward(self, x: Tensor) -> Tensor: + features = self.get_features(x) # concatenate multiscale outputs - block_cat = concatenate(results, 1) # Bx6xHxW + block_cat = concatenate(features, 1) # Bx6xHxW block_cat = self.block_cat(block_cat) # Bx1xHxW - # return results - results.append(block_cat) - return results + return block_cat diff --git a/kornia/models/__init__.py b/kornia/models/__init__.py index f7a598a13c..0de59dc90a 100644 --- a/kornia/models/__init__.py +++ b/kornia/models/__init__.py @@ -1 +1,3 @@ from . import detector +from . import segmentor +from . import tracker diff --git a/kornia/models/detector/__init__.py b/kornia/models/detector/__init__.py index c51e7ed66f..b511f133c2 100644 --- a/kornia/models/detector/__init__.py +++ b/kornia/models/detector/__init__.py @@ -1,2 +1,3 @@ +from .base import * from .rtdetr import * from .utils import * diff --git a/kornia/models/detector/base.py b/kornia/models/detector/base.py new file mode 100644 index 0000000000..2edf978eca --- /dev/null +++ b/kornia/models/detector/base.py @@ -0,0 +1,204 @@ +from __future__ import annotations + +import datetime +import logging +import os +from dataclasses import dataclass +from enum import Enum +from typing import Optional, Union + +import torch + +from kornia.core import Module, Tensor +from kornia.core.check import KORNIA_CHECK_SHAPE +from kornia.core.external import PILImage as Image +from kornia.core.external import numpy as np +from kornia.io import write_image +from kornia.utils.draw import draw_rectangle + +__all__ = [ + "BoundingBoxDataFormat", + "BoundingBox", + "results_from_detections", + "ObjectDetector", + "ObjectDetectorResult", +] + +logger = logging.getLogger(__name__) + + +class BoundingBoxDataFormat(Enum): + """Enum class that maps bounding box data format.""" + + XYWH = 0 + XYXY = 1 + CXCYWH = 2 + CENTER_XYWH = 2 + + +# NOTE: probably we should use a more generic name like BoundingBox2D +# and add a BoundingBox3D class for 3D bounding boxes. Also for serialization +# we should have an explicit class for each format to make it more production ready +# specially to serialize to protobuf and not saturate at a high rates. + + +@dataclass(frozen=True) +class BoundingBox: + """Bounding box data class. + + Useful for representing bounding boxes in different formats for object detection. + + Args: + data: tuple of bounding box data. The length of the tuple depends on the data format. + data_format: bounding box data format. + """ + + data: tuple[float, float, float, float] + data_format: BoundingBoxDataFormat + + +@dataclass(frozen=True) +class ObjectDetectorResult: + """Object detection result. + + Args: + class_id: class id of the detected object. + confidence: confidence score of the detected object. + bbox: bounding box of the detected object in xywh format. + """ + + class_id: int + confidence: float + bbox: BoundingBox + + +def results_from_detections(detections: Tensor, format: str | BoundingBoxDataFormat) -> list[ObjectDetectorResult]: + """Convert a detection tensor to a list of :py:class:`ObjectDetectorResult`. + + Args: + detections: tensor with shape :math:`(D, 6)`, where :math:`D` is the number of detections in the given image, + :math:`6` represents class id, score, and `xywh` bounding box. + + Returns: + list of :py:class:`ObjectDetectorResult`. + """ + KORNIA_CHECK_SHAPE(detections, ["D", "6"]) + + if isinstance(format, str): + format = BoundingBoxDataFormat[format.upper()] + + results: list[ObjectDetectorResult] = [] + for det in detections: + det = det.squeeze().tolist() + if len(det) != 6: + continue + results.append( + ObjectDetectorResult( + class_id=int(det[0]), + confidence=det[1], + bbox=BoundingBox(data=(det[2], det[3], det[4], det[5]), data_format=format), + ) + ) + return results + + +# TODO: move this to kornia.models as AlgorithmicModel api +class ObjectDetector(Module): + """This class wraps an object detection model and performs pre-processing and post-processing.""" + + def __init__(self, model: Module, pre_processor: Module, post_processor: Module) -> None: + """Construct an Object Detector object. + + Args: + model: an object detection model. + pre_processor: a pre-processing module + post_processor: a post-processing module. + """ + super().__init__() + self.model = model.eval() + self.pre_processor = pre_processor.eval() + self.post_processor = post_processor.eval() + + @torch.inference_mode() + def forward(self, images: Union[Tensor, list[Tensor]]) -> Tensor: + """Detect objects in a given list of images. + + Args: + images: If list of RGB images. Each image is a Tensor with shape :math:`(3, H, W)`. + If Tensor, a Tensor with shape :math:`(B, 3, H, W)`. + + Returns: + list of detections found in each image. For item in a batch, shape is :math:`(D, 6)`, where :math:`D` is the + number of detections in the given image, :math:`6` represents class id, score, and `xywh` bounding box. + """ + images, images_sizes = self.pre_processor(images) + logits, boxes = self.model(images) + detections = self.post_processor(logits, boxes, images_sizes) + return detections + + def draw( + self, images: Union[Tensor, list[Tensor]], detections: Optional[Tensor] = None, output_type: str = "torch" + ) -> Union[Tensor, list[Tensor], list[Image.Image]]: # type: ignore + """Very simple drawing. + + Needs to be more fancy later. + """ + if detections is None: + detections = self.forward(images) + output = [] + for image, detection in zip(images, detections): + out_img = image[None].clone() + for out in detection: + out_img = draw_rectangle( + out_img, + torch.Tensor([[[out[-4], out[-3], out[-4] + out[-2], out[-3] + out[-1]]]]), + ) + if output_type == "torch": + output.append(out_img[0]) + elif output_type == "pil": + output.append(Image.fromarray((out_img[0] * 255).permute(1, 2, 0).numpy().astype(np.uint8))) # type: ignore + else: + raise RuntimeError(f"Unsupported output type `{output_type}`.") + return output + + def save( + self, images: Union[Tensor, list[Tensor]], detections: Optional[Tensor] = None, directory: Optional[str] = None + ) -> None: + """Saves the output image(s) to a directory. + + Args: + name: Directory to save the images. + n_row: Number of images displayed in each row of the grid. + """ + if directory is None: + name = f"detection-{datetime.datetime.now(tz=datetime.timezone.utc).strftime('%Y%m%d%H%M%S')!s}" + directory = os.path.join("Kornia_outputs", name) + outputs = self.draw(images, detections) + os.makedirs(directory, exist_ok=True) + for i, out_image in enumerate(outputs): + write_image( + os.path.join(directory, f"{str(i).zfill(6)}.jpg"), + out_image.mul(255.0).byte(), + ) + logger.info(f"Outputs are saved in {directory}") + + def compile( + self, + *, + fullgraph: bool = False, + dynamic: bool = False, + backend: str = "inductor", + mode: Optional[str] = None, + options: Optional[dict[str, str | int | bool]] = None, + disable: bool = False, + ) -> None: + """Compile the internal object detection model with :py:func:`torch.compile()`.""" + self.model = torch.compile( # type: ignore + self.model, + fullgraph=fullgraph, + dynamic=dynamic, + backend=backend, + mode=mode, + options=options, + disable=disable, + ) diff --git a/kornia/models/detector/rtdetr.py b/kornia/models/detector/rtdetr.py index e4117127ed..36ffe78789 100644 --- a/kornia/models/detector/rtdetr.py +++ b/kornia/models/detector/rtdetr.py @@ -6,7 +6,8 @@ from kornia.contrib.models.rt_detr import DETRPostProcessor from kornia.contrib.models.rt_detr.model import RTDETR, RTDETRConfig -from kornia.contrib.object_detection import ObjectDetector, ResizePreProcessor +from kornia.models.detector.base import ObjectDetector +from kornia.models.utils import ResizePreProcessor from kornia.core import rand __all__ = ["RTDETRDetectorBuilder"] diff --git a/kornia/models/detector/utils.py b/kornia/models/detector/utils.py index fc2f9db34b..51a11eed5d 100644 --- a/kornia/models/detector/utils.py +++ b/kornia/models/detector/utils.py @@ -6,6 +6,13 @@ class BoxFiltering(Module, ONNXExportMixin): + """Filter boxes according to the desired threshold. + + Args: + confidence_threshold: an 0-d scalar that represents the desired threshold. + filter_as_zero: whether to filter boxes as zero. + """ + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, -1, 6] ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, -1, 6] ONNX_EXPORT_PSEUDO_SHAPE: ClassVar[List[int]] = [5, 20, 6] diff --git a/kornia/models/edge_detector/__init__.py b/kornia/models/edge_detector/__init__.py new file mode 100644 index 0000000000..c9359ff78d --- /dev/null +++ b/kornia/models/edge_detector/__init__.py @@ -0,0 +1,2 @@ +from .base import * +from .dexined import * diff --git a/kornia/models/edge_detector/base.py b/kornia/models/edge_detector/base.py new file mode 100644 index 0000000000..8f6b882e1e --- /dev/null +++ b/kornia/models/edge_detector/base.py @@ -0,0 +1,67 @@ +from typing import Union + +from kornia.core import Module, Tensor +from kornia.core.external import PILImage as Image + +__all__ = ["EdgeDetector"] + + +class EdgeDetector(Module): + """EdgeDetector is a module that wraps an edge detection model. + + This module uses EdgeDetectionModel library for edge detection. + """ + # NOTE: We need to implement this class for better visualization and user experience. + + def __init__(self, model: Module, pre_processor: Module, post_processor: Module) -> None: + """Construct an Object Detector object. + + Args: + model: an object detection model. + pre_processor: a pre-processing module + """ + super().__init__() + self.model = model + self.pre_processor = pre_processor + self.post_processor = post_processor + + def forward(self, images: Union[Tensor, list[Tensor]]) -> Tensor: + """Forward pass of the semantic segmentation model. + + Args: + x: input tensor. + + Returns: + output tensor. + """ + images, image_sizes = self.pre_processor(images) + out_images = self.model(images) + return self.post_processor(out_images, image_sizes) + + def draw( + self, images: Union[Tensor, list[Tensor]], output_type: str = "torch" + ) -> Union[Tensor, list[Tensor], list[Image.Image]]: # type: ignore + """Draw the segmentation results. + + Args: + images: input tensor. + output_type: type of the output. + + Returns: + output tensor. + """ + ... + + def save( + self, images: Union[Tensor, list[Tensor]], output_type: str = "torch" + ) -> None: + """Save the segmentation results. + + Args: + images: input tensor. + output_type: type of the output. + + Returns: + output tensor. + """ + ... diff --git a/kornia/models/edge_detector/dexined.py b/kornia/models/edge_detector/dexined.py new file mode 100644 index 0000000000..5e22e5cae1 --- /dev/null +++ b/kornia/models/edge_detector/dexined.py @@ -0,0 +1,55 @@ +from typing import Optional, Tuple +import torch +import torch.nn as nn + +from kornia.core import rand + +from kornia.filters.dexined import DexiNed +from kornia.models.edge_detector.base import EdgeDetector +from kornia.models.utils import ResizePreProcessor, ResizePostProcessor + + +class DexiNedBuilder: + + @staticmethod + def build( + pretrained: bool = True, image_size: Optional[int] = 352 + ) -> EdgeDetector: + model = DexiNed(pretrained=pretrained) + return EdgeDetector( + model, + ResizePreProcessor((image_size, image_size)) if image_size is not None else nn.Identity(), + ResizePostProcessor() if image_size is not None else nn.Identity(), + ) + + @staticmethod + def to_onnx( + onnx_name: Optional[str] = None, + pretrained: bool = True, + image_size: Optional[int] = 352, + ) -> Tuple[str, EdgeDetector]: + edge_detector = DexiNedBuilder.build(pretrained, image_size) + model_name = "DexiNed" + if onnx_name is None: + _model_name = model_name + onnx_name = f"Kornia-RTDETR-{_model_name}-{image_size}.onnx" + + if image_size is None: + val_image = rand(1, 3, 352, 352) + else: + val_image = rand(1, 3, image_size, image_size) + + dynamic_axes = {"input": {0: "batch_size", 2: "height", 3: "width"}, "output": {0: "batch_size"}} + torch.onnx.export( + edge_detector, + val_image, + onnx_name, + export_params=True, + opset_version=17, + do_constant_folding=True, + input_names=["input"], + output_names=["output"], + dynamic_axes=dynamic_axes, + ) + + return onnx_name, edge_detector diff --git a/kornia/models/segmentor/__init__.py b/kornia/models/segmentor/__init__.py index e69de29bb2..1bc6249871 100644 --- a/kornia/models/segmentor/__init__.py +++ b/kornia/models/segmentor/__init__.py @@ -0,0 +1 @@ +from .segmentation_models import * diff --git a/kornia/models/segmentor/base.py b/kornia/models/segmentor/base.py new file mode 100644 index 0000000000..ef117432da --- /dev/null +++ b/kornia/models/segmentor/base.py @@ -0,0 +1,63 @@ +from typing import Union + +from kornia.core import Module, Tensor +from kornia.core.external import PILImage as Image + + +class SemanticSegmentation(Module): + """Semantic Segmentation is a module that wraps a semantic segmentation model. + + This module uses SegmentationModel library for semantic segmentation. + """ + # NOTE: We need to implement this class for better visualization and user experience. + + def __init__(self, model: Module, pre_processor: Module) -> None: + """Construct an Object Detector object. + + Args: + model: an object detection model. + pre_processor: a pre-processing module + """ + super().__init__() + self.model = model + self.pre_processor = pre_processor + + def forward(self, images: Union[Tensor, list[Tensor]]) -> Tensor: + """Forward pass of the semantic segmentation model. + + Args: + x: input tensor. + + Returns: + output tensor. + """ + images = self.pre_processor(images) + return self.model(images) + + def draw( + self, images: Union[Tensor, list[Tensor]], output_type: str = "torch" + ) -> Union[Tensor, list[Tensor], list[Image.Image]]: # type: ignore + """Draw the segmentation results. + + Args: + images: input tensor. + output_type: type of the output. + + Returns: + output tensor. + """ + ... + + def save( + self, images: Union[Tensor, list[Tensor]], output_type: str = "torch" + ) -> None: + """Save the segmentation results. + + Args: + images: input tensor. + output_type: type of the output. + + Returns: + output tensor. + """ + ... diff --git a/kornia/models/segmentor/segmentation_models.py b/kornia/models/segmentor/segmentation_models.py index 09093aac11..0b9cdc7f26 100644 --- a/kornia/models/segmentor/segmentation_models.py +++ b/kornia/models/segmentor/segmentation_models.py @@ -1,7 +1,7 @@ from typing import Optional import kornia -from kornia.core import Module, Tensor, tensor, zeros_like, ones_like +from kornia.core import Module,Tensor, tensor, zeros_like, ones_like from kornia.core.module import ONNXExportMixin from kornia.core.external import segmentation_models_pytorch as smp diff --git a/kornia/models/tracker/boxmot_tracker.py b/kornia/models/tracker/boxmot_tracker.py index 97a09bf87b..9dc63ad00e 100644 --- a/kornia/models/tracker/boxmot_tracker.py +++ b/kornia/models/tracker/boxmot_tracker.py @@ -2,7 +2,7 @@ from pathlib import Path from typing import Union -from kornia.contrib.object_detection import ObjectDetector +from kornia.models.detector.base import ObjectDetector from kornia.core import Tensor from kornia.core.external import boxmot from kornia.core.external import numpy as np diff --git a/kornia/models/utils.py b/kornia/models/utils.py new file mode 100644 index 0000000000..514f8ba696 --- /dev/null +++ b/kornia/models/utils.py @@ -0,0 +1,78 @@ +from typing import Union + +import warnings +import torch +from torch import Tensor + +from kornia.core import Module, concatenate +from kornia.geometry.transform import resize + + +class ResizePreProcessor(Module): + """This module resizes a list of image tensors to the given size. + + Additionally, also returns the original image sizes for further post-processing. + """ + + def __init__(self, size: tuple[int, int], interpolation_mode: str = "bilinear") -> None: + """ + Args: + size: images will be resized to this value. If a 2-integer tuple is given, it is interpreted as + (height, width). + interpolation_mode: interpolation mode for image resizing. Supported values: ``nearest``, ``bilinear``, + ``bicubic``, ``area``, and ``nearest-exact``. + """ + super().__init__() + self.size = size + self.interpolation_mode = interpolation_mode + + def forward(self, imgs: Union[Tensor, list[Tensor]]) -> tuple[Tensor, Tensor]: + """ + Returns: + resized_imgs: resized images in a batch. + original_sizes: the original image sizes of (height, width). + """ + # TODO: support other input formats e.g. file path, numpy + resized_imgs: list[Tensor] = [] + + iters = len(imgs) if isinstance(imgs, list) else imgs.shape[0] + original_sizes = imgs[0].new_zeros((iters, 2)) + for i in range(iters): + img = imgs[i] + original_sizes[i, 0] = img.shape[-2] # Height + original_sizes[i, 1] = img.shape[-1] # Width + resized_imgs.append(resize(img[None], size=self.size, interpolation=self.interpolation_mode)) + return concatenate(resized_imgs), original_sizes + + +class ResizePostProcessor(Module): + def __init__(self, interpolation_mode: str = "bilinear") -> None: + super().__init__() + self.interpolation_mode = interpolation_mode + + def forward( + self, imgs: Union[Tensor, list[Tensor]], original_sizes: Tensor + ) -> Tensor: + """ + Returns: + resized_imgs: resized images in a batch. + original_sizes: the original image sizes of (height, width). + """ + # TODO: support other input formats e.g. file path, numpy + resized_imgs: list[Tensor] = [] + + if not torch.onnx.is_in_onnx_export(): + iters = len(imgs) if isinstance(imgs, list) else imgs.shape[0] + for i in range(iters): + img = imgs[i] + size = original_sizes[i] + resized_imgs.append(resize( + img[None], size=size.cpu().long().numpy().tolist(), interpolation=self.interpolation_mode)) + else: + warnings.warn( + "ResizePostProcessor is not supported in ONNX export. " + "The output will not be resized back to the original size." + ) + resized_imgs = imgs + + return resized_imgs \ No newline at end of file diff --git a/kornia/onnx/sequential.py b/kornia/onnx/sequential.py index d4dea8c26e..fad9cc8b51 100644 --- a/kornia/onnx/sequential.py +++ b/kornia/onnx/sequential.py @@ -137,9 +137,9 @@ def as_cpu(self) -> None: """Set the session to run on CPU.""" self._session = self.create_session(providers=['CPUExecutionProvider']) - def as_cuda(self, device_id: int = 0) -> None: + def as_cuda(self, device: int = 0) -> None: """Set the session to run on CUDA.""" - self._session = self.create_session(providers=[('CUDAExecutionProvider', {'device_id': device_id})]) + self._session = self.create_session(providers=[('CUDAExecutionProvider', {'device_id': device})]) def __call__(self, *inputs: "np.ndarray") -> List["np.ndarray"]: # type:ignore """Perform inference using the combined ONNX model. From e486f29624f0b35055bad0a4a1bff95f8c48b3fc Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 15 Sep 2024 05:53:43 +0000 Subject: [PATCH 60/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/contrib/object_detection.py | 38 +++++++++++++++---- kornia/models/__init__.py | 4 +- kornia/models/detector/rtdetr.py | 2 +- kornia/models/edge_detector/base.py | 5 +-- kornia/models/edge_detector/dexined.py | 11 ++---- kornia/models/segmentor/base.py | 5 +-- .../models/segmentor/segmentation_models.py | 10 +++-- kornia/models/tracker/boxmot_tracker.py | 2 +- kornia/models/utils.py | 13 +++---- kornia/onnx/sequential.py | 6 +-- 10 files changed, 57 insertions(+), 39 deletions(-) diff --git a/kornia/contrib/object_detection.py b/kornia/contrib/object_detection.py index d6a7de5ad7..2917834368 100644 --- a/kornia/contrib/object_detection.py +++ b/kornia/contrib/object_detection.py @@ -1,13 +1,21 @@ import warnings -from kornia.models.utils import ResizePreProcessor as ResizePreProcessorBase from kornia.models.detector.base import ( BoundingBox as BoundingBoxBase, +) +from kornia.models.detector.base import ( BoundingBoxDataFormat, - results_from_detections as results_from_detections_base, +) +from kornia.models.detector.base import ( ObjectDetector as ObjectDetectorBase, +) +from kornia.models.detector.base import ( ObjectDetectorResult as ObjectDetectorResultBase, ) +from kornia.models.detector.base import ( + results_from_detections as results_from_detections_base, +) +from kornia.models.utils import ResizePreProcessor as ResizePreProcessorBase __all__ = [ "BoundingBoxDataFormat", @@ -18,22 +26,38 @@ "ObjectDetectorResult", ] + class BoundingBox(BoundingBoxBase): - warnings.warn("BoundingBox is deprecated and will be removed in v0.8.0. Use kornia.models.detector.BoundingBox instead.", DeprecationWarning) + warnings.warn( + "BoundingBox is deprecated and will be removed in v0.8.0. Use kornia.models.detector.BoundingBox instead.", + DeprecationWarning, + ) def results_from_detections(*args, **kwargs): - warnings.warn("results_from_detections is deprecated and will be removed in v0.8.0. Use kornia.models.detector.results_from_detections instead.", DeprecationWarning) + warnings.warn( + "results_from_detections is deprecated and will be removed in v0.8.0. Use kornia.models.detector.results_from_detections instead.", + DeprecationWarning, + ) return results_from_detections_base(*args, **kwargs) class ResizePreProcessor(ResizePreProcessorBase): - warnings.warn("ResizePreProcessor is deprecated and will be removed in v0.8.0. Use kornia.models.utils.ResizePreProcessor instead.", DeprecationWarning) + warnings.warn( + "ResizePreProcessor is deprecated and will be removed in v0.8.0. Use kornia.models.utils.ResizePreProcessor instead.", + DeprecationWarning, + ) class ObjectDetector(ObjectDetectorBase): - warnings.warn("ObjectDetector is deprecated and will be removed in v0.8.0. Use kornia.models.detector.ObjectDetector instead.", DeprecationWarning) + warnings.warn( + "ObjectDetector is deprecated and will be removed in v0.8.0. Use kornia.models.detector.ObjectDetector instead.", + DeprecationWarning, + ) class ObjectDetectorResult(ObjectDetectorResultBase): - warnings.warn("ObjectDetectorResult is deprecated and will be removed in v0.8.0. Use kornia.models.detector.ObjectDetectorResult instead.", DeprecationWarning) + warnings.warn( + "ObjectDetectorResult is deprecated and will be removed in v0.8.0. Use kornia.models.detector.ObjectDetectorResult instead.", + DeprecationWarning, + ) diff --git a/kornia/models/__init__.py b/kornia/models/__init__.py index 0de59dc90a..ec936c3a85 100644 --- a/kornia/models/__init__.py +++ b/kornia/models/__init__.py @@ -1,3 +1 @@ -from . import detector -from . import segmentor -from . import tracker +from . import detector, segmentor, tracker diff --git a/kornia/models/detector/rtdetr.py b/kornia/models/detector/rtdetr.py index 36ffe78789..96ead94a89 100644 --- a/kornia/models/detector/rtdetr.py +++ b/kornia/models/detector/rtdetr.py @@ -6,9 +6,9 @@ from kornia.contrib.models.rt_detr import DETRPostProcessor from kornia.contrib.models.rt_detr.model import RTDETR, RTDETRConfig +from kornia.core import rand from kornia.models.detector.base import ObjectDetector from kornia.models.utils import ResizePreProcessor -from kornia.core import rand __all__ = ["RTDETRDetectorBuilder"] diff --git a/kornia/models/edge_detector/base.py b/kornia/models/edge_detector/base.py index 8f6b882e1e..1ca3477a80 100644 --- a/kornia/models/edge_detector/base.py +++ b/kornia/models/edge_detector/base.py @@ -11,6 +11,7 @@ class EdgeDetector(Module): This module uses EdgeDetectionModel library for edge detection. """ + # NOTE: We need to implement this class for better visualization and user experience. def __init__(self, model: Module, pre_processor: Module, post_processor: Module) -> None: @@ -52,9 +53,7 @@ def draw( """ ... - def save( - self, images: Union[Tensor, list[Tensor]], output_type: str = "torch" - ) -> None: + def save(self, images: Union[Tensor, list[Tensor]], output_type: str = "torch") -> None: """Save the segmentation results. Args: diff --git a/kornia/models/edge_detector/dexined.py b/kornia/models/edge_detector/dexined.py index 5e22e5cae1..ba8b21c228 100644 --- a/kornia/models/edge_detector/dexined.py +++ b/kornia/models/edge_detector/dexined.py @@ -1,20 +1,17 @@ from typing import Optional, Tuple + import torch -import torch.nn as nn +from torch import nn from kornia.core import rand - from kornia.filters.dexined import DexiNed from kornia.models.edge_detector.base import EdgeDetector -from kornia.models.utils import ResizePreProcessor, ResizePostProcessor +from kornia.models.utils import ResizePostProcessor, ResizePreProcessor class DexiNedBuilder: - @staticmethod - def build( - pretrained: bool = True, image_size: Optional[int] = 352 - ) -> EdgeDetector: + def build(pretrained: bool = True, image_size: Optional[int] = 352) -> EdgeDetector: model = DexiNed(pretrained=pretrained) return EdgeDetector( model, diff --git a/kornia/models/segmentor/base.py b/kornia/models/segmentor/base.py index ef117432da..ec6293f40c 100644 --- a/kornia/models/segmentor/base.py +++ b/kornia/models/segmentor/base.py @@ -9,6 +9,7 @@ class SemanticSegmentation(Module): This module uses SegmentationModel library for semantic segmentation. """ + # NOTE: We need to implement this class for better visualization and user experience. def __init__(self, model: Module, pre_processor: Module) -> None: @@ -48,9 +49,7 @@ def draw( """ ... - def save( - self, images: Union[Tensor, list[Tensor]], output_type: str = "torch" - ) -> None: + def save(self, images: Union[Tensor, list[Tensor]], output_type: str = "torch") -> None: """Save the segmentation results. Args: diff --git a/kornia/models/segmentor/segmentation_models.py b/kornia/models/segmentor/segmentation_models.py index 0b9cdc7f26..d50674e05b 100644 --- a/kornia/models/segmentor/segmentation_models.py +++ b/kornia/models/segmentor/segmentation_models.py @@ -1,9 +1,9 @@ from typing import Optional import kornia -from kornia.core import Module,Tensor, tensor, zeros_like, ones_like -from kornia.core.module import ONNXExportMixin +from kornia.core import Module, Tensor, ones_like, tensor, zeros_like from kornia.core.external import segmentation_models_pytorch as smp +from kornia.core.module import ONNXExportMixin class SegmentationModels(Module, ONNXExportMixin): @@ -53,7 +53,9 @@ def __init__( def preprocessing(self, input: Tensor) -> Tensor: # Ensure the color space transformation is ONNX-friendly input_space = self.preproc_params["input_space"] - input = kornia.color.rgb_to_bgr(input) if input_space == "BGR" else input # Assume input is already RGB if not BGR + input = ( + kornia.color.rgb_to_bgr(input) if input_space == "BGR" else input + ) # Assume input is already RGB if not BGR # Normalize input range if needed input_range = self.preproc_params["input_range"] @@ -74,7 +76,7 @@ def preprocessing(self, input: Tensor) -> Tensor: std = tensor([self.preproc_params["std"]]).to(input.device) else: std = ones_like(input) - + return kornia.enhance.normalize(input, mean, std) def forward(self, input: Tensor) -> Tensor: diff --git a/kornia/models/tracker/boxmot_tracker.py b/kornia/models/tracker/boxmot_tracker.py index 9dc63ad00e..4cc8be105e 100644 --- a/kornia/models/tracker/boxmot_tracker.py +++ b/kornia/models/tracker/boxmot_tracker.py @@ -2,10 +2,10 @@ from pathlib import Path from typing import Union -from kornia.models.detector.base import ObjectDetector from kornia.core import Tensor from kornia.core.external import boxmot from kornia.core.external import numpy as np +from kornia.models.detector.base import ObjectDetector from kornia.models.detector.rtdetr import RTDETRDetectorBuilder from kornia.utils.image import tensor_to_image diff --git a/kornia/models/utils.py b/kornia/models/utils.py index 514f8ba696..08676160c2 100644 --- a/kornia/models/utils.py +++ b/kornia/models/utils.py @@ -1,6 +1,6 @@ +import warnings from typing import Union -import warnings import torch from torch import Tensor @@ -50,9 +50,7 @@ def __init__(self, interpolation_mode: str = "bilinear") -> None: super().__init__() self.interpolation_mode = interpolation_mode - def forward( - self, imgs: Union[Tensor, list[Tensor]], original_sizes: Tensor - ) -> Tensor: + def forward(self, imgs: Union[Tensor, list[Tensor]], original_sizes: Tensor) -> Tensor: """ Returns: resized_imgs: resized images in a batch. @@ -66,8 +64,9 @@ def forward( for i in range(iters): img = imgs[i] size = original_sizes[i] - resized_imgs.append(resize( - img[None], size=size.cpu().long().numpy().tolist(), interpolation=self.interpolation_mode)) + resized_imgs.append( + resize(img[None], size=size.cpu().long().numpy().tolist(), interpolation=self.interpolation_mode) + ) else: warnings.warn( "ResizePostProcessor is not supported in ONNX export. " @@ -75,4 +74,4 @@ def forward( ) resized_imgs = imgs - return resized_imgs \ No newline at end of file + return resized_imgs diff --git a/kornia/onnx/sequential.py b/kornia/onnx/sequential.py index fad9cc8b51..c934a0550e 100644 --- a/kornia/onnx/sequential.py +++ b/kornia/onnx/sequential.py @@ -112,7 +112,7 @@ def create_session( session = ort.InferenceSession( # type:ignore self._combined_op.SerializeToString(), sess_options=sess_options, - providers=providers or ['CPUExecutionProvider'], + providers=providers or ["CPUExecutionProvider"], ) return session @@ -135,11 +135,11 @@ def get_session(self) -> "ort.InferenceSession": # type: ignore def as_cpu(self) -> None: """Set the session to run on CPU.""" - self._session = self.create_session(providers=['CPUExecutionProvider']) + self._session = self.create_session(providers=["CPUExecutionProvider"]) def as_cuda(self, device: int = 0) -> None: """Set the session to run on CUDA.""" - self._session = self.create_session(providers=[('CUDAExecutionProvider', {'device_id': device})]) + self._session = self.create_session(providers=[("CUDAExecutionProvider", {"device_id": device})]) def __call__(self, *inputs: "np.ndarray") -> List["np.ndarray"]: # type:ignore """Perform inference using the combined ONNX model. From aa707a534048feb7ea534172d7cb03d498d84c45 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Sun, 15 Sep 2024 08:56:55 +0300 Subject: [PATCH 61/84] update --- kornia/models/edge_detector/dexined.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/kornia/models/edge_detector/dexined.py b/kornia/models/edge_detector/dexined.py index ba8b21c228..841f1f1b74 100644 --- a/kornia/models/edge_detector/dexined.py +++ b/kornia/models/edge_detector/dexined.py @@ -26,10 +26,8 @@ def to_onnx( image_size: Optional[int] = 352, ) -> Tuple[str, EdgeDetector]: edge_detector = DexiNedBuilder.build(pretrained, image_size) - model_name = "DexiNed" if onnx_name is None: - _model_name = model_name - onnx_name = f"Kornia-RTDETR-{_model_name}-{image_size}.onnx" + onnx_name = f"Kornia-DexiNed-{image_size}.onnx" if image_size is None: val_image = rand(1, 3, 352, 352) From 1dfaf1a71205f8011f1a0913053f90088d510af6 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Sun, 15 Sep 2024 16:19:06 +0300 Subject: [PATCH 62/84] update --- kornia/__init__.py | 1 + kornia/config.py | 51 +++++++++++++++++++ .../contrib/models/rt_detr/post_processor.py | 2 +- kornia/contrib/object_detection.py | 10 ++-- kornia/core/external.py | 39 +++++++++----- kornia/models/__init__.py | 2 +- .../{detector => detection}/__init__.py | 0 kornia/models/{detector => detection}/base.py | 7 +-- .../models/{detector => detection}/rtdetr.py | 6 +-- .../models/{detector => detection}/utils.py | 0 .../__init__.py | 0 .../{edge_detector => edge_detection}/base.py | 0 .../dexined.py | 7 +-- .../{segmentor => segmentation}/__init__.py | 0 .../{segmentor => segmentation}/base.py | 0 .../segmentation_models.py | 0 .../models/{tracker => tracking}/__init__.py | 0 .../{tracker => tracking}/boxmot_tracker.py | 4 +- kornia/models/utils.py | 8 +-- kornia/onnx/sequential.py | 5 +- kornia/onnx/utils.py | 10 ++-- 21 files changed, 111 insertions(+), 41 deletions(-) create mode 100644 kornia/config.py rename kornia/models/{detector => detection}/__init__.py (100%) rename kornia/models/{detector => detection}/base.py (95%) rename kornia/models/{detector => detection}/rtdetr.py (97%) rename kornia/models/{detector => detection}/utils.py (100%) rename kornia/models/{edge_detector => edge_detection}/__init__.py (100%) rename kornia/models/{edge_detector => edge_detection}/base.py (100%) rename kornia/models/{edge_detector => edge_detection}/dexined.py (86%) rename kornia/models/{segmentor => segmentation}/__init__.py (100%) rename kornia/models/{segmentor => segmentation}/base.py (100%) rename kornia/models/{segmentor => segmentation}/segmentation_models.py (100%) rename kornia/models/{tracker => tracking}/__init__.py (100%) rename kornia/models/{tracker => tracking}/boxmot_tracker.py (97%) diff --git a/kornia/__init__.py b/kornia/__init__.py index 9bdc9c1d66..b1b9782747 100644 --- a/kornia/__init__.py +++ b/kornia/__init__.py @@ -10,6 +10,7 @@ color, contrib, core, + config, enhance, feature, io, diff --git a/kornia/config.py b/kornia/config.py new file mode 100644 index 0000000000..857a566e28 --- /dev/null +++ b/kornia/config.py @@ -0,0 +1,51 @@ +from enum import Enum +from dataclasses import dataclass, field + +__all__ = ["config", "InstallationMode"] + + +class InstallationMode(str, Enum): + # Ask the user if to install the dependencies + ASK = "ASK" + # Install the dependencies + AUTO = "AUTO" + # Raise an error if the dependencies are not installed + RAISE = "RAISE" + + def __eq__(self, other): + if isinstance(other, str): + return self.value.lower() == other.lower() # Case-insensitive comparison + return super().__eq__(other) + + +class LazyLoaderConfig: + _installation_mode: InstallationMode = InstallationMode.ASK + + @property + def installation_mode(self) -> InstallationMode: + return self._installation_mode + + @installation_mode.setter + def installation_mode(self, value: str): + # Allow setting via string by converting to the Enum + if isinstance(value, str): + try: + self._installation_mode = InstallationMode(value.upper()) + except ValueError: + raise ValueError(f"{value} is not a valid InstallationMode. Choose from: {list(InstallationMode)}") + elif isinstance(value, InstallationMode): + self._installation_mode = value + else: + raise TypeError("installation_mode must be a string or InstallationMode Enum.") + + +@dataclass +class KorniaConfig: + output_dir: str = "kornia_outputs" + hub_cache_dir: str = ".kornia_hub" + hub_models_dir: str = ".kornia_hub/models" + hub_onnx_dir: str = ".kornia_hub/onnx_models" + lazyloader: LazyLoaderConfig = field(default_factory=LazyLoaderConfig) + + +kornia_config = KorniaConfig() diff --git a/kornia/contrib/models/rt_detr/post_processor.py b/kornia/contrib/models/rt_detr/post_processor.py index ae9c4ac791..0df2c1190c 100644 --- a/kornia/contrib/models/rt_detr/post_processor.py +++ b/kornia/contrib/models/rt_detr/post_processor.py @@ -7,7 +7,7 @@ import torch from kornia.core import Module, Tensor, concatenate, tensor -from kornia.models.detector.utils import BoxFiltering +from kornia.models.detection.utils import BoxFiltering def mod(a: Tensor, b: int) -> Tensor: diff --git a/kornia/contrib/object_detection.py b/kornia/contrib/object_detection.py index 2917834368..c8c4e336e5 100644 --- a/kornia/contrib/object_detection.py +++ b/kornia/contrib/object_detection.py @@ -1,18 +1,18 @@ import warnings -from kornia.models.detector.base import ( +from kornia.models.detection.base import ( BoundingBox as BoundingBoxBase, ) -from kornia.models.detector.base import ( +from kornia.models.detection.base import ( BoundingBoxDataFormat, ) -from kornia.models.detector.base import ( +from kornia.models.detection.base import ( ObjectDetector as ObjectDetectorBase, ) -from kornia.models.detector.base import ( +from kornia.models.detection.base import ( ObjectDetectorResult as ObjectDetectorResultBase, ) -from kornia.models.detector.base import ( +from kornia.models.detection.base import ( results_from_detections as results_from_detections_base, ) from kornia.models.utils import ResizePreProcessor as ResizePreProcessorBase diff --git a/kornia/core/external.py b/kornia/core/external.py index de1544f7ca..39c0b789d1 100644 --- a/kornia/core/external.py +++ b/kornia/core/external.py @@ -5,6 +5,8 @@ from types import ModuleType from typing import List, Optional +from kornia.config import kornia_config, InstallationMode + logger = logging.getLogger(__name__) @@ -45,23 +47,36 @@ def _load(self) -> None: try: self.module = importlib.import_module(self.module_name) except ImportError as e: - if self.auto_install: + if kornia_config.lazyloader.installation_mode == InstallationMode.AUTO or self.auto_install: self._install_package(self.module_name) - else: + elif kornia_config.lazyloader.installation_mode == InstallationMode.ASK: + to_ask = True if_install = input( f"Optional dependency '{self.module_name}' is not installed. " + "You may silent this prompt by `kornia_config.lazyloader.installation_mode = 'auto'`. " "Do you wish to install the dependency? [Y]es, [N]o, [A]ll." ) - if if_install.lower() == "y": - self._install_package(self.module_name) - elif if_install.lower() == "a": - self.auto_install = True - self._install_package(self.module_name) - else: - raise ImportError( - f"Optional dependency '{self.module_name}' is not installed. " - f"Please install it to use this functionality." - ) from e + while to_ask: + if if_install.lower() == "y" or if_install.lower() == "yes": + self._install_package(self.module_name) + to_ask = False + elif if_install.lower() == "a" or if_install.lower() == "all": + self.auto_install = True + self._install_package(self.module_name) + to_ask = False + elif if_install.lower() == "n" or if_install.lower() == "no": + raise ImportError( + f"Optional dependency '{self.module_name}' is not installed. " + f"Please install it to use this functionality." + ) from e + else: + if_install = input("Invalid input. Please enter 'Y', 'N', or 'A'.") + + elif kornia_config.lazyloader.installation_mode == InstallationMode.RAISE: + raise ImportError( + f"Optional dependency '{self.module_name}' is not installed. " + f"Please install it to use this functionality." + ) from e def __getattr__(self, item: str) -> object: """Loads the module (if not already loaded) and returns the requested attribute. diff --git a/kornia/models/__init__.py b/kornia/models/__init__.py index ec936c3a85..9117c87c71 100644 --- a/kornia/models/__init__.py +++ b/kornia/models/__init__.py @@ -1 +1 @@ -from . import detector, segmentor, tracker +from . import detection, segmentation, tracking diff --git a/kornia/models/detector/__init__.py b/kornia/models/detection/__init__.py similarity index 100% rename from kornia/models/detector/__init__.py rename to kornia/models/detection/__init__.py diff --git a/kornia/models/detector/base.py b/kornia/models/detection/base.py similarity index 95% rename from kornia/models/detector/base.py rename to kornia/models/detection/base.py index 2edf978eca..744fc579d2 100644 --- a/kornia/models/detector/base.py +++ b/kornia/models/detection/base.py @@ -15,6 +15,7 @@ from kornia.core.external import numpy as np from kornia.io import write_image from kornia.utils.draw import draw_rectangle +from kornia.utils.image import tensor_to_image __all__ = [ "BoundingBoxDataFormat", @@ -156,7 +157,7 @@ def draw( if output_type == "torch": output.append(out_img[0]) elif output_type == "pil": - output.append(Image.fromarray((out_img[0] * 255).permute(1, 2, 0).numpy().astype(np.uint8))) # type: ignore + output.append(Image.fromarray((tensor_to_image(out_img[0]) * 255).astype(np.uint8))) # type: ignore else: raise RuntimeError(f"Unsupported output type `{output_type}`.") return output @@ -171,8 +172,8 @@ def save( n_row: Number of images displayed in each row of the grid. """ if directory is None: - name = f"detection-{datetime.datetime.now(tz=datetime.timezone.utc).strftime('%Y%m%d%H%M%S')!s}" - directory = os.path.join("Kornia_outputs", name) + name = f"detection_{datetime.datetime.now(tz=datetime.timezone.utc).strftime('%Y%m%d%H%M%S')!s}" + directory = os.path.join("kornia_outputs", name) outputs = self.draw(images, detections) os.makedirs(directory, exist_ok=True) for i, out_image in enumerate(outputs): diff --git a/kornia/models/detector/rtdetr.py b/kornia/models/detection/rtdetr.py similarity index 97% rename from kornia/models/detector/rtdetr.py rename to kornia/models/detection/rtdetr.py index 96ead94a89..60cc296f75 100644 --- a/kornia/models/detector/rtdetr.py +++ b/kornia/models/detection/rtdetr.py @@ -7,7 +7,7 @@ from kornia.contrib.models.rt_detr import DETRPostProcessor from kornia.contrib.models.rt_detr.model import RTDETR, RTDETRConfig from kornia.core import rand -from kornia.models.detector.base import ObjectDetector +from kornia.models.detection.base import ObjectDetector from kornia.models.utils import ResizePreProcessor __all__ = ["RTDETRDetectorBuilder"] @@ -128,10 +128,10 @@ def to_onnx( if onnx_name is None: _model_name = model_name if model_name is None and config is not None: - _model_name = "rtdetr-customized" + _model_name = "rtdetr_customized" elif model_name is None and config is None: _model_name = "rtdetr_r18vd" - onnx_name = f"Kornia-RTDETR-{_model_name}-{image_size}.onnx" + onnx_name = f"kornia_{_model_name}_{image_size}.onnx" if image_size is None: val_image = rand(1, 3, 640, 640) diff --git a/kornia/models/detector/utils.py b/kornia/models/detection/utils.py similarity index 100% rename from kornia/models/detector/utils.py rename to kornia/models/detection/utils.py diff --git a/kornia/models/edge_detector/__init__.py b/kornia/models/edge_detection/__init__.py similarity index 100% rename from kornia/models/edge_detector/__init__.py rename to kornia/models/edge_detection/__init__.py diff --git a/kornia/models/edge_detector/base.py b/kornia/models/edge_detection/base.py similarity index 100% rename from kornia/models/edge_detector/base.py rename to kornia/models/edge_detection/base.py diff --git a/kornia/models/edge_detector/dexined.py b/kornia/models/edge_detection/dexined.py similarity index 86% rename from kornia/models/edge_detector/dexined.py rename to kornia/models/edge_detection/dexined.py index 841f1f1b74..64e1bf4139 100644 --- a/kornia/models/edge_detector/dexined.py +++ b/kornia/models/edge_detection/dexined.py @@ -5,17 +5,18 @@ from kornia.core import rand from kornia.filters.dexined import DexiNed -from kornia.models.edge_detector.base import EdgeDetector +from kornia.models.edge_detection.base import EdgeDetector from kornia.models.utils import ResizePostProcessor, ResizePreProcessor class DexiNedBuilder: + @staticmethod def build(pretrained: bool = True, image_size: Optional[int] = 352) -> EdgeDetector: model = DexiNed(pretrained=pretrained) return EdgeDetector( model, - ResizePreProcessor((image_size, image_size)) if image_size is not None else nn.Identity(), + ResizePreProcessor(image_size, image_size) if image_size is not None else nn.Identity(), ResizePostProcessor() if image_size is not None else nn.Identity(), ) @@ -27,7 +28,7 @@ def to_onnx( ) -> Tuple[str, EdgeDetector]: edge_detector = DexiNedBuilder.build(pretrained, image_size) if onnx_name is None: - onnx_name = f"Kornia-DexiNed-{image_size}.onnx" + onnx_name = f"kornia_dexined_{image_size}.onnx" if image_size is None: val_image = rand(1, 3, 352, 352) diff --git a/kornia/models/segmentor/__init__.py b/kornia/models/segmentation/__init__.py similarity index 100% rename from kornia/models/segmentor/__init__.py rename to kornia/models/segmentation/__init__.py diff --git a/kornia/models/segmentor/base.py b/kornia/models/segmentation/base.py similarity index 100% rename from kornia/models/segmentor/base.py rename to kornia/models/segmentation/base.py diff --git a/kornia/models/segmentor/segmentation_models.py b/kornia/models/segmentation/segmentation_models.py similarity index 100% rename from kornia/models/segmentor/segmentation_models.py rename to kornia/models/segmentation/segmentation_models.py diff --git a/kornia/models/tracker/__init__.py b/kornia/models/tracking/__init__.py similarity index 100% rename from kornia/models/tracker/__init__.py rename to kornia/models/tracking/__init__.py diff --git a/kornia/models/tracker/boxmot_tracker.py b/kornia/models/tracking/boxmot_tracker.py similarity index 97% rename from kornia/models/tracker/boxmot_tracker.py rename to kornia/models/tracking/boxmot_tracker.py index 4cc8be105e..a9feeb25df 100644 --- a/kornia/models/tracker/boxmot_tracker.py +++ b/kornia/models/tracking/boxmot_tracker.py @@ -5,8 +5,8 @@ from kornia.core import Tensor from kornia.core.external import boxmot from kornia.core.external import numpy as np -from kornia.models.detector.base import ObjectDetector -from kornia.models.detector.rtdetr import RTDETRDetectorBuilder +from kornia.models.detection.base import ObjectDetector +from kornia.models.detection.rtdetr import RTDETRDetectorBuilder from kornia.utils.image import tensor_to_image __all__ = ["BoxMotTracker"] diff --git a/kornia/models/utils.py b/kornia/models/utils.py index 08676160c2..f0f4931531 100644 --- a/kornia/models/utils.py +++ b/kornia/models/utils.py @@ -14,16 +14,16 @@ class ResizePreProcessor(Module): Additionally, also returns the original image sizes for further post-processing. """ - def __init__(self, size: tuple[int, int], interpolation_mode: str = "bilinear") -> None: + def __init__(self, height: int, width: int, interpolation_mode: str = "bilinear") -> None: """ Args: - size: images will be resized to this value. If a 2-integer tuple is given, it is interpreted as - (height, width). + height: height of the resized image. + width: width of the resized image. interpolation_mode: interpolation mode for image resizing. Supported values: ``nearest``, ``bilinear``, ``bicubic``, ``area``, and ``nearest-exact``. """ super().__init__() - self.size = size + self.size = (height, width) self.interpolation_mode = interpolation_mode def forward(self, imgs: Union[Tensor, list[Tensor]]) -> tuple[Tensor, Tensor]: diff --git a/kornia/onnx/sequential.py b/kornia/onnx/sequential.py index c934a0550e..155fada053 100644 --- a/kornia/onnx/sequential.py +++ b/kornia/onnx/sequential.py @@ -3,6 +3,7 @@ from kornia.core.external import numpy as np from kornia.core.external import onnx from kornia.core.external import onnxruntime as ort +from kornia.config import kornia_config from .utils import ONNXLoader @@ -10,7 +11,7 @@ class ONNXSequential: - """ONNXSequential to chain multiple ONNX operators together. + f"""ONNXSequential to chain multiple ONNX operators together. Args: *args: A variable number of ONNX models (either ONNX ModelProto objects or file paths). @@ -24,7 +25,7 @@ class ONNXSequential: only one input and output node for each graph. If not None, `io_maps[0]` shall represent the `io_map` for combining the first and second ONNX models. cache_dir: The directory where ONNX models are cached locally (only for downloading from HuggingFace). - Defaults to None, which will use a default `.kornia_hub/onnx_models` directory. + Defaults to None, which will use a default `{kornia_config.hub_onnx_dir}` directory. """ def __init__( diff --git a/kornia/onnx/utils.py b/kornia/onnx/utils.py index 98cd774f30..669cbce79a 100644 --- a/kornia/onnx/utils.py +++ b/kornia/onnx/utils.py @@ -7,18 +7,18 @@ import requests from kornia.core.external import onnx - +from kornia.config import kornia_config __all__ = ["ONNXLoader"] logger = logging.getLogger(__name__) class ONNXLoader: - """Manages ONNX models, handling local caching, downloading from Hugging Face, and loading models. + f"""Manages ONNX models, handling local caching, downloading from Hugging Face, and loading models. Attributes: cache_dir: The directory where ONNX models are cached locally. - Defaults to None, which will use a default `.kornia_hub/onnx_models` directory. + Defaults to None, which will use a default `{kornia_config.hub_onnx_dir}` directory. """ def __init__(self, cache_dir: Optional[str] = None): @@ -39,7 +39,7 @@ def _get_file_path(self, model_name: str, cache_dir: Optional[str]) -> str: if self.cache_dir is not None: cache_dir = self.cache_dir else: - cache_dir = ".kornia_hub/onnx_models" + cache_dir = kornia_config.hub_onnx_dir # The filename is the model name (without directory path) file_name = f"{model_name.split('/')[-1]}.onnx" @@ -61,7 +61,7 @@ def load_model(self, model_name: str, download: bool = True, **kwargs) -> "onnx. """ if model_name.startswith("hf://"): model_name = model_name[len("hf://") :] - cache_dir = kwargs.get("cache_dir", None) or self.cache_dir + cache_dir = kwargs.get(kornia_config.hub_onnx_dir, None) or self.cache_dir file_path = self._get_file_path(model_name, cache_dir) if not os.path.exists(file_path): # Construct the raw URL for the ONNX file From d8a370e1e096aac21f93f8537b83f2481625ffe9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 15 Sep 2024 13:20:03 +0000 Subject: [PATCH 63/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/config.py | 2 +- kornia/core/external.py | 2 +- kornia/models/edge_detection/dexined.py | 1 - kornia/onnx/sequential.py | 2 +- kornia/onnx/utils.py | 3 ++- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/kornia/config.py b/kornia/config.py index 857a566e28..b993465e23 100644 --- a/kornia/config.py +++ b/kornia/config.py @@ -1,5 +1,5 @@ -from enum import Enum from dataclasses import dataclass, field +from enum import Enum __all__ = ["config", "InstallationMode"] diff --git a/kornia/core/external.py b/kornia/core/external.py index 39c0b789d1..2596ad0271 100644 --- a/kornia/core/external.py +++ b/kornia/core/external.py @@ -5,7 +5,7 @@ from types import ModuleType from typing import List, Optional -from kornia.config import kornia_config, InstallationMode +from kornia.config import InstallationMode, kornia_config logger = logging.getLogger(__name__) diff --git a/kornia/models/edge_detection/dexined.py b/kornia/models/edge_detection/dexined.py index 64e1bf4139..c4cab95cf6 100644 --- a/kornia/models/edge_detection/dexined.py +++ b/kornia/models/edge_detection/dexined.py @@ -10,7 +10,6 @@ class DexiNedBuilder: - @staticmethod def build(pretrained: bool = True, image_size: Optional[int] = 352) -> EdgeDetector: model = DexiNed(pretrained=pretrained) diff --git a/kornia/onnx/sequential.py b/kornia/onnx/sequential.py index 155fada053..0cbadc8eaa 100644 --- a/kornia/onnx/sequential.py +++ b/kornia/onnx/sequential.py @@ -1,9 +1,9 @@ from typing import List, Optional, Tuple, Union +from kornia.config import kornia_config from kornia.core.external import numpy as np from kornia.core.external import onnx from kornia.core.external import onnxruntime as ort -from kornia.config import kornia_config from .utils import ONNXLoader diff --git a/kornia/onnx/utils.py b/kornia/onnx/utils.py index 669cbce79a..51977bdf49 100644 --- a/kornia/onnx/utils.py +++ b/kornia/onnx/utils.py @@ -6,8 +6,9 @@ import requests -from kornia.core.external import onnx from kornia.config import kornia_config +from kornia.core.external import onnx + __all__ = ["ONNXLoader"] logger = logging.getLogger(__name__) From 01b53a9e028a4b4f519e326a5486f7eb33e56b00 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Sun, 15 Sep 2024 17:32:02 +0300 Subject: [PATCH 64/84] update --- kornia/models/base.py | 23 ++++++ kornia/models/detection/base.py | 69 ++++++++++++---- kornia/models/detection/rtdetr.py | 78 +------------------ kornia/models/edge_detection/base.py | 22 +----- kornia/models/edge_detection/dexined.py | 13 +++- kornia/models/segmentation/base.py | 21 ++--- .../segmentation/segmentation_models.py | 6 +- requirements/requirements-dev.txt | 2 +- 8 files changed, 100 insertions(+), 134 deletions(-) create mode 100644 kornia/models/base.py diff --git a/kornia/models/base.py b/kornia/models/base.py new file mode 100644 index 0000000000..8712b5dc3f --- /dev/null +++ b/kornia/models/base.py @@ -0,0 +1,23 @@ +from kornia.core import Module +from typing import Optional + + +class ModelBase(Module): + """This class wraps a model and performs pre-processing and post-processing.""" + + name: str = "model" + + def __init__(self, model: Module, pre_processor: Module, post_processor: Module, name: Optional[str] = None) -> None: + """Construct an Object Detector object. + + Args: + model: an object detection model. + pre_processor: a pre-processing module + post_processor: a post-processing module. + """ + super().__init__() + self.model = model.eval() + self.pre_processor = pre_processor.eval() + self.post_processor = post_processor.eval() + if name is not None: + self.name = name diff --git a/kornia/models/detection/base.py b/kornia/models/detection/base.py index 744fc579d2..7bb7b66c53 100644 --- a/kornia/models/detection/base.py +++ b/kornia/models/detection/base.py @@ -5,17 +5,18 @@ import os from dataclasses import dataclass from enum import Enum -from typing import Optional, Union +from typing import Optional, Tuple, Union import torch -from kornia.core import Module, Tensor +from kornia.core import Tensor, rand from kornia.core.check import KORNIA_CHECK_SHAPE from kornia.core.external import PILImage as Image from kornia.core.external import numpy as np from kornia.io import write_image from kornia.utils.draw import draw_rectangle from kornia.utils.image import tensor_to_image +from kornia.models.base import ModelBase __all__ = [ "BoundingBoxDataFormat", @@ -103,22 +104,10 @@ def results_from_detections(detections: Tensor, format: str | BoundingBoxDataFor return results -# TODO: move this to kornia.models as AlgorithmicModel api -class ObjectDetector(Module): +class ObjectDetector(ModelBase): """This class wraps an object detection model and performs pre-processing and post-processing.""" - def __init__(self, model: Module, pre_processor: Module, post_processor: Module) -> None: - """Construct an Object Detector object. - - Args: - model: an object detection model. - pre_processor: a pre-processing module - post_processor: a post-processing module. - """ - super().__init__() - self.model = model.eval() - self.pre_processor = pre_processor.eval() - self.post_processor = post_processor.eval() + name: str = "ObjectDetector" @torch.inference_mode() def forward(self, images: Union[Tensor, list[Tensor]]) -> Tensor: @@ -183,6 +172,54 @@ def save( ) logger.info(f"Outputs are saved in {directory}") + def to_onnx( + self, + onnx_name: Optional[str] = None, + image_size: Optional[int] = 640, + include_pre_and_post_processor: bool = True, + ) -> Tuple[str, ObjectDetector]: + """Exports an RT-DETR object detection model to ONNX format. + + Either `model_name` or `config` must be provided. If neither is provided, + a default pretrained model (`rtdetr_r18vd`) will be built. + + Args: + onnx_name: + The name of the ONNX model. + image_size: + The size to which input images will be resized during preprocessing. + If None, image_size will be dynamic. + For RTDETR, recommended scales include [480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800]. + include_pre_and_post_processor: + Whether to include the pre-processor and post-processor in the exported model. + + Returns: + - The name of the ONNX model. + """ + if onnx_name is None: + onnx_name = f"kornia_{self.name}_{image_size}.onnx" + + if image_size is None: + val_image = rand(1, 3, 640, 640) + dynamic_axes = {"input": {0: "batch_size"}, "output": {0: "batch_size"}} + else: + val_image = rand(1, 3, image_size, image_size) + dynamic_axes = {"input": {0: "batch_size", 2: "height", 3: "width"}, "output": {0: "batch_size"}} + + torch.onnx.export( + self if include_pre_and_post_processor else self.model, + val_image, + onnx_name, + export_params=True, + opset_version=17, + do_constant_folding=True, + input_names=["input"], + output_names=["output"], + dynamic_axes=dynamic_axes, + ) + + return onnx_name + def compile( self, *, diff --git a/kornia/models/detection/rtdetr.py b/kornia/models/detection/rtdetr.py index 60cc296f75..99dc466fd8 100644 --- a/kornia/models/detection/rtdetr.py +++ b/kornia/models/detection/rtdetr.py @@ -1,12 +1,11 @@ import warnings -from typing import Optional, Tuple +from typing import Optional import torch from torch import nn from kornia.contrib.models.rt_detr import DETRPostProcessor from kornia.contrib.models.rt_detr.model import RTDETR, RTDETRConfig -from kornia.core import rand from kornia.models.detection.base import ObjectDetector from kornia.models.utils import ResizePreProcessor @@ -71,7 +70,7 @@ def build( return ObjectDetector( model, - ResizePreProcessor((image_size, image_size)) if image_size is not None else nn.Identity(), + ResizePreProcessor(image_size, image_size) if image_size is not None else nn.Identity(), DETRPostProcessor( confidence_threshold=confidence_threshold, confidence_filtering=confidence_filtering or not torch.onnx.is_in_onnx_export(), @@ -79,76 +78,3 @@ def build( num_top_queries=model.decoder.num_queries, ), ) - - @staticmethod - def to_onnx( - model_name: Optional[str] = None, - onnx_name: Optional[str] = None, - config: Optional[RTDETRConfig] = None, - pretrained: bool = True, - image_size: Optional[int] = 640, - confidence_threshold: Optional[float] = None, - confidence_filtering: Optional[bool] = None, - ) -> Tuple[str, ObjectDetector]: - """Exports an RT-DETR object detection model to ONNX format. - - Either `model_name` or `config` must be provided. If neither is provided, - a default pretrained model (`rtdetr_r18vd`) will be built. - - Args: - model_name: - Name of the RT-DETR model to load. Can be one of the available pretrained models. - config: - A custom configuration object for building the RT-DETR model. - pretrained: - Whether to load a pretrained version of the model (applies when `model_name` is provided). - image_size: - The size to which input images will be resized during preprocessing. - If None, image_size will be dynamic. Recommended scales include - [480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800]. - confidence_threshold: - The confidence threshold used during post-processing to filter detections. - confidence_filtering: - If to perform filtering on resulting boxes. If None, the filtering will be blocked when exporting - to ONNX, while it would perform as per confidence_threshold when build the model. - - Returns: - - The name of the ONNX model. - - The exported torch model. - """ - - detector = RTDETRDetectorBuilder.build( - model_name=model_name, - config=config, - pretrained=pretrained, - image_size=image_size, - confidence_threshold=confidence_threshold, - confidence_filtering=confidence_filtering, - ) - if onnx_name is None: - _model_name = model_name - if model_name is None and config is not None: - _model_name = "rtdetr_customized" - elif model_name is None and config is None: - _model_name = "rtdetr_r18vd" - onnx_name = f"kornia_{_model_name}_{image_size}.onnx" - - if image_size is None: - val_image = rand(1, 3, 640, 640) - else: - val_image = rand(1, 3, image_size, image_size) - - dynamic_axes = {"input": {0: "batch_size", 2: "height", 3: "width"}, "output": {0: "batch_size"}} - torch.onnx.export( - detector, - val_image, - onnx_name, - export_params=True, - opset_version=17, - do_constant_folding=True, - input_names=["input"], - output_names=["output"], - dynamic_axes=dynamic_axes, - ) - - return onnx_name, detector diff --git a/kornia/models/edge_detection/base.py b/kornia/models/edge_detection/base.py index 1ca3477a80..8c1ff02eb3 100644 --- a/kornia/models/edge_detection/base.py +++ b/kornia/models/edge_detection/base.py @@ -1,36 +1,22 @@ from typing import Union -from kornia.core import Module, Tensor +from kornia.core import Tensor from kornia.core.external import PILImage as Image - +from kornia.models.base import ModelBase __all__ = ["EdgeDetector"] -class EdgeDetector(Module): +class EdgeDetector(ModelBase): """EdgeDetector is a module that wraps an edge detection model. This module uses EdgeDetectionModel library for edge detection. """ - # NOTE: We need to implement this class for better visualization and user experience. - - def __init__(self, model: Module, pre_processor: Module, post_processor: Module) -> None: - """Construct an Object Detector object. - - Args: - model: an object detection model. - pre_processor: a pre-processing module - """ - super().__init__() - self.model = model - self.pre_processor = pre_processor - self.post_processor = post_processor - def forward(self, images: Union[Tensor, list[Tensor]]) -> Tensor: """Forward pass of the semantic segmentation model. Args: - x: input tensor. + images: input tensor. Returns: output tensor. diff --git a/kornia/models/edge_detection/dexined.py b/kornia/models/edge_detection/dexined.py index c4cab95cf6..6b20a9dd41 100644 --- a/kornia/models/edge_detection/dexined.py +++ b/kornia/models/edge_detection/dexined.py @@ -11,8 +11,12 @@ class DexiNedBuilder: @staticmethod - def build(pretrained: bool = True, image_size: Optional[int] = 352) -> EdgeDetector: - model = DexiNed(pretrained=pretrained) + def build(model_name: str = "dexined", pretrained: bool = True, image_size: Optional[int] = 352) -> EdgeDetector: + if model_name.lower() == "dexined": + model = DexiNed(pretrained=pretrained) + else: + raise ValueError(f"Model {model_name} not found. Please choose from 'DexiNed'.") + return EdgeDetector( model, ResizePreProcessor(image_size, image_size) if image_size is not None else nn.Identity(), @@ -21,13 +25,14 @@ def build(pretrained: bool = True, image_size: Optional[int] = 352) -> EdgeDetec @staticmethod def to_onnx( + model_name: str = "dexined", onnx_name: Optional[str] = None, pretrained: bool = True, image_size: Optional[int] = 352, ) -> Tuple[str, EdgeDetector]: - edge_detector = DexiNedBuilder.build(pretrained, image_size) + edge_detector = DexiNedBuilder.build(model_name, pretrained, image_size) if onnx_name is None: - onnx_name = f"kornia_dexined_{image_size}.onnx" + onnx_name = f"kornia_{model_name.lower()}_{image_size}.onnx" if image_size is None: val_image = rand(1, 3, 352, 352) diff --git a/kornia/models/segmentation/base.py b/kornia/models/segmentation/base.py index ec6293f40c..6dceee9252 100644 --- a/kornia/models/segmentation/base.py +++ b/kornia/models/segmentation/base.py @@ -1,28 +1,16 @@ from typing import Union -from kornia.core import Module, Tensor +from kornia.core import Tensor from kornia.core.external import PILImage as Image +from kornia.models.base import ModelBase -class SemanticSegmentation(Module): +class SemanticSegmentation(ModelBase): """Semantic Segmentation is a module that wraps a semantic segmentation model. This module uses SegmentationModel library for semantic segmentation. """ - # NOTE: We need to implement this class for better visualization and user experience. - - def __init__(self, model: Module, pre_processor: Module) -> None: - """Construct an Object Detector object. - - Args: - model: an object detection model. - pre_processor: a pre-processing module - """ - super().__init__() - self.model = model - self.pre_processor = pre_processor - def forward(self, images: Union[Tensor, list[Tensor]]) -> Tensor: """Forward pass of the semantic segmentation model. @@ -33,7 +21,8 @@ def forward(self, images: Union[Tensor, list[Tensor]]) -> Tensor: output tensor. """ images = self.pre_processor(images) - return self.model(images) + output = self.model(images) + return self.post_processor(output) def draw( self, images: Union[Tensor, list[Tensor]], output_type: str = "torch" diff --git a/kornia/models/segmentation/segmentation_models.py b/kornia/models/segmentation/segmentation_models.py index d50674e05b..96e8a8ed7a 100644 --- a/kornia/models/segmentation/segmentation_models.py +++ b/kornia/models/segmentation/segmentation_models.py @@ -54,7 +54,7 @@ def preprocessing(self, input: Tensor) -> Tensor: # Ensure the color space transformation is ONNX-friendly input_space = self.preproc_params["input_space"] input = ( - kornia.color.rgb_to_bgr(input) if input_space == "BGR" else input + kornia.color.rgb.rgb_to_bgr(input) if input_space == "BGR" else input ) # Assume input is already RGB if not BGR # Normalize input range if needed @@ -68,12 +68,12 @@ def preprocessing(self, input: Tensor) -> Tensor: # Handle mean and std normalization if self.preproc_params["mean"] is not None: - mean = tensor([self.preproc_params["mean"]]).to(input.device) + mean = tensor([self.preproc_params["mean"]], device=input.device) else: mean = zeros_like(input) if self.preproc_params["std"] is not None: - std = tensor([self.preproc_params["std"]]).to(input.device) + std = tensor([self.preproc_params["std"]], device=input.device) else: std = ones_like(input) diff --git a/requirements/requirements-dev.txt b/requirements/requirements-dev.txt index f52cc1aa94..52c1bdb8f1 100644 --- a/requirements/requirements-dev.txt +++ b/requirements/requirements-dev.txt @@ -3,7 +3,7 @@ boxmot coverage diffusers mypy -numpy<2 +numpy>1.24.4,<2 onnx onnxruntime pillow From 6c0ef12451313671e7f776f5a2d99ea3080dad22 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 15 Sep 2024 14:32:37 +0000 Subject: [PATCH 65/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/models/base.py | 7 +++++-- kornia/models/detection/base.py | 2 +- kornia/models/edge_detection/base.py | 1 + 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/kornia/models/base.py b/kornia/models/base.py index 8712b5dc3f..89c05a1535 100644 --- a/kornia/models/base.py +++ b/kornia/models/base.py @@ -1,13 +1,16 @@ -from kornia.core import Module from typing import Optional +from kornia.core import Module + class ModelBase(Module): """This class wraps a model and performs pre-processing and post-processing.""" name: str = "model" - def __init__(self, model: Module, pre_processor: Module, post_processor: Module, name: Optional[str] = None) -> None: + def __init__( + self, model: Module, pre_processor: Module, post_processor: Module, name: Optional[str] = None + ) -> None: """Construct an Object Detector object. Args: diff --git a/kornia/models/detection/base.py b/kornia/models/detection/base.py index 7bb7b66c53..96fb080b48 100644 --- a/kornia/models/detection/base.py +++ b/kornia/models/detection/base.py @@ -14,9 +14,9 @@ from kornia.core.external import PILImage as Image from kornia.core.external import numpy as np from kornia.io import write_image +from kornia.models.base import ModelBase from kornia.utils.draw import draw_rectangle from kornia.utils.image import tensor_to_image -from kornia.models.base import ModelBase __all__ = [ "BoundingBoxDataFormat", diff --git a/kornia/models/edge_detection/base.py b/kornia/models/edge_detection/base.py index 8c1ff02eb3..1b345979f0 100644 --- a/kornia/models/edge_detection/base.py +++ b/kornia/models/edge_detection/base.py @@ -3,6 +3,7 @@ from kornia.core import Tensor from kornia.core.external import PILImage as Image from kornia.models.base import ModelBase + __all__ = ["EdgeDetector"] From 8cb9bf3320bd409bbe734a6f1974909eb1c4dfa7 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Mon, 16 Sep 2024 00:50:51 +0300 Subject: [PATCH 66/84] update --- docs/source/index.rst | 30 +++++++++++++ kornia/config.py | 2 +- kornia/enhance/normalize.py | 4 +- kornia/filters/dexined.py | 6 +-- kornia/models/base.py | 46 +++++++++++++++++++- kornia/models/detection/base.py | 41 +++++------------- kornia/models/detection/rtdetr.py | 7 +++ kornia/models/edge_detection/base.py | 29 ++++++++++--- kornia/models/edge_detection/dexined.py | 17 +++++++- kornia/models/segmentation/base.py | 8 ++-- kornia/models/tracking/boxmot_tracker.py | 55 ++++++++++++++++++++---- kornia/models/utils.py | 2 +- requirements/requirements-dev.txt | 2 +- 13 files changed, 188 insertions(+), 61 deletions(-) diff --git a/docs/source/index.rst b/docs/source/index.rst index 4f18cee559..b3e8287a7c 100644 --- a/docs/source/index.rst +++ b/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,6 +95,7 @@ Join the community io image losses + models metrics morphology nerf diff --git a/kornia/config.py b/kornia/config.py index b993465e23..45688bb779 100644 --- a/kornia/config.py +++ b/kornia/config.py @@ -1,7 +1,7 @@ from dataclasses import dataclass, field from enum import Enum -__all__ = ["config", "InstallationMode"] +__all__ = ["kornia_config", "InstallationMode"] class InstallationMode(str, Enum): diff --git a/kornia/enhance/normalize.py b/kornia/enhance/normalize.py index a291fa58cf..f46748edaa 100644 --- a/kornia/enhance/normalize.py +++ b/kornia/enhance/normalize.py @@ -54,10 +54,10 @@ def __init__( std = torch.tensor([std]) if isinstance(mean, (tuple, list)): - mean = torch.tensor(mean) + mean = torch.tensor(mean)[None] if isinstance(std, (tuple, list)): - std = torch.tensor(std) + std = torch.tensor(std)[None] self.mean = mean self.std = std diff --git a/kornia/filters/dexined.py b/kornia/filters/dexined.py index 12bfac591c..1f77393253 100644 --- a/kornia/filters/dexined.py +++ b/kornia/filters/dexined.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections import OrderedDict -from typing import Optional +from typing import ClassVar, Optional, List import torch import torch.nn.functional as F @@ -179,8 +179,8 @@ class DexiNed(Module): torch.Size([1, 1, 320, 320]) """ - # TODO: Handle multiple inputs and outputs models later - ONNX_EXPORTABLE = False + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 1, -1, -1] def __init__(self, pretrained: bool) -> None: super().__init__() diff --git a/kornia/models/base.py b/kornia/models/base.py index 89c05a1535..7ed6f12ef9 100644 --- a/kornia/models/base.py +++ b/kornia/models/base.py @@ -1,4 +1,15 @@ -from typing import Optional +import os +import datetime +import logging +from typing import Optional, Union + +from kornia.core import Module, Tensor, stack +from kornia.core.external import PILImage as Image +from kornia.core.external import numpy as np +from kornia.utils.image import tensor_to_image +from kornia.io import write_image + +logger = logging.getLogger(__name__) from kornia.core import Module @@ -24,3 +35,36 @@ def __init__( self.post_processor = post_processor.eval() if name is not None: self.name = name + + def _tensor_to_type( + self, output: list[Tensor], output_type: str, is_batch: bool = False + ) -> Union[Tensor, list[Tensor], list[Image.Image]]: # type: ignore + if output_type == "torch": + if is_batch: + return stack(output) + elif output_type == "pil": + return [Image.fromarray((tensor_to_image(out_img) * 255).astype(np.uint8)) for out_img in output] # type: ignore + + raise RuntimeError(f"Unsupported output type `{output_type}`.") + + def _save_outputs( + self, outputs: Union[Tensor, list[Tensor]], directory: Optional[str] = None, + suffix: str = "" + ) -> None: + """Save the output image(s) to a directory. + + Args: + outputs: output tensor. + directory: directory to save the images. + """ + if directory is None: + name = f"{self.name}_{datetime.datetime.now(tz=datetime.timezone.utc).strftime('%Y%m%d%H%M%S')!s}" + directory = os.path.join("kornia_outputs", name) + + os.makedirs(directory, exist_ok=True) + for i, out_image in enumerate(outputs): + write_image( + os.path.join(directory, f"{str(i).zfill(6)}{suffix}.jpg"), + out_image.mul(255.0).byte(), + ) + logger.info(f"Outputs are saved in {directory}") \ No newline at end of file diff --git a/kornia/models/detection/base.py b/kornia/models/detection/base.py index 96fb080b48..2f8911a854 100644 --- a/kornia/models/detection/base.py +++ b/kornia/models/detection/base.py @@ -1,8 +1,5 @@ from __future__ import annotations -import datetime -import logging -import os from dataclasses import dataclass from enum import Enum from typing import Optional, Tuple, Union @@ -12,11 +9,8 @@ from kornia.core import Tensor, rand from kornia.core.check import KORNIA_CHECK_SHAPE from kornia.core.external import PILImage as Image -from kornia.core.external import numpy as np -from kornia.io import write_image from kornia.models.base import ModelBase from kornia.utils.draw import draw_rectangle -from kornia.utils.image import tensor_to_image __all__ = [ "BoundingBoxDataFormat", @@ -26,7 +20,6 @@ "ObjectDetectorResult", ] -logger = logging.getLogger(__name__) class BoundingBoxDataFormat(Enum): @@ -107,10 +100,10 @@ def results_from_detections(detections: Tensor, format: str | BoundingBoxDataFor class ObjectDetector(ModelBase): """This class wraps an object detection model and performs pre-processing and post-processing.""" - name: str = "ObjectDetector" + name: str = "detection" @torch.inference_mode() - def forward(self, images: Union[Tensor, list[Tensor]]) -> Tensor: + def forward(self, images: Union[Tensor, list[Tensor]]) -> Union[Tensor, list[Tensor]]: """Detect objects in a given list of images. Args: @@ -126,7 +119,7 @@ def forward(self, images: Union[Tensor, list[Tensor]]) -> Tensor: detections = self.post_processor(logits, boxes, images_sizes) return detections - def draw( + def visualize( self, images: Union[Tensor, list[Tensor]], detections: Optional[Tensor] = None, output_type: str = "torch" ) -> Union[Tensor, list[Tensor], list[Image.Image]]: # type: ignore """Very simple drawing. @@ -143,13 +136,9 @@ def draw( out_img, torch.Tensor([[[out[-4], out[-3], out[-4] + out[-2], out[-3] + out[-1]]]]), ) - if output_type == "torch": - output.append(out_img[0]) - elif output_type == "pil": - output.append(Image.fromarray((tensor_to_image(out_img[0]) * 255).astype(np.uint8))) # type: ignore - else: - raise RuntimeError(f"Unsupported output type `{output_type}`.") - return output + output.append(out_img[0]) + + return self._tensor_to_type(output, output_type, is_batch=isinstance(images, Tensor)) def save( self, images: Union[Tensor, list[Tensor]], detections: Optional[Tensor] = None, directory: Optional[str] = None @@ -157,20 +146,12 @@ def save( """Saves the output image(s) to a directory. Args: - name: Directory to save the images. - n_row: Number of images displayed in each row of the grid. + images: input tensor. + detections: detection tensor. + directory: directory to save the images. """ - if directory is None: - name = f"detection_{datetime.datetime.now(tz=datetime.timezone.utc).strftime('%Y%m%d%H%M%S')!s}" - directory = os.path.join("kornia_outputs", name) - outputs = self.draw(images, detections) - os.makedirs(directory, exist_ok=True) - for i, out_image in enumerate(outputs): - write_image( - os.path.join(directory, f"{str(i).zfill(6)}.jpg"), - out_image.mul(255.0).byte(), - ) - logger.info(f"Outputs are saved in {directory}") + outputs = self.visualize(images, detections) + self._save_outputs(outputs, directory) def to_onnx( self, diff --git a/kornia/models/detection/rtdetr.py b/kornia/models/detection/rtdetr.py index 99dc466fd8..9f54ddca78 100644 --- a/kornia/models/detection/rtdetr.py +++ b/kornia/models/detection/rtdetr.py @@ -18,6 +18,13 @@ class RTDETRDetectorBuilder: This class provides static methods to: - Build an object detection model from a model name or configuration. - Export the model to ONNX format for inference. + + .. code-block:: python + + image = kornia.utils.sample.get_sample_images()[0][None] + model = RTDETRDetectorBuilder.build() + model.save(image) + """ @staticmethod diff --git a/kornia/models/edge_detection/base.py b/kornia/models/edge_detection/base.py index 1b345979f0..9fff5ec35c 100644 --- a/kornia/models/edge_detection/base.py +++ b/kornia/models/edge_detection/base.py @@ -1,5 +1,6 @@ -from typing import Union +from typing import Union, Optional +from kornia.color.gray import grayscale_to_rgb from kornia.core import Tensor from kornia.core.external import PILImage as Image from kornia.models.base import ModelBase @@ -13,7 +14,9 @@ class EdgeDetector(ModelBase): This module uses EdgeDetectionModel library for edge detection. """ - def forward(self, images: Union[Tensor, list[Tensor]]) -> Tensor: + name: str = "edge_detection" + + def forward(self, images: Union[Tensor, list[Tensor]]) -> Union[Tensor, list[Tensor]]: """Forward pass of the semantic segmentation model. Args: @@ -26,8 +29,9 @@ def forward(self, images: Union[Tensor, list[Tensor]]) -> Tensor: out_images = self.model(images) return self.post_processor(out_images, image_sizes) - def draw( - self, images: Union[Tensor, list[Tensor]], output_type: str = "torch" + def visualize( + self, images: Union[Tensor, list[Tensor]], edge_maps: Optional[Union[Tensor, list[Tensor]]] = None, + output_type: str = "torch" ) -> Union[Tensor, list[Tensor], list[Image.Image]]: # type: ignore """Draw the segmentation results. @@ -38,9 +42,18 @@ def draw( Returns: output tensor. """ - ... + if edge_maps is None: + edge_maps = self.forward(images) + output = [] + for edge_map in edge_maps: + output.append(grayscale_to_rgb(edge_map)[0]) + + return self._tensor_to_type(output, output_type, is_batch=isinstance(images, Tensor)) - def save(self, images: Union[Tensor, list[Tensor]], output_type: str = "torch") -> None: + def save( + self, images: Union[Tensor, list[Tensor]], edge_maps: Optional[Union[Tensor, list[Tensor]]] = None, + directory: Optional[str] = None, output_type: str = "torch" + ) -> None: """Save the segmentation results. Args: @@ -50,4 +63,6 @@ def save(self, images: Union[Tensor, list[Tensor]], output_type: str = "torch") Returns: output tensor. """ - ... + outputs = self.visualize(images, edge_maps, output_type) + self._save_outputs(images, directory, suffix="_src") + self._save_outputs(outputs, directory, suffix="_edge") diff --git a/kornia/models/edge_detection/dexined.py b/kornia/models/edge_detection/dexined.py index 6b20a9dd41..8b2223e8c0 100644 --- a/kornia/models/edge_detection/dexined.py +++ b/kornia/models/edge_detection/dexined.py @@ -3,17 +3,30 @@ import torch from torch import nn -from kornia.core import rand +from kornia.core import rand, tensor +from kornia.enhance.normalize import Normalize from kornia.filters.dexined import DexiNed from kornia.models.edge_detection.base import EdgeDetector from kornia.models.utils import ResizePostProcessor, ResizePreProcessor class DexiNedBuilder: + """DexiNedBuilder is a class that builds a DexiNed model. + + .. code-block:: python + + image = kornia.utils.sample.get_sample_images()[0][None] + model = DexiNedBuilder.build() + model.save(image) + + """ + @staticmethod def build(model_name: str = "dexined", pretrained: bool = True, image_size: Optional[int] = 352) -> EdgeDetector: if model_name.lower() == "dexined": - model = DexiNed(pretrained=pretrained) + # Normalize then scale to [0, 255] + norm = Normalize(mean=tensor([[0.485, 0.456, 0.406]]), std=tensor([[1. / 255.] * 3])) + model = nn.Sequential(norm, DexiNed(pretrained=pretrained), nn.Sigmoid()) else: raise ValueError(f"Model {model_name} not found. Please choose from 'DexiNed'.") diff --git a/kornia/models/segmentation/base.py b/kornia/models/segmentation/base.py index 6dceee9252..053eb5c42b 100644 --- a/kornia/models/segmentation/base.py +++ b/kornia/models/segmentation/base.py @@ -11,7 +11,7 @@ class SemanticSegmentation(ModelBase): This module uses SegmentationModel library for semantic segmentation. """ - def forward(self, images: Union[Tensor, list[Tensor]]) -> Tensor: + def forward(self, images: Union[Tensor, list[Tensor]]) -> Union[Tensor, list[Tensor]]: """Forward pass of the semantic segmentation model. Args: @@ -24,7 +24,7 @@ def forward(self, images: Union[Tensor, list[Tensor]]) -> Tensor: output = self.model(images) return self.post_processor(output) - def draw( + def visualize( self, images: Union[Tensor, list[Tensor]], output_type: str = "torch" ) -> Union[Tensor, list[Tensor], list[Image.Image]]: # type: ignore """Draw the segmentation results. @@ -36,7 +36,7 @@ def draw( Returns: output tensor. """ - ... + raise NotImplementedError("Visualization is not implemented for this model.") def save(self, images: Union[Tensor, list[Tensor]], output_type: str = "torch") -> None: """Save the segmentation results. @@ -48,4 +48,4 @@ def save(self, images: Union[Tensor, list[Tensor]], output_type: str = "torch") Returns: output tensor. """ - ... + raise NotImplementedError("Saving is not implemented for this model.") diff --git a/kornia/models/tracking/boxmot_tracker.py b/kornia/models/tracking/boxmot_tracker.py index a9feeb25df..af4533fa40 100644 --- a/kornia/models/tracking/boxmot_tracker.py +++ b/kornia/models/tracking/boxmot_tracker.py @@ -1,16 +1,22 @@ import os +import datetime +import logging from pathlib import Path -from typing import Union +from typing import Union, Optional -from kornia.core import Tensor +from kornia.core import Tensor, tensor from kornia.core.external import boxmot from kornia.core.external import numpy as np from kornia.models.detection.base import ObjectDetector from kornia.models.detection.rtdetr import RTDETRDetectorBuilder from kornia.utils.image import tensor_to_image +from kornia.config import kornia_config +from kornia.io import write_image __all__ = ["BoxMotTracker"] +logger = logging.getLogger(__name__) + class BoxMotTracker: """BoxMotTracker is a module that wraps a detector and a tracker model. @@ -45,12 +51,26 @@ class BoxMotTracker: frame_rate: Frame rate of the video being processed. Used to scale the track buffer size. fuse_first_associate: Whether to fuse appearance and motion information during the first association step. with_reid: Whether to use ReID (Re-Identification) features for association. + + .. code-block:: python + + import kornia + image = kornia.utils.sample.get_sample_images()[0][None] + model = BoxMotTracker() + for i in range(4): # At least 4 frames are needed to initialize the tracking position + model.update(image) + model.save(image) + + .. note:: + At least 4 frames are needed to initialize the tracking position. """ + name: str = "boxmot_tracker" + def __init__( self, detector: Union[ObjectDetector, str] = "rtdetr_r18vd", - tracker_model_name: str = "BoTSORT", + tracker_model_name: str = "DeepOCSORT", tracker_model_weights: str = "osnet_x0_25_msmt17.pt", device: str = "cpu", fp16: bool = False, @@ -65,9 +85,9 @@ def __init__( f"Detector `{detector}` not available. You may pass an ObjectDetector instance instead." ) self.detector = detector - os.makedirs(".kornia_hub/models/boxmot", exist_ok=True) + os.makedirs(f"{kornia_config.hub_models_dir}/boxmot", exist_ok=True) self.tracker = getattr(boxmot, tracker_model_name)( - model_weights=Path(os.path.join(".kornia_hub/models/boxmot", tracker_model_weights)), + model_weights=Path(os.path.join(f"{kornia_config.hub_models_dir}/boxmot", tracker_model_weights)), device=device, fp16=fp16, **kwargs, @@ -106,11 +126,10 @@ def update(self, image: Tensor) -> None: detections = np.empty((0, 6)) # type: ignore frame_raw = (tensor_to_image(image) * 255).astype(np.uint8) # type: ignore - # --> M X (x, y, x, y, id, conf, cls, ind) return self.tracker.update(detections, frame_raw) # type: ignore - def visualize(self, image: Tensor, show_trajectories: bool = True) -> np.ndarray: # type: ignore + def visualize(self, image: Tensor, show_trajectories: bool = True) -> Tensor: """Visualize the results of the tracker. Args: @@ -121,7 +140,25 @@ def visualize(self, image: Tensor, show_trajectories: bool = True) -> np.ndarray The image with the results of the tracker. """ frame_raw = (tensor_to_image(image) * 255).astype(np.uint8) - self.update(image) self.tracker.plot_results(frame_raw, show_trajectories=show_trajectories) - return frame_raw + return tensor(frame_raw).permute(2, 0, 1) + + def save( + self, image: Tensor, show_trajectories: bool = True, directory: Optional[str] = None + ) -> None: + """Save the model to ONNX format. + + Args: + image: The input image. + """ + if directory is None: + name = f"{self.name}_{datetime.datetime.now(tz=datetime.timezone.utc).strftime('%Y%m%d%H%M%S')!s}" + directory = os.path.join("kornia_outputs", name) + output = self.visualize(image, show_trajectories=show_trajectories) + + os.makedirs(directory, exist_ok=True) + write_image( + os.path.join(directory, f"{str(0).zfill(6)}.jpg"), output.byte(), + ) + logger.info(f"Outputs are saved in {directory}") diff --git a/kornia/models/utils.py b/kornia/models/utils.py index f0f4931531..b5977321e3 100644 --- a/kornia/models/utils.py +++ b/kornia/models/utils.py @@ -50,7 +50,7 @@ def __init__(self, interpolation_mode: str = "bilinear") -> None: super().__init__() self.interpolation_mode = interpolation_mode - def forward(self, imgs: Union[Tensor, list[Tensor]], original_sizes: Tensor) -> Tensor: + def forward(self, imgs: Union[Tensor, list[Tensor]], original_sizes: Tensor) -> Union[Tensor, list[Tensor]]: """ Returns: resized_imgs: resized images in a batch. diff --git a/requirements/requirements-dev.txt b/requirements/requirements-dev.txt index 52c1bdb8f1..38964f9735 100644 --- a/requirements/requirements-dev.txt +++ b/requirements/requirements-dev.txt @@ -3,7 +3,7 @@ boxmot coverage diffusers mypy -numpy>1.24.4,<2 +numpy<3 onnx onnxruntime pillow From 5b88c072711584deac0d998b304d46645a29ed5a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 15 Sep 2024 21:52:51 +0000 Subject: [PATCH 67/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/filters/dexined.py | 2 +- kornia/models/base.py | 13 +++++-------- kornia/models/detection/base.py | 1 - kornia/models/detection/rtdetr.py | 3 +-- kornia/models/edge_detection/base.py | 15 ++++++++++----- kornia/models/edge_detection/dexined.py | 3 +-- kornia/models/tracking/boxmot_tracker.py | 19 +++++++++---------- 7 files changed, 27 insertions(+), 29 deletions(-) diff --git a/kornia/filters/dexined.py b/kornia/filters/dexined.py index 1f77393253..c683f5f678 100644 --- a/kornia/filters/dexined.py +++ b/kornia/filters/dexined.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections import OrderedDict -from typing import ClassVar, Optional, List +from typing import ClassVar, List, Optional import torch import torch.nn.functional as F diff --git a/kornia/models/base.py b/kornia/models/base.py index 7ed6f12ef9..b9ca13d61f 100644 --- a/kornia/models/base.py +++ b/kornia/models/base.py @@ -1,18 +1,16 @@ -import os import datetime import logging +import os from typing import Optional, Union from kornia.core import Module, Tensor, stack from kornia.core.external import PILImage as Image from kornia.core.external import numpy as np -from kornia.utils.image import tensor_to_image from kornia.io import write_image +from kornia.utils.image import tensor_to_image logger = logging.getLogger(__name__) -from kornia.core import Module - class ModelBase(Module): """This class wraps a model and performs pre-processing and post-processing.""" @@ -35,7 +33,7 @@ def __init__( self.post_processor = post_processor.eval() if name is not None: self.name = name - + def _tensor_to_type( self, output: list[Tensor], output_type: str, is_batch: bool = False ) -> Union[Tensor, list[Tensor], list[Image.Image]]: # type: ignore @@ -48,8 +46,7 @@ def _tensor_to_type( raise RuntimeError(f"Unsupported output type `{output_type}`.") def _save_outputs( - self, outputs: Union[Tensor, list[Tensor]], directory: Optional[str] = None, - suffix: str = "" + self, outputs: Union[Tensor, list[Tensor]], directory: Optional[str] = None, suffix: str = "" ) -> None: """Save the output image(s) to a directory. @@ -67,4 +64,4 @@ def _save_outputs( os.path.join(directory, f"{str(i).zfill(6)}{suffix}.jpg"), out_image.mul(255.0).byte(), ) - logger.info(f"Outputs are saved in {directory}") \ No newline at end of file + logger.info(f"Outputs are saved in {directory}") diff --git a/kornia/models/detection/base.py b/kornia/models/detection/base.py index 2f8911a854..ca54699f39 100644 --- a/kornia/models/detection/base.py +++ b/kornia/models/detection/base.py @@ -21,7 +21,6 @@ ] - class BoundingBoxDataFormat(Enum): """Enum class that maps bounding box data format.""" diff --git a/kornia/models/detection/rtdetr.py b/kornia/models/detection/rtdetr.py index 9f54ddca78..27c6c51be0 100644 --- a/kornia/models/detection/rtdetr.py +++ b/kornia/models/detection/rtdetr.py @@ -18,13 +18,12 @@ class RTDETRDetectorBuilder: This class provides static methods to: - Build an object detection model from a model name or configuration. - Export the model to ONNX format for inference. - + .. code-block:: python image = kornia.utils.sample.get_sample_images()[0][None] model = RTDETRDetectorBuilder.build() model.save(image) - """ @staticmethod diff --git a/kornia/models/edge_detection/base.py b/kornia/models/edge_detection/base.py index 9fff5ec35c..17b8cf3a7b 100644 --- a/kornia/models/edge_detection/base.py +++ b/kornia/models/edge_detection/base.py @@ -1,4 +1,4 @@ -from typing import Union, Optional +from typing import Optional, Union from kornia.color.gray import grayscale_to_rgb from kornia.core import Tensor @@ -30,8 +30,10 @@ def forward(self, images: Union[Tensor, list[Tensor]]) -> Union[Tensor, list[Ten return self.post_processor(out_images, image_sizes) def visualize( - self, images: Union[Tensor, list[Tensor]], edge_maps: Optional[Union[Tensor, list[Tensor]]] = None, - output_type: str = "torch" + self, + images: Union[Tensor, list[Tensor]], + edge_maps: Optional[Union[Tensor, list[Tensor]]] = None, + output_type: str = "torch", ) -> Union[Tensor, list[Tensor], list[Image.Image]]: # type: ignore """Draw the segmentation results. @@ -51,8 +53,11 @@ def visualize( return self._tensor_to_type(output, output_type, is_batch=isinstance(images, Tensor)) def save( - self, images: Union[Tensor, list[Tensor]], edge_maps: Optional[Union[Tensor, list[Tensor]]] = None, - directory: Optional[str] = None, output_type: str = "torch" + self, + images: Union[Tensor, list[Tensor]], + edge_maps: Optional[Union[Tensor, list[Tensor]]] = None, + directory: Optional[str] = None, + output_type: str = "torch", ) -> None: """Save the segmentation results. diff --git a/kornia/models/edge_detection/dexined.py b/kornia/models/edge_detection/dexined.py index 8b2223e8c0..f56904847b 100644 --- a/kornia/models/edge_detection/dexined.py +++ b/kornia/models/edge_detection/dexined.py @@ -18,14 +18,13 @@ class DexiNedBuilder: image = kornia.utils.sample.get_sample_images()[0][None] model = DexiNedBuilder.build() model.save(image) - """ @staticmethod def build(model_name: str = "dexined", pretrained: bool = True, image_size: Optional[int] = 352) -> EdgeDetector: if model_name.lower() == "dexined": # Normalize then scale to [0, 255] - norm = Normalize(mean=tensor([[0.485, 0.456, 0.406]]), std=tensor([[1. / 255.] * 3])) + norm = Normalize(mean=tensor([[0.485, 0.456, 0.406]]), std=tensor([[1.0 / 255.0] * 3])) model = nn.Sequential(norm, DexiNed(pretrained=pretrained), nn.Sigmoid()) else: raise ValueError(f"Model {model_name} not found. Please choose from 'DexiNed'.") diff --git a/kornia/models/tracking/boxmot_tracker.py b/kornia/models/tracking/boxmot_tracker.py index af4533fa40..5970122e08 100644 --- a/kornia/models/tracking/boxmot_tracker.py +++ b/kornia/models/tracking/boxmot_tracker.py @@ -1,17 +1,17 @@ -import os import datetime import logging +import os from pathlib import Path -from typing import Union, Optional +from typing import Optional, Union +from kornia.config import kornia_config from kornia.core import Tensor, tensor from kornia.core.external import boxmot from kornia.core.external import numpy as np +from kornia.io import write_image from kornia.models.detection.base import ObjectDetector from kornia.models.detection.rtdetr import RTDETRDetectorBuilder from kornia.utils.image import tensor_to_image -from kornia.config import kornia_config -from kornia.io import write_image __all__ = ["BoxMotTracker"] @@ -51,7 +51,7 @@ class BoxMotTracker: frame_rate: Frame rate of the video being processed. Used to scale the track buffer size. fuse_first_associate: Whether to fuse appearance and motion information during the first association step. with_reid: Whether to use ReID (Re-Identification) features for association. - + .. code-block:: python import kornia @@ -60,7 +60,7 @@ class BoxMotTracker: for i in range(4): # At least 4 frames are needed to initialize the tracking position model.update(image) model.save(image) - + .. note:: At least 4 frames are needed to initialize the tracking position. """ @@ -144,9 +144,7 @@ def visualize(self, image: Tensor, show_trajectories: bool = True) -> Tensor: return tensor(frame_raw).permute(2, 0, 1) - def save( - self, image: Tensor, show_trajectories: bool = True, directory: Optional[str] = None - ) -> None: + def save(self, image: Tensor, show_trajectories: bool = True, directory: Optional[str] = None) -> None: """Save the model to ONNX format. Args: @@ -159,6 +157,7 @@ def save( os.makedirs(directory, exist_ok=True) write_image( - os.path.join(directory, f"{str(0).zfill(6)}.jpg"), output.byte(), + os.path.join(directory, f"{str(0).zfill(6)}.jpg"), + output.byte(), ) logger.info(f"Outputs are saved in {directory}") From e6295561618d43c00acbb89c86e1ced612549230 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Mon, 16 Sep 2024 01:01:31 +0300 Subject: [PATCH 68/84] update --- kornia/contrib/object_detection.py | 15 ++++++++++----- kornia/filters/dexined.py | 4 ++-- kornia/models/detection/base.py | 4 ++-- requirements/requirements-dev.txt | 1 + 4 files changed, 15 insertions(+), 9 deletions(-) diff --git a/kornia/contrib/object_detection.py b/kornia/contrib/object_detection.py index c8c4e336e5..b2a9ba20a4 100644 --- a/kornia/contrib/object_detection.py +++ b/kornia/contrib/object_detection.py @@ -29,14 +29,16 @@ class BoundingBox(BoundingBoxBase): warnings.warn( - "BoundingBox is deprecated and will be removed in v0.8.0. Use kornia.models.detector.BoundingBox instead.", + "BoundingBox is deprecated and will be removed in v0.8.0. " + "Use kornia.models.detector.BoundingBox instead.", DeprecationWarning, ) def results_from_detections(*args, **kwargs): warnings.warn( - "results_from_detections is deprecated and will be removed in v0.8.0. Use kornia.models.detector.results_from_detections instead.", + "results_from_detections is deprecated and will be removed in v0.8.0. " + "Use kornia.models.detector.results_from_detections instead.", DeprecationWarning, ) return results_from_detections_base(*args, **kwargs) @@ -44,20 +46,23 @@ def results_from_detections(*args, **kwargs): class ResizePreProcessor(ResizePreProcessorBase): warnings.warn( - "ResizePreProcessor is deprecated and will be removed in v0.8.0. Use kornia.models.utils.ResizePreProcessor instead.", + "ResizePreProcessor is deprecated and will be removed in v0.8.0. " + "Use kornia.models.utils.ResizePreProcessor instead.", DeprecationWarning, ) class ObjectDetector(ObjectDetectorBase): warnings.warn( - "ObjectDetector is deprecated and will be removed in v0.8.0. Use kornia.models.detector.ObjectDetector instead.", + "ObjectDetector is deprecated and will be removed in v0.8.0. " + "Use kornia.models.detector.ObjectDetector instead.", DeprecationWarning, ) class ObjectDetectorResult(ObjectDetectorResultBase): warnings.warn( - "ObjectDetectorResult is deprecated and will be removed in v0.8.0. Use kornia.models.detector.ObjectDetectorResult instead.", + "ObjectDetectorResult is deprecated and will be removed in v0.8.0. " + "Use kornia.models.detector.ObjectDetectorResult instead.", DeprecationWarning, ) diff --git a/kornia/filters/dexined.py b/kornia/filters/dexined.py index c683f5f678..ba4dc2f719 100644 --- a/kornia/filters/dexined.py +++ b/kornia/filters/dexined.py @@ -179,8 +179,8 @@ class DexiNed(Module): torch.Size([1, 1, 320, 320]) """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, 1, -1, -1] + ONNX_DEFAULT_INPUTSHAPE: ClassVar[list[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[list[int]] = [-1, 1, -1, -1] def __init__(self, pretrained: bool) -> None: super().__init__() diff --git a/kornia/models/detection/base.py b/kornia/models/detection/base.py index ca54699f39..73d106c73d 100644 --- a/kornia/models/detection/base.py +++ b/kornia/models/detection/base.py @@ -2,7 +2,7 @@ from dataclasses import dataclass from enum import Enum -from typing import Optional, Tuple, Union +from typing import Optional, Union import torch @@ -157,7 +157,7 @@ def to_onnx( onnx_name: Optional[str] = None, image_size: Optional[int] = 640, include_pre_and_post_processor: bool = True, - ) -> Tuple[str, ObjectDetector]: + ) -> tuple[str, ObjectDetector]: """Exports an RT-DETR object detection model to ONNX format. Either `model_name` or `config` must be provided. If neither is provided, diff --git a/requirements/requirements-dev.txt b/requirements/requirements-dev.txt index 38964f9735..0886bbedbb 100644 --- a/requirements/requirements-dev.txt +++ b/requirements/requirements-dev.txt @@ -1,6 +1,7 @@ accelerate boxmot coverage +distutils diffusers mypy numpy<3 From d092f65924f7d51d7600c9fd24f1cde9ff441700 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 15 Sep 2024 22:02:15 +0000 Subject: [PATCH 69/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/contrib/object_detection.py | 3 +-- kornia/filters/dexined.py | 2 +- requirements/requirements-dev.txt | 2 +- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/kornia/contrib/object_detection.py b/kornia/contrib/object_detection.py index b2a9ba20a4..40dfbd0962 100644 --- a/kornia/contrib/object_detection.py +++ b/kornia/contrib/object_detection.py @@ -29,8 +29,7 @@ class BoundingBox(BoundingBoxBase): warnings.warn( - "BoundingBox is deprecated and will be removed in v0.8.0. " - "Use kornia.models.detector.BoundingBox instead.", + "BoundingBox is deprecated and will be removed in v0.8.0. " "Use kornia.models.detector.BoundingBox instead.", DeprecationWarning, ) diff --git a/kornia/filters/dexined.py b/kornia/filters/dexined.py index ba4dc2f719..adb3d2437c 100644 --- a/kornia/filters/dexined.py +++ b/kornia/filters/dexined.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections import OrderedDict -from typing import ClassVar, List, Optional +from typing import ClassVar, Optional import torch import torch.nn.functional as F diff --git a/requirements/requirements-dev.txt b/requirements/requirements-dev.txt index 0886bbedbb..28b35e28dc 100644 --- a/requirements/requirements-dev.txt +++ b/requirements/requirements-dev.txt @@ -1,8 +1,8 @@ accelerate boxmot coverage -distutils diffusers +distutils mypy numpy<3 onnx From 385e9cd876448a4781ec5902bf9fdb7d58ab78c2 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Mon, 16 Sep 2024 18:42:40 +0300 Subject: [PATCH 70/84] update --- requirements/requirements-dev.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/requirements-dev.txt b/requirements/requirements-dev.txt index 28b35e28dc..ff0bb25e5a 100644 --- a/requirements/requirements-dev.txt +++ b/requirements/requirements-dev.txt @@ -4,7 +4,7 @@ coverage diffusers distutils mypy -numpy<3 +numpy<3,>=2.19.0 onnx onnxruntime pillow From abd2ffa435ff7f9d3fd927ce90232e11c572bba2 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Mon, 16 Sep 2024 18:43:09 +0300 Subject: [PATCH 71/84] update --- requirements/requirements-dev.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/requirements/requirements-dev.txt b/requirements/requirements-dev.txt index ff0bb25e5a..cc1aaceb33 100644 --- a/requirements/requirements-dev.txt +++ b/requirements/requirements-dev.txt @@ -2,7 +2,6 @@ accelerate boxmot coverage diffusers -distutils mypy numpy<3,>=2.19.0 onnx From de1e19efdcc771de8665c2ab4673102e9350b9fb Mon Sep 17 00:00:00 2001 From: shijianjian Date: Tue, 17 Sep 2024 00:44:03 +0300 Subject: [PATCH 72/84] update --- requirements/requirements-dev.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/requirements-dev.txt b/requirements/requirements-dev.txt index cc1aaceb33..9dd54df9ed 100644 --- a/requirements/requirements-dev.txt +++ b/requirements/requirements-dev.txt @@ -3,7 +3,7 @@ boxmot coverage diffusers mypy -numpy<3,>=2.19.0 +numpy==1.26.4 # For Boxmot onnx onnxruntime pillow From 81cae978e0fb4ca57f73a3d5be49e084b3878fe3 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Tue, 17 Sep 2024 00:49:28 +0300 Subject: [PATCH 73/84] update --- requirements/requirements-dev.txt | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/requirements/requirements-dev.txt b/requirements/requirements-dev.txt index 9dd54df9ed..21866c5bfb 100644 --- a/requirements/requirements-dev.txt +++ b/requirements/requirements-dev.txt @@ -1,9 +1,8 @@ accelerate -boxmot coverage diffusers mypy -numpy==1.26.4 # For Boxmot +numpy<3 onnx onnxruntime pillow From 6824af80683e7cf254edf79d5daed40172be0934 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Tue, 17 Sep 2024 01:19:10 +0300 Subject: [PATCH 74/84] update --- kornia/config.py | 4 ++-- kornia/models/base.py | 8 ++++---- kornia/models/detection/base.py | 7 +++---- kornia/models/segmentation/segmentation_models.py | 6 +++--- kornia/models/tracking/boxmot_tracker.py | 6 +++--- kornia/onnx/sequential.py | 4 ++-- 6 files changed, 17 insertions(+), 18 deletions(-) diff --git a/kornia/config.py b/kornia/config.py index 45688bb779..c21bc25b25 100644 --- a/kornia/config.py +++ b/kornia/config.py @@ -12,7 +12,7 @@ class InstallationMode(str, Enum): # Raise an error if the dependencies are not installed RAISE = "RAISE" - def __eq__(self, other): + def __eq__(self, other: object) -> bool: if isinstance(other, str): return self.value.lower() == other.lower() # Case-insensitive comparison return super().__eq__(other) @@ -26,7 +26,7 @@ def installation_mode(self) -> InstallationMode: return self._installation_mode @installation_mode.setter - def installation_mode(self, value: str): + def installation_mode(self, value: str) -> None: # Allow setting via string by converting to the Enum if isinstance(value, str): try: diff --git a/kornia/models/base.py b/kornia/models/base.py index b9ca13d61f..56e4b2681d 100644 --- a/kornia/models/base.py +++ b/kornia/models/base.py @@ -1,7 +1,7 @@ import datetime import logging import os -from typing import Optional, Union +from typing import Optional, Union, List from kornia.core import Module, Tensor, stack from kornia.core.external import PILImage as Image @@ -35,8 +35,8 @@ def __init__( self.name = name def _tensor_to_type( - self, output: list[Tensor], output_type: str, is_batch: bool = False - ) -> Union[Tensor, list[Tensor], list[Image.Image]]: # type: ignore + self, output: List[Tensor], output_type: str, is_batch: bool = False + ) -> Union[Tensor, List[Tensor], List[Image.Image]]: # type: ignore if output_type == "torch": if is_batch: return stack(output) @@ -46,7 +46,7 @@ def _tensor_to_type( raise RuntimeError(f"Unsupported output type `{output_type}`.") def _save_outputs( - self, outputs: Union[Tensor, list[Tensor]], directory: Optional[str] = None, suffix: str = "" + self, outputs: Union[Tensor, List[Tensor]], directory: Optional[str] = None, suffix: str = "" ) -> None: """Save the output image(s) to a directory. diff --git a/kornia/models/detection/base.py b/kornia/models/detection/base.py index 73d106c73d..b47025ab7e 100644 --- a/kornia/models/detection/base.py +++ b/kornia/models/detection/base.py @@ -125,10 +125,9 @@ def visualize( Needs to be more fancy later. """ - if detections is None: - detections = self.forward(images) + dets = detections or self.forward(images) output = [] - for image, detection in zip(images, detections): + for image, detection in zip(images, dets): out_img = image[None].clone() for out in detection: out_img = draw_rectangle( @@ -157,7 +156,7 @@ def to_onnx( onnx_name: Optional[str] = None, image_size: Optional[int] = 640, include_pre_and_post_processor: bool = True, - ) -> tuple[str, ObjectDetector]: + ) -> str: """Exports an RT-DETR object detection model to ONNX format. Either `model_name` or `config` must be provided. If neither is provided, diff --git a/kornia/models/segmentation/segmentation_models.py b/kornia/models/segmentation/segmentation_models.py index 96e8a8ed7a..03b2f1da61 100644 --- a/kornia/models/segmentation/segmentation_models.py +++ b/kornia/models/segmentation/segmentation_models.py @@ -1,4 +1,4 @@ -from typing import Optional +from typing import Optional, ClassVar, List import kornia from kornia.core import Module, Tensor, ones_like, tensor, zeros_like @@ -28,8 +28,8 @@ class SegmentationModels(Module, ONNXExportMixin): Pretrained weights for the whole model are not available. """ - ONNX_DEFAULT_INPUTSHAPE = (-1, 3, -1, -1) - ONNX_DEFAULT_OUTPUTSHAPE = (-1, -1, -1, -1) + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = (-1, 3, -1, -1) + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = (-1, -1, -1, -1) def __init__( self, diff --git a/kornia/models/tracking/boxmot_tracker.py b/kornia/models/tracking/boxmot_tracker.py index 5970122e08..260838c7be 100644 --- a/kornia/models/tracking/boxmot_tracker.py +++ b/kornia/models/tracking/boxmot_tracker.py @@ -108,7 +108,7 @@ def update(self, image: Tensor) -> None: detections: Union[Tensor, list[Tensor]] = self.detector(image) - detections = detections[0].cpu().numpy() # Batch size is 1 + detections = detections.cpu().numpy()[0] # Batch size is 1 detections = np.array( # type: ignore [ @@ -125,9 +125,9 @@ def update(self, image: Tensor) -> None: # empty N X (x, y, x, y, conf, cls) detections = np.empty((0, 6)) # type: ignore - frame_raw = (tensor_to_image(image) * 255).astype(np.uint8) # type: ignore + frame_raw = (tensor_to_image(image) * 255).astype(np.uint8) # --> M X (x, y, x, y, id, conf, cls, ind) - return self.tracker.update(detections, frame_raw) # type: ignore + return self.tracker.update(detections, frame_raw) def visualize(self, image: Tensor, show_trajectories: bool = True) -> Tensor: """Visualize the results of the tracker. diff --git a/kornia/onnx/sequential.py b/kornia/onnx/sequential.py index 0cbadc8eaa..8c90e5eed1 100644 --- a/kornia/onnx/sequential.py +++ b/kornia/onnx/sequential.py @@ -136,11 +136,11 @@ def get_session(self) -> "ort.InferenceSession": # type: ignore def as_cpu(self) -> None: """Set the session to run on CPU.""" - self._session = self.create_session(providers=["CPUExecutionProvider"]) + self._session.set_providers(["CPUExecutionProvider"]) def as_cuda(self, device: int = 0) -> None: """Set the session to run on CUDA.""" - self._session = self.create_session(providers=[("CUDAExecutionProvider", {"device_id": device})]) + self._session.set_providers(["CUDAExecutionProvider"], provider_options=[{"device_id": device}]) def __call__(self, *inputs: "np.ndarray") -> List["np.ndarray"]: # type:ignore """Perform inference using the combined ONNX model. From 3daa151ebf292bd4ecd340105d59528d8282a204 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 16 Sep 2024 22:19:33 +0000 Subject: [PATCH 75/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/models/base.py | 2 +- kornia/models/segmentation/segmentation_models.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/kornia/models/base.py b/kornia/models/base.py index 56e4b2681d..20c1b4d9c0 100644 --- a/kornia/models/base.py +++ b/kornia/models/base.py @@ -1,7 +1,7 @@ import datetime import logging import os -from typing import Optional, Union, List +from typing import List, Optional, Union from kornia.core import Module, Tensor, stack from kornia.core.external import PILImage as Image diff --git a/kornia/models/segmentation/segmentation_models.py b/kornia/models/segmentation/segmentation_models.py index 03b2f1da61..249028de90 100644 --- a/kornia/models/segmentation/segmentation_models.py +++ b/kornia/models/segmentation/segmentation_models.py @@ -1,4 +1,4 @@ -from typing import Optional, ClassVar, List +from typing import ClassVar, List, Optional import kornia from kornia.core import Module, Tensor, ones_like, tensor, zeros_like From 30cffadddea8919ab75cd4cda58eaf1dfd533bb2 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Tue, 17 Sep 2024 14:49:38 +0300 Subject: [PATCH 76/84] update --- kornia/contrib/object_detection.py | 3 ++- kornia/models/segmentation/segmentation_models.py | 8 ++++---- kornia/models/tracking/boxmot_tracker.py | 8 ++++---- kornia/models/utils.py | 8 ++++---- 4 files changed, 14 insertions(+), 13 deletions(-) diff --git a/kornia/contrib/object_detection.py b/kornia/contrib/object_detection.py index 40dfbd0962..f89fba96d4 100644 --- a/kornia/contrib/object_detection.py +++ b/kornia/contrib/object_detection.py @@ -1,4 +1,5 @@ import warnings +from typing import Any from kornia.models.detection.base import ( BoundingBox as BoundingBoxBase, @@ -34,7 +35,7 @@ class BoundingBox(BoundingBoxBase): ) -def results_from_detections(*args, **kwargs): +def results_from_detections(*args: Any, **kwargs: Any): warnings.warn( "results_from_detections is deprecated and will be removed in v0.8.0. " "Use kornia.models.detector.results_from_detections instead.", diff --git a/kornia/models/segmentation/segmentation_models.py b/kornia/models/segmentation/segmentation_models.py index 249028de90..dc0469feff 100644 --- a/kornia/models/segmentation/segmentation_models.py +++ b/kornia/models/segmentation/segmentation_models.py @@ -1,4 +1,4 @@ -from typing import ClassVar, List, Optional +from typing import Any, ClassVar, List, Optional import kornia from kornia.core import Module, Tensor, ones_like, tensor, zeros_like @@ -28,8 +28,8 @@ class SegmentationModels(Module, ONNXExportMixin): Pretrained weights for the whole model are not available. """ - ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = (-1, 3, -1, -1) - ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = (-1, -1, -1, -1) + ONNX_DEFAULT_INPUTSHAPE: ClassVar[List[int]] = [-1, 3, -1, -1] + ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, -1, -1, -1] def __init__( self, @@ -38,7 +38,7 @@ def __init__( encoder_weights: Optional[str] = "imagenet", in_channels: int = 3, classes: int = 1, - **kwargs, + **kwargs: Any, ) -> None: super().__init__() self.preproc_params = smp.encoders.get_preprocessing_params(encoder_name) # type: ignore diff --git a/kornia/models/tracking/boxmot_tracker.py b/kornia/models/tracking/boxmot_tracker.py index 260838c7be..92549e8471 100644 --- a/kornia/models/tracking/boxmot_tracker.py +++ b/kornia/models/tracking/boxmot_tracker.py @@ -2,7 +2,7 @@ import logging import os from pathlib import Path -from typing import Optional, Union +from typing import Any,Optional, Union from kornia.config import kornia_config from kornia.core import Tensor, tensor @@ -74,7 +74,7 @@ def __init__( tracker_model_weights: str = "osnet_x0_25_msmt17.pt", device: str = "cpu", fp16: bool = False, - **kwargs, + **kwargs: Any, ) -> None: super().__init__() if isinstance(detector, str): @@ -106,9 +106,9 @@ def update(self, image: Tensor) -> None: if image.ndim == 3: image = image.unsqueeze(0) - detections: Union[Tensor, list[Tensor]] = self.detector(image) + detections_raw: Union[Tensor, list[Tensor]] = self.detector(image) - detections = detections.cpu().numpy()[0] # Batch size is 1 + detections = detections_raw[0].cpu().numpy() # Batch size is 1 detections = np.array( # type: ignore [ diff --git a/kornia/models/utils.py b/kornia/models/utils.py index b5977321e3..a9d0b654f4 100644 --- a/kornia/models/utils.py +++ b/kornia/models/utils.py @@ -1,5 +1,5 @@ import warnings -from typing import Union +from typing import Union, List, Tuple import torch from torch import Tensor @@ -26,7 +26,7 @@ def __init__(self, height: int, width: int, interpolation_mode: str = "bilinear" self.size = (height, width) self.interpolation_mode = interpolation_mode - def forward(self, imgs: Union[Tensor, list[Tensor]]) -> tuple[Tensor, Tensor]: + def forward(self, imgs: Union[Tensor, List[Tensor]]) -> Tuple[Tensor, Tensor]: """ Returns: resized_imgs: resized images in a batch. @@ -50,14 +50,14 @@ def __init__(self, interpolation_mode: str = "bilinear") -> None: super().__init__() self.interpolation_mode = interpolation_mode - def forward(self, imgs: Union[Tensor, list[Tensor]], original_sizes: Tensor) -> Union[Tensor, list[Tensor]]: + def forward(self, imgs: Union[Tensor, List[Tensor]], original_sizes: Tensor) -> Union[Tensor, List[Tensor]]: """ Returns: resized_imgs: resized images in a batch. original_sizes: the original image sizes of (height, width). """ # TODO: support other input formats e.g. file path, numpy - resized_imgs: list[Tensor] = [] + resized_imgs: Union[Tensor, list[Tensor]] = [] if not torch.onnx.is_in_onnx_export(): iters = len(imgs) if isinstance(imgs, list) else imgs.shape[0] From 842aeef7b7263dc52e32affe8cd99d5470112c0f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 17 Sep 2024 11:50:20 +0000 Subject: [PATCH 77/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/models/tracking/boxmot_tracker.py | 2 +- kornia/models/utils.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/kornia/models/tracking/boxmot_tracker.py b/kornia/models/tracking/boxmot_tracker.py index 92549e8471..d10c5aea01 100644 --- a/kornia/models/tracking/boxmot_tracker.py +++ b/kornia/models/tracking/boxmot_tracker.py @@ -2,7 +2,7 @@ import logging import os from pathlib import Path -from typing import Any,Optional, Union +from typing import Any, Optional, Union from kornia.config import kornia_config from kornia.core import Tensor, tensor diff --git a/kornia/models/utils.py b/kornia/models/utils.py index a9d0b654f4..9e535c156a 100644 --- a/kornia/models/utils.py +++ b/kornia/models/utils.py @@ -1,5 +1,5 @@ import warnings -from typing import Union, List, Tuple +from typing import List, Tuple, Union import torch from torch import Tensor From 837388928509e6358b0a15f59c5162fc65df26f8 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Tue, 17 Sep 2024 15:15:46 +0300 Subject: [PATCH 78/84] update --- kornia/models/tracking/boxmot_tracker.py | 14 +++++++------- kornia/models/utils.py | 21 ++++++++++----------- tests/contrib/test_object_detector.py | 4 ++-- tests/enhance/test_normalize.py | 4 ++-- tests/filters/test_filters.py | 5 +++-- 5 files changed, 24 insertions(+), 24 deletions(-) diff --git a/kornia/models/tracking/boxmot_tracker.py b/kornia/models/tracking/boxmot_tracker.py index d10c5aea01..254063fd08 100644 --- a/kornia/models/tracking/boxmot_tracker.py +++ b/kornia/models/tracking/boxmot_tracker.py @@ -112,16 +112,16 @@ def update(self, image: Tensor) -> None: detections = np.array( # type: ignore [ - detections[:, 2], # type: ignore - detections[:, 3], # type: ignore - detections[:, 2] + detections[:, 4], # type: ignore - detections[:, 3] + detections[:, 5], # type: ignore - detections[:, 1], # type: ignore - detections[:, 0], # type: ignore + detections[:, 2], + detections[:, 3], + detections[:, 2] + detections[:, 4], + detections[:, 3] + detections[:, 5], + detections[:, 1], + detections[:, 0], ] ).T - if detections.shape[0] == 0: # type: ignore + if detections.shape[0] == 0: # empty N X (x, y, x, y, conf, cls) detections = np.empty((0, 6)) # type: ignore diff --git a/kornia/models/utils.py b/kornia/models/utils.py index 9e535c156a..f12a5e3b09 100644 --- a/kornia/models/utils.py +++ b/kornia/models/utils.py @@ -57,21 +57,20 @@ def forward(self, imgs: Union[Tensor, List[Tensor]], original_sizes: Tensor) -> original_sizes: the original image sizes of (height, width). """ # TODO: support other input formats e.g. file path, numpy - resized_imgs: Union[Tensor, list[Tensor]] = [] + resized_imgs: list[Tensor] = [] - if not torch.onnx.is_in_onnx_export(): - iters = len(imgs) if isinstance(imgs, list) else imgs.shape[0] - for i in range(iters): - img = imgs[i] - size = original_sizes[i] - resized_imgs.append( - resize(img[None], size=size.cpu().long().numpy().tolist(), interpolation=self.interpolation_mode) - ) - else: + if torch.onnx.is_in_onnx_export(): warnings.warn( "ResizePostProcessor is not supported in ONNX export. " "The output will not be resized back to the original size." ) - resized_imgs = imgs + return imgs + iters = len(imgs) if isinstance(imgs, list) else imgs.shape[0] + for i in range(iters): + img = imgs[i] + size = original_sizes[i] + resized_imgs.append( + resize(img[None], size=size.cpu().long().numpy().tolist(), interpolation=self.interpolation_mode) + ) return resized_imgs diff --git a/tests/contrib/test_object_detector.py b/tests/contrib/test_object_detector.py index 3a3b71b500..954e59c2dd 100644 --- a/tests/contrib/test_object_detector.py +++ b/tests/contrib/test_object_detector.py @@ -16,7 +16,7 @@ def test_smoke(self, device, dtype): confidence = 0.3 config = RTDETRConfig("resnet50d", 10, head_num_queries=10) model = RTDETR.from_config(config).to(device, dtype).eval() - pre_processor = kornia.contrib.object_detection.ResizePreProcessor((32, 32)) + pre_processor = kornia.contrib.object_detection.ResizePreProcessor(32, 32) post_processor = DETRPostProcessor(confidence, num_top_queries=3).to(device, dtype).eval() detector = kornia.contrib.ObjectDetector(model, pre_processor, post_processor) @@ -39,7 +39,7 @@ def test_smoke(self, device, dtype): def test_onnx(self, device, dtype, tmp_path: Path, variant: str): config = RTDETRConfig(variant, 1) model = RTDETR.from_config(config).to(device=device, dtype=dtype).eval() - pre_processor = kornia.contrib.object_detection.ResizePreProcessor((640, 640)) + pre_processor = kornia.contrib.object_detection.ResizePreProcessor(640, 640) post_processor = DETRPostProcessor(0.3, num_top_queries=3) detector = kornia.contrib.ObjectDetector(model, pre_processor, post_processor) diff --git a/tests/enhance/test_normalize.py b/tests/enhance/test_normalize.py index 51dc469f8f..d6f6deeaca 100644 --- a/tests/enhance/test_normalize.py +++ b/tests/enhance/test_normalize.py @@ -10,7 +10,7 @@ class TestNormalize(BaseTester): def test_smoke(self, device, dtype): mean = [0.5] std = [0.1] - repr = "Normalize(mean=tensor([0.5000]), std=tensor([0.1000]))" + repr = "Normalize(mean=tensor([[0.5000]]), std=tensor([[0.1000]]))" assert str(kornia.enhance.Normalize(mean, std)) == repr def test_normalize(self, device, dtype): @@ -137,7 +137,7 @@ def test_random_normalize_different_parameter_types(self, mean, std): def test_random_normalize_invalid_parameter_shape(self, mean, std): f = kornia.enhance.Normalize(mean=mean, std=std) inputs = torch.arange(0.0, 16.0, step=1).reshape(1, 4, 4).unsqueeze(0) - with pytest.raises(ValueError): + with pytest.raises((ValueError, RuntimeError)): f(inputs) @pytest.mark.skip(reason="not implemented yet") diff --git a/tests/filters/test_filters.py b/tests/filters/test_filters.py index 77e669d50d..7482bffa34 100644 --- a/tests/filters/test_filters.py +++ b/tests/filters/test_filters.py @@ -729,9 +729,10 @@ class TestDexiNed(BaseTester): def test_smoke(self, device, dtype): img = torch.rand(2, 3, 32, 32, device=device, dtype=dtype) net = DexiNed(pretrained=False).to(device, dtype) + feat = net.get_features(img) + assert len(feat) == 6 out = net(img) - assert len(out) == 7 - assert out[-1].shape == (2, 1, 32, 32) + assert out.shape == (2, 1, 32, 32) @pytest.mark.slow @pytest.mark.parametrize("data", ["dexined"], indirect=True) From e927844ee341515bf4880e5594e41c6767c78d85 Mon Sep 17 00:00:00 2001 From: shijianjian Date: Tue, 17 Sep 2024 15:21:19 +0300 Subject: [PATCH 79/84] bug fix --- kornia/contrib/object_detection.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/kornia/contrib/object_detection.py b/kornia/contrib/object_detection.py index f89fba96d4..d85619dd6b 100644 --- a/kornia/contrib/object_detection.py +++ b/kornia/contrib/object_detection.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import warnings from typing import Any @@ -35,7 +37,7 @@ class BoundingBox(BoundingBoxBase): ) -def results_from_detections(*args: Any, **kwargs: Any): +def results_from_detections(*args: Any, **kwargs: Any) -> list[ObjectDetectorResultBase]: warnings.warn( "results_from_detections is deprecated and will be removed in v0.8.0. " "Use kornia.models.detector.results_from_detections instead.", From 20753678aa0e0b670c696bdc9c3fd2c1744b2429 Mon Sep 17 00:00:00 2001 From: Jian S Date: Fri, 20 Sep 2024 05:16:27 +0800 Subject: [PATCH 80/84] tests fix --- kornia/contrib/object_detection.py | 48 ++++++++++++++++----------- kornia/models/__init__.py | 1 + kornia/models/utils.py | 2 ++ kornia/onnx/utils.py | 4 +-- tests/contrib/test_object_detector.py | 8 ++--- tests/onnx/test_utils.py | 26 +++++++-------- 6 files changed, 50 insertions(+), 39 deletions(-) diff --git a/kornia/contrib/object_detection.py b/kornia/contrib/object_detection.py index d85619dd6b..96a648b427 100644 --- a/kornia/contrib/object_detection.py +++ b/kornia/contrib/object_detection.py @@ -31,40 +31,48 @@ class BoundingBox(BoundingBoxBase): - warnings.warn( - "BoundingBox is deprecated and will be removed in v0.8.0. " "Use kornia.models.detector.BoundingBox instead.", - DeprecationWarning, - ) + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + warnings.warn( + "BoundingBox is deprecated and will be removed in v0.8.0. " "Use kornia.models.detection.BoundingBox instead.", + DeprecationWarning, + ) def results_from_detections(*args: Any, **kwargs: Any) -> list[ObjectDetectorResultBase]: warnings.warn( "results_from_detections is deprecated and will be removed in v0.8.0. " - "Use kornia.models.detector.results_from_detections instead.", + "Use kornia.models.detection.results_from_detections instead.", DeprecationWarning, ) return results_from_detections_base(*args, **kwargs) class ResizePreProcessor(ResizePreProcessorBase): - warnings.warn( - "ResizePreProcessor is deprecated and will be removed in v0.8.0. " - "Use kornia.models.utils.ResizePreProcessor instead.", - DeprecationWarning, - ) + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + warnings.warn( + "ResizePreProcessor is deprecated and will be removed in v0.8.0. " + "Use kornia.models.utils.ResizePreProcessor instead.", + DeprecationWarning, + ) class ObjectDetector(ObjectDetectorBase): - warnings.warn( - "ObjectDetector is deprecated and will be removed in v0.8.0. " - "Use kornia.models.detector.ObjectDetector instead.", - DeprecationWarning, - ) + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + warnings.warn( + "ObjectDetector is deprecated and will be removed in v0.8.0. " + "Use kornia.models.detection.ObjectDetector instead.", + DeprecationWarning, + ) class ObjectDetectorResult(ObjectDetectorResultBase): - warnings.warn( - "ObjectDetectorResult is deprecated and will be removed in v0.8.0. " - "Use kornia.models.detector.ObjectDetectorResult instead.", - DeprecationWarning, - ) + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + warnings.warn( + "ObjectDetectorResult is deprecated and will be removed in v0.8.0. " + "Use kornia.models.detection.ObjectDetectorResult instead.", + DeprecationWarning, + ) diff --git a/kornia/models/__init__.py b/kornia/models/__init__.py index 9117c87c71..3f2f71da48 100644 --- a/kornia/models/__init__.py +++ b/kornia/models/__init__.py @@ -1 +1,2 @@ from . import detection, segmentation, tracking +from .utils import * diff --git a/kornia/models/utils.py b/kornia/models/utils.py index f12a5e3b09..34e95b87e7 100644 --- a/kornia/models/utils.py +++ b/kornia/models/utils.py @@ -7,6 +7,8 @@ from kornia.core import Module, concatenate from kornia.geometry.transform import resize +__all__ = ["ResizePreProcessor", "ResizePostProcessor"] + class ResizePreProcessor(Module): """This module resizes a list of image tensors to the given size. diff --git a/kornia/onnx/utils.py b/kornia/onnx/utils.py index 51977bdf49..ac40beae81 100644 --- a/kornia/onnx/utils.py +++ b/kornia/onnx/utils.py @@ -43,8 +43,8 @@ def _get_file_path(self, model_name: str, cache_dir: Optional[str]) -> str: cache_dir = kornia_config.hub_onnx_dir # The filename is the model name (without directory path) - file_name = f"{model_name.split('/')[-1]}.onnx" - file_path = os.path.join(cache_dir, "/".join(model_name.split("/")[:-1]), file_name) + file_name = f"{os.path.split(model_name)[-1]}.onnx" + file_path = os.path.join(*os.path.split(cache_dir), *os.path.split(model_name)[:-1], file_name) return file_path def load_model(self, model_name: str, download: bool = True, **kwargs) -> "onnx.ModelProto": # type:ignore diff --git a/tests/contrib/test_object_detector.py b/tests/contrib/test_object_detector.py index 954e59c2dd..072a56635f 100644 --- a/tests/contrib/test_object_detector.py +++ b/tests/contrib/test_object_detector.py @@ -16,9 +16,9 @@ def test_smoke(self, device, dtype): confidence = 0.3 config = RTDETRConfig("resnet50d", 10, head_num_queries=10) model = RTDETR.from_config(config).to(device, dtype).eval() - pre_processor = kornia.contrib.object_detection.ResizePreProcessor(32, 32) + pre_processor = kornia.models.utils.ResizePreProcessor(32, 32) post_processor = DETRPostProcessor(confidence, num_top_queries=3).to(device, dtype).eval() - detector = kornia.contrib.ObjectDetector(model, pre_processor, post_processor) + detector = kornia.models.detection.ObjectDetector(model, pre_processor, post_processor) sizes = torch.randint(5, 10, (batch_size, 2)) * 32 imgs = [torch.randn(3, h, w, device=device, dtype=dtype) for h, w in sizes] @@ -39,9 +39,9 @@ def test_smoke(self, device, dtype): def test_onnx(self, device, dtype, tmp_path: Path, variant: str): config = RTDETRConfig(variant, 1) model = RTDETR.from_config(config).to(device=device, dtype=dtype).eval() - pre_processor = kornia.contrib.object_detection.ResizePreProcessor(640, 640) + pre_processor = kornia.models.utils.ResizePreProcessor(640, 640) post_processor = DETRPostProcessor(0.3, num_top_queries=3) - detector = kornia.contrib.ObjectDetector(model, pre_processor, post_processor) + detector = kornia.models.detection.ObjectDetector(model, pre_processor, post_processor) data = torch.rand(3, 400, 640, device=device, dtype=dtype) diff --git a/tests/onnx/test_utils.py b/tests/onnx/test_utils.py index 43e70f2799..bbc6f91adf 100644 --- a/tests/onnx/test_utils.py +++ b/tests/onnx/test_utils.py @@ -14,14 +14,14 @@ def loader(self): return ONNXLoader(cache_dir=".test_cache") def test_get_file_path_with_custom_cache_dir(self, loader): - model_name = "operators/some_model" - expected_path = ".test_cache/operators/some_model.onnx" + model_name = os.path.join("operators", "some_model") + expected_path = os.path.join(".test_cache", "operators", "some_model.onnx") assert loader._get_file_path(model_name, loader.cache_dir) == expected_path def test_get_file_path_with_default_cache_dir(self): loader = ONNXLoader() - model_name = "operators/some_model" - expected_path = ".kornia_hub/onnx_models/operators/some_model.onnx" + model_name = os.path.join("operators", "some_model") + expected_path = os.path.join(".kornia_hub", "onnx_models", "operators", "some_model.onnx") assert loader._get_file_path(model_name, None) == expected_path @mock.patch("onnx.load") @@ -53,7 +53,7 @@ def test_load_model_download(self, mock_exists, mock_urlretrieve, loader): assert model == mock_model mock_urlretrieve.assert_called_once_with( "https://huggingface.co/kornia/ONNX_models/resolve/main/operators/some_model.onnx", - ".test_cache/operators/some_model.onnx", + os.path.join(".test_cache", "operators", "some_model.onnx"), ) def test_load_model_file_not_found(self, loader): @@ -66,7 +66,7 @@ def test_load_model_file_not_found(self, loader): @mock.patch("os.makedirs") def test_download_success(self, mock_makedirs, mock_urlretrieve, loader): url = "https://huggingface.co/some_model.onnx" - file_path = ".test_cache/some_model.onnx" + file_path = os.path.join(".test_cache", "some_model.onnx") loader.download(url, file_path) @@ -79,7 +79,7 @@ def test_download_success(self, mock_makedirs, mock_urlretrieve, loader): ) def test_download_failure(self, mock_urlretrieve, loader): url = "https://huggingface.co/non_existent_model.onnx" - file_path = ".test_cache/non_existent_model.onnx" + file_path = os.path.join(".test_cache", "non_existent_model.onnx") with pytest.raises(ValueError, match="Error in resolving"): loader.download(url, file_path) @@ -88,11 +88,11 @@ def test_download_failure(self, mock_urlretrieve, loader): def test_fetch_repo_contents_success(self, mock_get): mock_response = mock.Mock() mock_response.status_code = 200 - mock_response.json.return_value = [{"path": "operators/model.onnx"}] + mock_response.json.return_value = [{"path": os.path.join("operators", "model.onnx")}] mock_get.return_value = mock_response contents = ONNXLoader._fetch_repo_contents("operators") - assert contents == [{"path": "operators/model.onnx"}] + assert contents == [{"path": os.path.join("operators", "model.onnx")}] @mock.patch("requests.get") def test_fetch_repo_contents_failure(self, mock_get): @@ -105,18 +105,18 @@ def test_fetch_repo_contents_failure(self, mock_get): @mock.patch("kornia.onnx.utils.ONNXLoader._fetch_repo_contents") def test_list_operators(self, mock_fetch_repo_contents, capsys): - mock_fetch_repo_contents.return_value = [{"path": "operators/some_model.onnx"}] + mock_fetch_repo_contents.return_value = [{"path": os.path.join("operators", "some_model.onnx")}] ONNXLoader.list_operators() captured = capsys.readouterr() - assert "operators/some_model.onnx" in captured.out + assert os.path.join("operators", "some_model.onnx").replace('\\', "\\\\") in captured.out # .replace() for Windows @mock.patch("kornia.onnx.utils.ONNXLoader._fetch_repo_contents") def test_list_models(self, mock_fetch_repo_contents, capsys): - mock_fetch_repo_contents.return_value = [{"path": "models/some_model.onnx"}] + mock_fetch_repo_contents.return_value = [{"path": os.path.join("operators", "some_model.onnx")}] ONNXLoader.list_models() captured = capsys.readouterr() - assert "models/some_model.onnx" in captured.out + assert os.path.join("operators", "some_model.onnx").replace('\\', "\\\\") in captured.out # .replace() for Windows From 44f54a88b31f643709c0f68d0b4c7c6d499e45bc Mon Sep 17 00:00:00 2001 From: shijianjian Date: Fri, 20 Sep 2024 00:31:36 +0300 Subject: [PATCH 81/84] updated --- docs/source/models.rst | 135 +++++++++++++++++++++++++++++++++++++++++ kornia/onnx/utils.py | 6 ++ 2 files changed, 141 insertions(+) create mode 100644 docs/source/models.rst diff --git a/docs/source/models.rst b/docs/source/models.rst new file mode 100644 index 0000000000..223453037f --- /dev/null +++ b/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. diff --git a/kornia/onnx/utils.py b/kornia/onnx/utils.py index 51977bdf49..9980f26ac2 100644 --- a/kornia/onnx/utils.py +++ b/kornia/onnx/utils.py @@ -54,6 +54,7 @@ def load_model(self, model_name: str, download: bool = True, **kwargs) -> "onnx. model_name: The name of the ONNX model or operator. For Hugging Face-hosted models, use the format 'hf://model_name'. Valid `model_name` can be found on https://huggingface.co/kornia/ONNX_models. + Or a URL to the ONNX model. download: If True, the model will be downloaded from Hugging Face if it's not already in the local cache. **kwargs: Additional arguments to pass to the download method, if needed. @@ -72,6 +73,11 @@ def load_model(self, model_name: str, download: bool = True, **kwargs) -> "onnx. else: raise ValueError(f"`{model_name}` is not found in `{file_path}`. You may set `download=True`.") return onnx.load(file_path) # type:ignore + elif model_name.startswith("https://"): + cache_dir = kwargs.get(kornia_config.hub_onnx_dir, None) or self.cache_dir + file_path = self._get_file_path(model_name, cache_dir) + self.download(model_name, file_path) + return onnx.load(file_path) # type:ignore if os.path.exists(model_name): return onnx.load(model_name) # type:ignore From 82f4750a3d96d93af5f70af69f8b52c67dae937e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 19 Sep 2024 21:32:21 +0000 Subject: [PATCH 82/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/contrib/object_detection.py | 3 ++- tests/onnx/test_utils.py | 8 ++++++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/kornia/contrib/object_detection.py b/kornia/contrib/object_detection.py index 96a648b427..e43d97394d 100644 --- a/kornia/contrib/object_detection.py +++ b/kornia/contrib/object_detection.py @@ -34,7 +34,8 @@ class BoundingBox(BoundingBoxBase): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) warnings.warn( - "BoundingBox is deprecated and will be removed in v0.8.0. " "Use kornia.models.detection.BoundingBox instead.", + "BoundingBox is deprecated and will be removed in v0.8.0. " + "Use kornia.models.detection.BoundingBox instead.", DeprecationWarning, ) diff --git a/tests/onnx/test_utils.py b/tests/onnx/test_utils.py index bbc6f91adf..82be0c1629 100644 --- a/tests/onnx/test_utils.py +++ b/tests/onnx/test_utils.py @@ -110,7 +110,9 @@ def test_list_operators(self, mock_fetch_repo_contents, capsys): ONNXLoader.list_operators() captured = capsys.readouterr() - assert os.path.join("operators", "some_model.onnx").replace('\\', "\\\\") in captured.out # .replace() for Windows + assert ( + os.path.join("operators", "some_model.onnx").replace("\\", "\\\\") in captured.out + ) # .replace() for Windows @mock.patch("kornia.onnx.utils.ONNXLoader._fetch_repo_contents") def test_list_models(self, mock_fetch_repo_contents, capsys): @@ -119,4 +121,6 @@ def test_list_models(self, mock_fetch_repo_contents, capsys): ONNXLoader.list_models() captured = capsys.readouterr() - assert os.path.join("operators", "some_model.onnx").replace('\\', "\\\\") in captured.out # .replace() for Windows + assert ( + os.path.join("operators", "some_model.onnx").replace("\\", "\\\\") in captured.out + ) # .replace() for Windows From 87a697fd66f4891087be731cf1d44b0698ae09c5 Mon Sep 17 00:00:00 2001 From: Jian S Date: Fri, 20 Sep 2024 06:44:37 +0800 Subject: [PATCH 83/84] update --- kornia/models/detection/utils.py | 72 +++++++++++---- .../segmentation/segmentation_models.py | 9 +- kornia/onnx/sequential.py | 41 ++++++++- tests/models/box_filtering.py | 90 +++++++++++++++++++ 4 files changed, 188 insertions(+), 24 deletions(-) create mode 100644 tests/models/box_filtering.py diff --git a/kornia/models/detection/utils.py b/kornia/models/detection/utils.py index 51a11eed5d..6ec0d823e2 100644 --- a/kornia/models/detection/utils.py +++ b/kornia/models/detection/utils.py @@ -1,6 +1,6 @@ from typing import Any, ClassVar, List, Optional, Tuple, Union -from kornia.core import Module, ONNXExportMixin, Tensor, rand +from kornia.core import Module, ONNXExportMixin, Tensor, tensor, rand __all__ = ["BoxFiltering"] @@ -10,6 +10,7 @@ class BoxFiltering(Module, ONNXExportMixin): Args: confidence_threshold: an 0-d scalar that represents the desired threshold. + classes_to_keep: a 1-d list of classes to keep. If None, keep all classes. filter_as_zero: whether to filter boxes as zero. """ @@ -17,35 +18,72 @@ class BoxFiltering(Module, ONNXExportMixin): ONNX_DEFAULT_OUTPUTSHAPE: ClassVar[List[int]] = [-1, -1, 6] ONNX_EXPORT_PSEUDO_SHAPE: ClassVar[List[int]] = [5, 20, 6] - def __init__(self, confidence_threshold: Optional[Tensor] = None, filter_as_zero: bool = False) -> None: + def __init__( + self, + confidence_threshold: Optional[Union[Tensor, float]] = None, + classes_to_keep: Optional[Union[Tensor, List[int]]] = None, + filter_as_zero: bool = False + ) -> None: super().__init__() self.filter_as_zero = filter_as_zero - self.confidence_threshold = confidence_threshold + self.classes_to_keep = None + self.confidence_threshold = None + if classes_to_keep is not None: + self.classes_to_keep = classes_to_keep if isinstance( + classes_to_keep, Tensor) else tensor(classes_to_keep) + if confidence_threshold is not None: + self.confidence_threshold = confidence_threshold or confidence_threshold if isinstance( + confidence_threshold, Tensor) else tensor(confidence_threshold) - def forward(self, boxes: Tensor, confidence_threshold: Optional[Tensor] = None) -> Union[Tensor, List[Tensor]]: + def forward( + self, boxes: Tensor, confidence_threshold: Optional[Tensor] = None, classes_to_keep: Optional[Tensor] = None + ) -> Union[Tensor, List[Tensor]]: """Filter boxes according to the desired threshold. + + To be ONNX-friendly, the inputs for direct forwarding need to be all tensors. Args: boxes: [B, D, 6], where B is the batchsize, D is the number of detections in the image, 6 represent (class_id, confidence_score, x, y, w, h). confidence_threshold: an 0-d scalar that represents the desired threshold. + classes_to_keep: a 1-d tensor of classes to keep. If None, keep all classes. + + Returns: + Union[Tensor, List[Tensor]] + If `filter_as_zero` is True, return a tensor of shape [D, 6], where D is the total number of + detections as input. + If `filter_as_zero` is False, return a list of tensors of shape [D, 6], where D is the number of + valid detections for each element in the batch. """ - if confidence_threshold is None: - confidence_threshold = self.confidence_threshold - if confidence_threshold is None: - raise ValueError("`confidence_threshold` must be provided if not set in the constructor.") + # Apply confidence filtering + confidence_threshold = confidence_threshold or self.confidence_threshold or 0. # If None, use 0 as threshold + confidence_mask = boxes[:, :, 1] > confidence_threshold # [B, D] - filtered_boxes: Union[Tensor, List[Tensor]] - if self.filter_as_zero: - box_mask = (boxes[:, :, 1] > confidence_threshold).unsqueeze(-1).expand_as(boxes) - filtered_boxes = boxes * box_mask.float() + # Apply class filtering + classes_to_keep = classes_to_keep or self.classes_to_keep + if classes_to_keep is not None: + class_ids = boxes[:, :, 0:1] # [B, D, 1] + classes_to_keep = classes_to_keep.view(1, 1, -1) # [1, 1, C] for broadcasting + class_mask = (class_ids == classes_to_keep).any(dim=-1) # [B, D] else: - filtered_boxes = [] - for i in range(boxes.shape[0]): - box = boxes[i][(boxes[i, :, 1] > confidence_threshold).unsqueeze(-1).expand_as(boxes[i])] - filtered_boxes.append(box.view(-1, boxes.shape[-1])) + # If no class filtering is needed, just use a mask of all `True` + class_mask = (confidence_mask * 0 + 1).bool() + + # Combine the confidence and class masks + combined_mask = confidence_mask & class_mask # [B, D] + + if self.filter_as_zero: + filtered_boxes = boxes * combined_mask[:, :, None] + return filtered_boxes + + filtered_boxes_list = [] + for i in range(boxes.shape[0]): + box = boxes[i] + mask = combined_mask[i] # [D] + valid_boxes = box[mask] + filtered_boxes_list.append(valid_boxes) - return filtered_boxes + return filtered_boxes_list def _create_dummy_input(self, input_shape: List[int]) -> Union[Tuple[Any, ...], Tensor]: pseudo_input = rand( diff --git a/kornia/models/segmentation/segmentation_models.py b/kornia/models/segmentation/segmentation_models.py index dc0469feff..78693e84a4 100644 --- a/kornia/models/segmentation/segmentation_models.py +++ b/kornia/models/segmentation/segmentation_models.py @@ -53,9 +53,12 @@ def __init__( def preprocessing(self, input: Tensor) -> Tensor: # Ensure the color space transformation is ONNX-friendly input_space = self.preproc_params["input_space"] - input = ( - kornia.color.rgb.rgb_to_bgr(input) if input_space == "BGR" else input - ) # Assume input is already RGB if not BGR + if input_space == "BGR": + input = kornia.color.rgb.bgr_to_rgb(input) + elif input_space == "RGB": + pass + else: + raise ValueError(f"Unsupported input space: {input_space}") # Normalize input range if needed input_range = self.preproc_params["input_range"] diff --git a/kornia/onnx/sequential.py b/kornia/onnx/sequential.py index 8c90e5eed1..1918b2eb99 100644 --- a/kornia/onnx/sequential.py +++ b/kornia/onnx/sequential.py @@ -1,4 +1,4 @@ -from typing import List, Optional, Tuple, Union +from typing import Any, List, Optional, Tuple, Union from kornia.config import kornia_config from kornia.core.external import numpy as np @@ -138,9 +138,42 @@ def as_cpu(self) -> None: """Set the session to run on CPU.""" self._session.set_providers(["CPUExecutionProvider"]) - def as_cuda(self, device: int = 0) -> None: - """Set the session to run on CUDA.""" - self._session.set_providers(["CUDAExecutionProvider"], provider_options=[{"device_id": device}]) + def as_cuda(self, device_id: int = 0, **kwargs: Any) -> None: + """Set the session to run on CUDA. + + We set the ONNX runtime session to use CUDAExecutionProvider. For other CUDAExecutionProvider configurations, + or CUDA/cuDNN/ONNX version issues, + you may refer to https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html. + + Args: + device_id: Select GPU to execute. + """ + self._session.set_providers(["CUDAExecutionProvider"], provider_options=[{"device_id": device_id, **kwargs}]) + + def as_tensorrt(self, device_id: int = 0, **kwargs: Any) -> None: + """Set the session to run on TensorRT. + + We set the ONNX runtime session to use TensorrtExecutionProvider. For other TensorrtExecutionProvider configurations, + or CUDA/cuDNN/ONNX/TensorRT version issues, + you may refer to https://onnxruntime.ai/docs/execution-providers/TensorRT-ExecutionProvider.html. + + Args: + device_id: select GPU to execute. + """ + self._session.set_providers(["TensorrtExecutionProvider"], provider_options=[{"device_id": device_id, **kwargs}]) + + def as_openvino(self, device_type: str = "GPU", **kwargs: Any) -> None: + """Set the session to run on TensorRT. + + We set the ONNX runtime session to use OpenVINOExecutionProvider. For other OpenVINOExecutionProvider configurations, + or CUDA/cuDNN/ONNX/TensorRT version issues, + you may refer to https://onnxruntime.ai/docs/execution-providers/OpenVINO-ExecutionProvider.html. + + Args: + device_type: CPU, NPU, GPU, GPU.0, GPU.1 based on the avaialable GPUs, NPU, Any valid Hetero combination, + Any valid Multi or Auto devices combination. + """ + self._session.set_providers(["OpenVINOExecutionProvider"], provider_options=[{"device_type": device_type, **kwargs}]) def __call__(self, *inputs: "np.ndarray") -> List["np.ndarray"]: # type:ignore """Perform inference using the combined ONNX model. diff --git a/tests/models/box_filtering.py b/tests/models/box_filtering.py new file mode 100644 index 0000000000..1ca33d0fcd --- /dev/null +++ b/tests/models/box_filtering.py @@ -0,0 +1,90 @@ +import torch +import pytest +from numpy.testing import assert_almost_equal +from kornia.core import tensor +from kornia.models.detection.utils import BoxFiltering # Replace with the actual module path + +class TestBoxFiltering: + @pytest.fixture + def sample_boxes(self): + # Setup some sample boxes with the format [class_id, confidence_score, x, y, w, h] + return tensor([ + [[1, 0.9, 10, 10, 20, 20], # High confidence, class 1 + [2, 0.7, 15, 15, 25, 25], # Medium confidence, class 2 + [3, 0.7, 15, 15, 25, 25], # Medium confidence, class 3 + [4, 0.3, 5, 5, 10, 10]], # Low confidence, class 4 + [[1, 0.95, 12, 12, 18, 18], # High confidence, class 1 + [2, 0.5, 13, 13, 20, 20], # Low confidence, class 2 + [3, 0.5, 13, 13, 20, 20], # Low confidence, class 3 + [4, 0.2, 7, 7, 14, 14]], # Very low confidence, class 4 + [[1, 0.1, 12, 12, 18, 18], # Very Low confidence, class 1 + [2, 0.1, 13, 13, 20, 20], # Very Low confidence, class 2 + [3, 0.1, 13, 13, 20, 20], # Very Low confidence, class 3 + [4, 0.1, 7, 7, 14, 14]] # Very Low confidence, class 4 + ]) # Shape: [3, 4, 6], i.e., [B, D, 6] + + def test_confidence_filtering(self, sample_boxes): + """Test filtering based on confidence threshold.""" + # Set a confidence threshold of 0.7 + filter = BoxFiltering(confidence_threshold=0.7) + filtered_boxes = filter(sample_boxes) + + # Expected output: only boxes with confidence > 0.7 should be kept + assert len(filtered_boxes[0]) == 1 # Only one box in the first batch + assert_almost_equal(filtered_boxes[0][0][1].item(), 0.9) # Box with confidence 0.9 + assert len(filtered_boxes[1]) == 1 # Only one box in the second batch + assert_almost_equal(filtered_boxes[1][0][1].item(), 0.95) # Box with confidence 0.95 + assert len(filtered_boxes[2]) == 0 # No boxes in the third batch + + def test_class_filtering(self, sample_boxes): + """Test filtering based on class IDs.""" + # Set classes_to_keep to [1, 2] + filter = BoxFiltering(classes_to_keep=tensor([1, 2])) + filtered_boxes = filter(sample_boxes) + + # Expected output: only boxes with class_id 1 and 2 should be kept + assert len(filtered_boxes[0]) == 2 # Two boxes in the first batch + assert filtered_boxes[0][0][0].item() == 1 # Box with class_id 1 + assert filtered_boxes[0][1][0].item() == 2 # Box with class_id 2 + assert len(filtered_boxes[1]) == 2 # Two boxes in the second batch + assert filtered_boxes[1][0][0].item() == 1 # Box with class_id 1 + assert filtered_boxes[1][1][0].item() == 2 # Box with class_id 2 + assert len(filtered_boxes[2]) == 2 # Two boxes in the third batch + assert filtered_boxes[2][0][0].item() == 1 # Box with class_id 1 + assert filtered_boxes[2][1][0].item() == 2 # Box with class_id 2 + + def test_combined_confidence_and_class_filtering(self, sample_boxes): + """Test filtering based on both confidence and class IDs.""" + # Set confidence threshold to 0.6 and classes_to_keep to [1, 3] + filter = BoxFiltering(confidence_threshold=0.6, classes_to_keep=tensor([1, 3])) + filtered_boxes = filter(sample_boxes) + + # Expected output: only boxes with confidence > 0.6 and class_id in [1, 3] should be kept + assert len(filtered_boxes[0]) == 2 # Two boxes in the first batch + assert filtered_boxes[0][0][0].item() == 1 # Class_id 1 + assert filtered_boxes[0][1][0].item() == 3 # Class_id 3 + assert filtered_boxes[1][0][0].item() == 1 # Class_id 1 + assert len(filtered_boxes[1]) == 1 # No boxes in the second batch + assert len(filtered_boxes[2]) == 0 # No boxes in the third batch + + def test_filter_as_zero(self, sample_boxes): + """Test filtering boxes as zero when filter_as_zero is True.""" + filter = BoxFiltering(confidence_threshold=0.8, filter_as_zero=True) + filtered_boxes = filter(sample_boxes) + + # Expected output: boxes with confidence <= 0.8 should be zeroed out + assert torch.all(filtered_boxes[0][0] != 0) # Box with confidence 0.9 should remain + assert torch.all(filtered_boxes[0][1:] == 0) # Remaining boxes should be zeroed + assert torch.all(filtered_boxes[1][0] != 0) # Box with confidence 0.95 should remain + assert torch.all(filtered_boxes[1][1:] == 0) # Remaining boxes should be zeroed + assert torch.all(filtered_boxes[2] == 0) # All boxes in the third batch should be zeroed + + def test_no_class_or_confidence_filtering(self, sample_boxes): + """Test when no class or confidence filtering is applied.""" + filter = BoxFiltering() # No thresholds set + filtered_boxes = filter(sample_boxes) + + # Expected output: all boxes should be returned as-is + assert len(filtered_boxes[0]) == 4 # All boxes in the first batch should be kept + assert len(filtered_boxes[1]) == 4 # All boxes in the second batch should be kept + assert len(filtered_boxes[2]) == 4 # All boxes in the third batch should be kept From 3aef7d8128fa9d69166f34700efd2ddeadcf0bed Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 19 Sep 2024 22:46:15 +0000 Subject: [PATCH 84/84] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- kornia/models/detection/utils.py | 20 +++++++------ kornia/onnx/sequential.py | 14 ++++++---- tests/models/box_filtering.py | 48 +++++++++++++++++++------------- 3 files changed, 49 insertions(+), 33 deletions(-) diff --git a/kornia/models/detection/utils.py b/kornia/models/detection/utils.py index 6ec0d823e2..8b23686246 100644 --- a/kornia/models/detection/utils.py +++ b/kornia/models/detection/utils.py @@ -1,6 +1,6 @@ from typing import Any, ClassVar, List, Optional, Tuple, Union -from kornia.core import Module, ONNXExportMixin, Tensor, tensor, rand +from kornia.core import Module, ONNXExportMixin, Tensor, rand, tensor __all__ = ["BoxFiltering"] @@ -22,24 +22,26 @@ def __init__( self, confidence_threshold: Optional[Union[Tensor, float]] = None, classes_to_keep: Optional[Union[Tensor, List[int]]] = None, - filter_as_zero: bool = False + filter_as_zero: bool = False, ) -> None: super().__init__() self.filter_as_zero = filter_as_zero self.classes_to_keep = None self.confidence_threshold = None if classes_to_keep is not None: - self.classes_to_keep = classes_to_keep if isinstance( - classes_to_keep, Tensor) else tensor(classes_to_keep) + self.classes_to_keep = classes_to_keep if isinstance(classes_to_keep, Tensor) else tensor(classes_to_keep) if confidence_threshold is not None: - self.confidence_threshold = confidence_threshold or confidence_threshold if isinstance( - confidence_threshold, Tensor) else tensor(confidence_threshold) + self.confidence_threshold = ( + confidence_threshold or confidence_threshold + if isinstance(confidence_threshold, Tensor) + else tensor(confidence_threshold) + ) def forward( self, boxes: Tensor, confidence_threshold: Optional[Tensor] = None, classes_to_keep: Optional[Tensor] = None ) -> Union[Tensor, List[Tensor]]: """Filter boxes according to the desired threshold. - + To be ONNX-friendly, the inputs for direct forwarding need to be all tensors. Args: @@ -47,7 +49,7 @@ def forward( 6 represent (class_id, confidence_score, x, y, w, h). confidence_threshold: an 0-d scalar that represents the desired threshold. classes_to_keep: a 1-d tensor of classes to keep. If None, keep all classes. - + Returns: Union[Tensor, List[Tensor]] If `filter_as_zero` is True, return a tensor of shape [D, 6], where D is the total number of @@ -56,7 +58,7 @@ def forward( valid detections for each element in the batch. """ # Apply confidence filtering - confidence_threshold = confidence_threshold or self.confidence_threshold or 0. # If None, use 0 as threshold + confidence_threshold = confidence_threshold or self.confidence_threshold or 0.0 # If None, use 0 as threshold confidence_mask = boxes[:, :, 1] > confidence_threshold # [B, D] # Apply class filtering diff --git a/kornia/onnx/sequential.py b/kornia/onnx/sequential.py index 1918b2eb99..52a3806d83 100644 --- a/kornia/onnx/sequential.py +++ b/kornia/onnx/sequential.py @@ -140,7 +140,7 @@ def as_cpu(self) -> None: def as_cuda(self, device_id: int = 0, **kwargs: Any) -> None: """Set the session to run on CUDA. - + We set the ONNX runtime session to use CUDAExecutionProvider. For other CUDAExecutionProvider configurations, or CUDA/cuDNN/ONNX version issues, you may refer to https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html. @@ -152,7 +152,7 @@ def as_cuda(self, device_id: int = 0, **kwargs: Any) -> None: def as_tensorrt(self, device_id: int = 0, **kwargs: Any) -> None: """Set the session to run on TensorRT. - + We set the ONNX runtime session to use TensorrtExecutionProvider. For other TensorrtExecutionProvider configurations, or CUDA/cuDNN/ONNX/TensorRT version issues, you may refer to https://onnxruntime.ai/docs/execution-providers/TensorRT-ExecutionProvider.html. @@ -160,11 +160,13 @@ def as_tensorrt(self, device_id: int = 0, **kwargs: Any) -> None: Args: device_id: select GPU to execute. """ - self._session.set_providers(["TensorrtExecutionProvider"], provider_options=[{"device_id": device_id, **kwargs}]) + self._session.set_providers( + ["TensorrtExecutionProvider"], provider_options=[{"device_id": device_id, **kwargs}] + ) def as_openvino(self, device_type: str = "GPU", **kwargs: Any) -> None: """Set the session to run on TensorRT. - + We set the ONNX runtime session to use OpenVINOExecutionProvider. For other OpenVINOExecutionProvider configurations, or CUDA/cuDNN/ONNX/TensorRT version issues, you may refer to https://onnxruntime.ai/docs/execution-providers/OpenVINO-ExecutionProvider.html. @@ -173,7 +175,9 @@ def as_openvino(self, device_type: str = "GPU", **kwargs: Any) -> None: device_type: CPU, NPU, GPU, GPU.0, GPU.1 based on the avaialable GPUs, NPU, Any valid Hetero combination, Any valid Multi or Auto devices combination. """ - self._session.set_providers(["OpenVINOExecutionProvider"], provider_options=[{"device_type": device_type, **kwargs}]) + self._session.set_providers( + ["OpenVINOExecutionProvider"], provider_options=[{"device_type": device_type, **kwargs}] + ) def __call__(self, *inputs: "np.ndarray") -> List["np.ndarray"]: # type:ignore """Perform inference using the combined ONNX model. diff --git a/tests/models/box_filtering.py b/tests/models/box_filtering.py index 1ca33d0fcd..6d489d27a6 100644 --- a/tests/models/box_filtering.py +++ b/tests/models/box_filtering.py @@ -1,27 +1,37 @@ -import torch import pytest +import torch from numpy.testing import assert_almost_equal + from kornia.core import tensor from kornia.models.detection.utils import BoxFiltering # Replace with the actual module path + class TestBoxFiltering: @pytest.fixture def sample_boxes(self): # Setup some sample boxes with the format [class_id, confidence_score, x, y, w, h] - return tensor([ - [[1, 0.9, 10, 10, 20, 20], # High confidence, class 1 - [2, 0.7, 15, 15, 25, 25], # Medium confidence, class 2 - [3, 0.7, 15, 15, 25, 25], # Medium confidence, class 3 - [4, 0.3, 5, 5, 10, 10]], # Low confidence, class 4 - [[1, 0.95, 12, 12, 18, 18], # High confidence, class 1 - [2, 0.5, 13, 13, 20, 20], # Low confidence, class 2 - [3, 0.5, 13, 13, 20, 20], # Low confidence, class 3 - [4, 0.2, 7, 7, 14, 14]], # Very low confidence, class 4 - [[1, 0.1, 12, 12, 18, 18], # Very Low confidence, class 1 - [2, 0.1, 13, 13, 20, 20], # Very Low confidence, class 2 - [3, 0.1, 13, 13, 20, 20], # Very Low confidence, class 3 - [4, 0.1, 7, 7, 14, 14]] # Very Low confidence, class 4 - ]) # Shape: [3, 4, 6], i.e., [B, D, 6] + return tensor( + [ + [ + [1, 0.9, 10, 10, 20, 20], # High confidence, class 1 + [2, 0.7, 15, 15, 25, 25], # Medium confidence, class 2 + [3, 0.7, 15, 15, 25, 25], # Medium confidence, class 3 + [4, 0.3, 5, 5, 10, 10], + ], # Low confidence, class 4 + [ + [1, 0.95, 12, 12, 18, 18], # High confidence, class 1 + [2, 0.5, 13, 13, 20, 20], # Low confidence, class 2 + [3, 0.5, 13, 13, 20, 20], # Low confidence, class 3 + [4, 0.2, 7, 7, 14, 14], + ], # Very low confidence, class 4 + [ + [1, 0.1, 12, 12, 18, 18], # Very Low confidence, class 1 + [2, 0.1, 13, 13, 20, 20], # Very Low confidence, class 2 + [3, 0.1, 13, 13, 20, 20], # Very Low confidence, class 3 + [4, 0.1, 7, 7, 14, 14], + ], # Very Low confidence, class 4 + ] + ) # Shape: [3, 4, 6], i.e., [B, D, 6] def test_confidence_filtering(self, sample_boxes): """Test filtering based on confidence threshold.""" @@ -41,7 +51,7 @@ def test_class_filtering(self, sample_boxes): # Set classes_to_keep to [1, 2] filter = BoxFiltering(classes_to_keep=tensor([1, 2])) filtered_boxes = filter(sample_boxes) - + # Expected output: only boxes with class_id 1 and 2 should be kept assert len(filtered_boxes[0]) == 2 # Two boxes in the first batch assert filtered_boxes[0][0][0].item() == 1 # Box with class_id 1 @@ -65,13 +75,13 @@ def test_combined_confidence_and_class_filtering(self, sample_boxes): assert filtered_boxes[0][1][0].item() == 3 # Class_id 3 assert filtered_boxes[1][0][0].item() == 1 # Class_id 1 assert len(filtered_boxes[1]) == 1 # No boxes in the second batch - assert len(filtered_boxes[2]) == 0 # No boxes in the third batch + assert len(filtered_boxes[2]) == 0 # No boxes in the third batch def test_filter_as_zero(self, sample_boxes): """Test filtering boxes as zero when filter_as_zero is True.""" filter = BoxFiltering(confidence_threshold=0.8, filter_as_zero=True) filtered_boxes = filter(sample_boxes) - + # Expected output: boxes with confidence <= 0.8 should be zeroed out assert torch.all(filtered_boxes[0][0] != 0) # Box with confidence 0.9 should remain assert torch.all(filtered_boxes[0][1:] == 0) # Remaining boxes should be zeroed @@ -83,7 +93,7 @@ def test_no_class_or_confidence_filtering(self, sample_boxes): """Test when no class or confidence filtering is applied.""" filter = BoxFiltering() # No thresholds set filtered_boxes = filter(sample_boxes) - + # Expected output: all boxes should be returned as-is assert len(filtered_boxes[0]) == 4 # All boxes in the first batch should be kept assert len(filtered_boxes[1]) == 4 # All boxes in the second batch should be kept