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