diff --git a/src/otx/algo/__init__.py b/src/otx/algo/__init__.py index 968d579e5f7..29312f92f25 100644 --- a/src/otx/algo/__init__.py +++ b/src/otx/algo/__init__.py @@ -3,6 +3,24 @@ # """Module for OTX custom algorithms, e.g., model, losses, hook, etc...""" -from . import action_classification, classification, detection, segmentation, visual_prompting +from . import ( + accelerators, + action_classification, + classification, + detection, + plugins, + segmentation, + strategies, + visual_prompting, +) -__all__ = ["action_classification", "classification", "detection", "segmentation", "visual_prompting"] +__all__ = [ + "action_classification", + "classification", + "detection", + "segmentation", + "visual_prompting", + "strategies", + "accelerators", + "plugins", +] diff --git a/src/otx/algo/accelerators/__init__.py b/src/otx/algo/accelerators/__init__.py new file mode 100644 index 00000000000..5fc4b9d9d1d --- /dev/null +++ b/src/otx/algo/accelerators/__init__.py @@ -0,0 +1,8 @@ +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# +"""Lightning accelerator for XPU device.""" + +from .xpu import XPUAccelerator + +__all__ = ["XPUAccelerator"] diff --git a/src/otx/algo/accelerators/xpu.py b/src/otx/algo/accelerators/xpu.py new file mode 100644 index 00000000000..f5969336ab4 --- /dev/null +++ b/src/otx/algo/accelerators/xpu.py @@ -0,0 +1,88 @@ +"""Lightning accelerator for XPU device.""" +# Copyright (C) 2023 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# +from __future__ import annotations + +from typing import Any, Union + +import numpy as np +import torch +from lightning.pytorch.accelerators import AcceleratorRegistry +from lightning.pytorch.accelerators.accelerator import Accelerator +from mmcv.ops.nms import NMSop +from mmcv.ops.roi_align import RoIAlign +from mmengine.structures import instance_data + +from otx.algo.detection.utils import monkey_patched_nms, monkey_patched_roi_align +from otx.utils.utils import is_xpu_available + + +class XPUAccelerator(Accelerator): + """Support for a XPU, optimized for large-scale machine learning.""" + + accelerator_name = "xpu" + + def setup_device(self, device: torch.device) -> None: + """Sets up the specified device.""" + if device.type != "xpu": + msg = f"Device should be xpu, got {device} instead" + raise RuntimeError(msg) + + torch.xpu.set_device(device) + self.patch_packages_xpu() + + @staticmethod + def parse_devices(devices: str | list | torch.device) -> list: + """Parses devices for multi-GPU training.""" + if isinstance(devices, list): + return devices + return [devices] + + @staticmethod + def get_parallel_devices(devices: list) -> list[torch.device]: + """Generates a list of parrallel devices.""" + return [torch.device("xpu", idx) for idx in devices] + + @staticmethod + def auto_device_count() -> int: + """Returns number of XPU devices available.""" + return torch.xpu.device_count() + + @staticmethod + def is_available() -> bool: + """Checks if XPU available.""" + return is_xpu_available() + + def get_device_stats(self, device: str | torch.device) -> dict[str, Any]: + """Returns XPU devices stats.""" + return {} + + def teardown(self) -> None: + """Cleans-up XPU-related resources.""" + self.revert_packages_xpu() + + def patch_packages_xpu(self) -> None: + """Patch packages when xpu is available.""" + # patch instance_data from mmengie + long_type_tensor = Union[torch.LongTensor, torch.xpu.LongTensor] + bool_type_tensor = Union[torch.BoolTensor, torch.xpu.BoolTensor] + instance_data.IndexType = Union[str, slice, int, list, long_type_tensor, bool_type_tensor, np.ndarray] + + # patch nms and roi_align + self._nms_op_forward = NMSop.forward + self._roi_align_forward = RoIAlign.forward + NMSop.forward = monkey_patched_nms + RoIAlign.forward = monkey_patched_roi_align + + def revert_packages_xpu(self) -> None: + """Revert packages when xpu is available.""" + NMSop.forward = self._nms_op_forward + RoIAlign.forward = self._roi_align_forward + + +AcceleratorRegistry.register( + XPUAccelerator.accelerator_name, + XPUAccelerator, + description="Accelerator supports XPU devices", +) diff --git a/src/otx/algo/detection/utils/__init__.py b/src/otx/algo/detection/utils/__init__.py index a8647baa335..2ab46a64ac4 100644 --- a/src/otx/algo/detection/utils/__init__.py +++ b/src/otx/algo/detection/utils/__init__.py @@ -1,3 +1,8 @@ # Copyright (C) 2024 Intel Corporation # SPDX-License-Identifier: Apache-2.0 -"""Data structures for detection task.""" +# +"""utils for detection task.""" + +from .mmcv_patched_ops import monkey_patched_nms, monkey_patched_roi_align + +__all__ = ["monkey_patched_nms", "monkey_patched_roi_align"] diff --git a/src/otx/algo/detection/utils/mmcv_patched_ops.py b/src/otx/algo/detection/utils/mmcv_patched_ops.py new file mode 100644 index 00000000000..ec3a884232d --- /dev/null +++ b/src/otx/algo/detection/utils/mmcv_patched_ops.py @@ -0,0 +1,73 @@ +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# +"""utils for detection task.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from mmcv.utils import ext_loader +from torchvision.ops import nms as tv_nms +from torchvision.ops import roi_align as tv_roi_align + +if TYPE_CHECKING: + from mmcv.ops.nms import NMSop + from mmcv.ops.roi_align import RoIAlign + +ext_module = ext_loader.load_ext("_ext", ["nms", "softnms", "nms_match", "nms_rotated", "nms_quadri"]) + + +def monkey_patched_nms( + ctx: NMSop, + bboxes: torch.Tensor, + scores: torch.Tensor, + iou_threshold: float, + offset: float, + score_threshold: float, + max_num: int, +) -> torch.Tensor: + """Runs MMCVs NMS with torchvision.nms, or forces NMS from MMCV to run on CPU.""" + _ = ctx + is_filtering_by_score = score_threshold > 0 + if is_filtering_by_score: + valid_mask = scores > score_threshold + bboxes, scores = bboxes[valid_mask], scores[valid_mask] + valid_inds = torch.nonzero(valid_mask, as_tuple=False).squeeze(dim=1) + + if bboxes.dtype == torch.bfloat16: + bboxes = bboxes.to(torch.float32) + if scores.dtype == torch.bfloat16: + scores = scores.to(torch.float32) + + if offset == 0: + inds = tv_nms(bboxes, scores, float(iou_threshold)) + else: + device = bboxes.device + bboxes = bboxes.to("cpu") + scores = scores.to("cpu") + inds = ext_module.nms(bboxes, scores, iou_threshold=float(iou_threshold), offset=offset) + bboxes = bboxes.to(device) + scores = scores.to(device) + + if max_num > 0: + inds = inds[:max_num] + if is_filtering_by_score: + inds = valid_inds[inds] + return inds + + +def monkey_patched_roi_align(self: RoIAlign, _input: torch.Tensor, rois: torch.Tensor) -> torch.Tensor: + """Replaces MMCVs roi align with the one from torchvision. + + Args: + self: patched instance + _input: NCHW images + rois: Bx5 boxes. First column is the index into N. The other 4 columns are xyxy. + """ + if "aligned" in tv_roi_align.__code__.co_varnames: + return tv_roi_align(_input, rois, self.output_size, self.spatial_scale, self.sampling_ratio, self.aligned) + if self.aligned: + rois -= rois.new_tensor([0.0] + [0.5 / self.spatial_scale] * 4) + return tv_roi_align(_input, rois, self.output_size, self.spatial_scale, self.sampling_ratio) diff --git a/src/otx/algo/plugins/__init__.py b/src/otx/algo/plugins/__init__.py new file mode 100644 index 00000000000..91be640aea6 --- /dev/null +++ b/src/otx/algo/plugins/__init__.py @@ -0,0 +1,8 @@ +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# +"""Plugin for mixed-precision training on XPU.""" + +from .xpu_precision import MixedPrecisionXPUPlugin + +__all__ = ["MixedPrecisionXPUPlugin"] diff --git a/src/otx/algo/plugins/xpu_precision.py b/src/otx/algo/plugins/xpu_precision.py new file mode 100644 index 00000000000..fb2a08eb182 --- /dev/null +++ b/src/otx/algo/plugins/xpu_precision.py @@ -0,0 +1,117 @@ +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# +"""Plugin for mixed-precision training on XPU.""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any, Callable, Generator + +import torch +from lightning.pytorch.plugins.precision.precision import Precision +from lightning.pytorch.utilities import GradClipAlgorithmType +from lightning.pytorch.utilities.exceptions import MisconfigurationException +from torch import Tensor +from torch.optim import LBFGS, Optimizer + +if TYPE_CHECKING: + import lightning.pytorch as pl + from lightning_fabric.utilities.types import Optimizable + + +class MixedPrecisionXPUPlugin(Precision): + """Plugin for Automatic Mixed Precision (AMP) training with ``torch.xpu.autocast``. + + Args: + scaler: An optional :class:`torch.cuda.amp.GradScaler` to use. + """ + + def __init__(self, scaler: torch.cuda.amp.GradScaler | None = None) -> None: + self.scaler = scaler + + def pre_backward(self, tensor: Tensor, module: pl.LightningModule) -> Tensor: + """Apply grad scaler before backward.""" + if self.scaler is not None: + tensor = self.scaler.scale(tensor) + return super().pre_backward(tensor, module) + + def optimizer_step( # type: ignore[override] + self, + optimizer: Optimizable, + model: pl.LightningModule, + closure: Callable, + **kwargs: dict, + ) -> None | dict: + """Make an optimizer step using scaler if it was passed.""" + if self.scaler is None: + # skip scaler logic, as bfloat16 does not require scaler + return super().optimizer_step( + optimizer, + model=model, + closure=closure, + **kwargs, + ) + if isinstance(optimizer, LBFGS): + msg = "Native AMP and the LBFGS optimizer are not compatible." + raise MisconfigurationException( + msg, + ) + closure_result = closure() + + if not _optimizer_handles_unscaling(optimizer): + # Unscaling needs to be performed here in case we are going to apply gradient clipping. + # Optimizers that perform unscaling in their `.step()` method are not supported (e.g., fused Adam). + # Note: `unscale` happens after the closure is executed, but before the `on_before_optimizer_step` hook. + self.scaler.unscale_(optimizer) + + self._after_closure(model, optimizer) + skipped_backward = closure_result is None + # in manual optimization, the closure does not return a value + if not model.automatic_optimization or not skipped_backward: + # note: the scaler will skip the `optimizer.step` if nonfinite gradients are found + step_output = self.scaler.step(optimizer, **kwargs) + self.scaler.update() + return step_output + return closure_result + + def clip_gradients( + self, + optimizer: Optimizer, + clip_val: int | float = 0.0, + gradient_clip_algorithm: GradClipAlgorithmType = GradClipAlgorithmType.NORM, + ) -> None: + """Handle grad clipping with scaler.""" + if clip_val > 0 and _optimizer_handles_unscaling(optimizer): + msg = f"The current optimizer, {type(optimizer).__qualname__}, does not allow for gradient clipping" + " because it performs unscaling of gradients internally. HINT: Are you using a 'fused' optimizer?" + raise RuntimeError( + msg, + ) + super().clip_gradients(optimizer=optimizer, clip_val=clip_val, gradient_clip_algorithm=gradient_clip_algorithm) + + @contextmanager + def forward_context(self) -> Generator[None, None, None]: + """Enable autocast context.""" + with torch.xpu.autocast(True): + yield + + def state_dict(self) -> dict[str, Any]: + """Returns state dict of the plugin.""" + if self.scaler is not None: + return self.scaler.state_dict() + return {} + + def load_state_dict(self, state_dict: dict[str, torch.Tensor]) -> None: + """Loads state dict to the plugin.""" + if self.scaler is not None: + self.scaler.load_state_dict(state_dict) + + +def _optimizer_handles_unscaling(optimizer: torch.optim.Optimizer) -> bool: + """Determines if a PyTorch optimizer handles unscaling gradients in the step method ratherthan through the scaler. + + Since, the current implementation of this function checks a PyTorch internal variable on the optimizer, the return + value will only be reliable for built-in PyTorch optimizers. + """ + return getattr(optimizer, "_step_supports_amp_scaling", False) diff --git a/src/otx/algo/strategies/__init__.py b/src/otx/algo/strategies/__init__.py new file mode 100644 index 00000000000..392a1b82b22 --- /dev/null +++ b/src/otx/algo/strategies/__init__.py @@ -0,0 +1,8 @@ +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# +"""Lightning strategy for single XPU device.""" + +from .xpu_single import SingleXPUStrategy + +__all__ = ["SingleXPUStrategy"] diff --git a/src/otx/algo/strategies/xpu_single.py b/src/otx/algo/strategies/xpu_single.py new file mode 100644 index 00000000000..4b9501dd36f --- /dev/null +++ b/src/otx/algo/strategies/xpu_single.py @@ -0,0 +1,73 @@ +"""Lightning strategy for single XPU device.""" + +# Copyright (C) 2023 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from lightning.pytorch.strategies import StrategyRegistry +from lightning.pytorch.strategies.single_device import SingleDeviceStrategy +from lightning.pytorch.utilities.exceptions import MisconfigurationException + +from otx.utils.utils import is_xpu_available + +if TYPE_CHECKING: + import lightning.pytorch as pl + from lightning.pytorch.plugins.precision import PrecisionPlugin + from lightning_fabric.plugins import CheckpointIO + from lightning_fabric.utilities.types import _DEVICE + + +class SingleXPUStrategy(SingleDeviceStrategy): + """Strategy for training on single XPU device.""" + + strategy_name = "xpu_single" + + def __init__( + self, + device: _DEVICE = "xpu:0", + accelerator: pl.accelerators.Accelerator | None = None, + checkpoint_io: CheckpointIO | None = None, + precision_plugin: PrecisionPlugin | None = None, + ): + if not is_xpu_available(): + msg = "`SingleXPUStrategy` requires XPU devices to run" + raise MisconfigurationException(msg) + + super().__init__( + accelerator=accelerator, + device=device, + checkpoint_io=checkpoint_io, + precision_plugin=precision_plugin, + ) + + @property + def is_distributed(self) -> bool: + """Returns true if the strategy supports distributed training.""" + return False + + def setup_optimizers(self, trainer: pl.Trainer) -> None: + """Sets up optimizers.""" + super().setup_optimizers(trainer) + if len(self.optimizers) > 1: # type: ignore[has-type] + msg = "XPU strategy doesn't support multiple optimizers" + raise RuntimeError(msg) + if trainer.task != "SEMANTIC_SEGMENTATION": + if len(self.optimizers) == 1: # type: ignore[has-type] + model, optimizer = torch.xpu.optimize(trainer.model, optimizer=self.optimizers[0]) # type: ignore[has-type] + self.optimizers = [optimizer] + self.model = model + else: # for inference + trainer.model.eval() + self.model = torch.xpu.optimize(trainer.model) + + +StrategyRegistry.register( + SingleXPUStrategy.strategy_name, + SingleXPUStrategy, + description="Strategy that enables training on single XPU", +) diff --git a/src/otx/cli/install.py b/src/otx/cli/install.py index 7d0fd49d45f..ae9454089dd 100644 --- a/src/otx/cli/install.py +++ b/src/otx/cli/install.py @@ -64,10 +64,15 @@ def add_install_parser(subcommands_action: _ActionSubCommands) -> None: help="Do not install PyTorch. Choose this option if you already install PyTorch.", action="store_true", ) + subcommands_action.add_subcommand("install", parser, help="Install OTX requirements.") -def otx_install(option: str | None = None, verbose: bool = False, do_not_install_torch: bool = False) -> int: +def otx_install( + option: str | None = None, + verbose: bool = False, + do_not_install_torch: bool = False, +) -> int: """Install OTX requirements. Args: diff --git a/src/otx/core/exporter/base.py b/src/otx/core/exporter/base.py index 410f4a898d0..dcefc5fd902 100644 --- a/src/otx/core/exporter/base.py +++ b/src/otx/core/exporter/base.py @@ -206,7 +206,6 @@ def to_exportable_code( arch.write(str(path_to_model), Path("model") / "model.xml") arch.write(path_to_model.with_suffix(".bin"), Path("model") / "model.bin") - arch.writestr( str(Path("model") / "config.json"), json.dumps(parameters, ensure_ascii=False, indent=4), diff --git a/src/otx/core/exporter/exportable_code/demo/demo_package/__init__.py b/src/otx/core/exporter/exportable_code/demo/demo_package/__init__.py index 3afe9c9f203..604e5969747 100644 --- a/src/otx/core/exporter/exportable_code/demo/demo_package/__init__.py +++ b/src/otx/core/exporter/exportable_code/demo/demo_package/__init__.py @@ -6,10 +6,22 @@ from .executors import AsyncExecutor, SyncExecutor from .model_wrapper import ModelWrapper from .utils import create_visualizer +from .visualizers import ( + BaseVisualizer, + ClassificationVisualizer, + InstanceSegmentationVisualizer, + ObjectDetectionVisualizer, + SemanticSegmentationVisualizer, +) __all__ = [ "SyncExecutor", "AsyncExecutor", "create_visualizer", "ModelWrapper", + "BaseVisualizer", + "ClassificationVisualizer", + "SemanticSegmentationVisualizer", + "InstanceSegmentationVisualizer", + "ObjectDetectionVisualizer", ] diff --git a/src/otx/core/exporter/exportable_code/demo/demo_package/executors/asynchronous.py b/src/otx/core/exporter/exportable_code/demo/demo_package/executors/asynchronous.py index fc0cd328131..4549922ef5e 100644 --- a/src/otx/core/exporter/exportable_code/demo/demo_package/executors/asynchronous.py +++ b/src/otx/core/exporter/exportable_code/demo/demo_package/executors/asynchronous.py @@ -14,6 +14,7 @@ import numpy as np from demo_package.model_wrapper import ModelWrapper + from demo_package.streamer import get_streamer from demo_package.visualizers import BaseVisualizer, dump_frames diff --git a/src/otx/core/exporter/exportable_code/demo/demo_package/executors/synchronous.py b/src/otx/core/exporter/exportable_code/demo/demo_package/executors/synchronous.py index ea280841aad..06fd9035f8d 100644 --- a/src/otx/core/exporter/exportable_code/demo/demo_package/executors/synchronous.py +++ b/src/otx/core/exporter/exportable_code/demo/demo_package/executors/synchronous.py @@ -12,7 +12,7 @@ from demo_package.model_wrapper import ModelWrapper from demo_package.visualizers import BaseVisualizer -from demo_package.streamer import get_streamer +from demo_package.streamer.streamer import get_streamer from demo_package.visualizers import dump_frames diff --git a/src/otx/core/exporter/exportable_code/demo/demo_package/visualizers/vis_utils.py b/src/otx/core/exporter/exportable_code/demo/demo_package/visualizers/vis_utils.py index 0edea1cbb94..b8d9662f541 100644 --- a/src/otx/core/exporter/exportable_code/demo/demo_package/visualizers/vis_utils.py +++ b/src/otx/core/exporter/exportable_code/demo/demo_package/visualizers/vis_utils.py @@ -95,7 +95,7 @@ def __init__(self, num_classes: int, rng: random.Random | None = None) -> None: Returns: None """ - if num_classes == 0: + if num_classes <= 0: msg = "ColorPalette accepts only the positive number of colors" raise ValueError(msg) if rng is None: diff --git a/src/otx/core/exporter/exportable_code/demo/demo_package/visualizers/visualizer.py b/src/otx/core/exporter/exportable_code/demo/demo_package/visualizers/visualizer.py index 764eb23cf1d..dc9ee8f90ba 100644 --- a/src/otx/core/exporter/exportable_code/demo/demo_package/visualizers/visualizer.py +++ b/src/otx/core/exporter/exportable_code/demo/demo_package/visualizers/visualizer.py @@ -126,6 +126,10 @@ def draw( Output image with annotations. """ predictions = predictions.top_labels + if not any(predictions): + log.warning("There are no predictions.") + return frame + class_label = predictions[0][1] font_scale = 0.7 label_height = cv2.getTextSize(class_label, cv2.FONT_HERSHEY_COMPLEX, font_scale, 2)[0][1] diff --git a/src/otx/core/types/device.py b/src/otx/core/types/device.py index 0d11e0393f7..bd87a5721df 100644 --- a/src/otx/core/types/device.py +++ b/src/otx/core/types/device.py @@ -10,7 +10,7 @@ class DeviceType(str, Enum): """OTX Device type definition.""" - # ("cpu", "gpu", "tpu", "ipu", "hpu", "mps", "auto") + # ("cpu", "gpu", "tpu", "ipu", "hpu", "mps", "xpu", "auto") auto = "auto" gpu = "gpu" @@ -19,3 +19,4 @@ class DeviceType(str, Enum): ipu = "ipu" hpu = "hpu" mps = "mps" + xpu = "xpu" diff --git a/src/otx/engine/engine.py b/src/otx/engine/engine.py index 4a90faa3ea2..cb76f982c29 100644 --- a/src/otx/engine/engine.py +++ b/src/otx/engine/engine.py @@ -16,6 +16,7 @@ import torch from lightning import Trainer, seed_everything +from otx.algo.plugins import MixedPrecisionXPUPlugin from otx.core.config.device import DeviceConfig from otx.core.config.explain import ExplainConfig from otx.core.config.hpo import HpoConfig @@ -27,6 +28,7 @@ from otx.core.types.precision import OTXPrecisionType from otx.core.types.task import OTXTaskType from otx.core.utils.cache import TrainerArgumentsCache +from otx.utils.utils import is_xpu_available from .hpo import execute_hpo, update_hyper_parameter from .utils.auto_configurator import DEFAULT_CONFIG_PER_TASK, AutoConfigurator @@ -879,6 +881,8 @@ def device(self) -> DeviceConfig: @device.setter def device(self, device: DeviceType) -> None: + if is_xpu_available() and device == DeviceType.auto: + device = DeviceType.xpu self._device = DeviceConfig(accelerator=device) self._cache.update(accelerator=self._device.accelerator, devices=self._device.devices) self._cache.is_trainer_args_identical = False @@ -901,6 +905,14 @@ def _build_trainer(self, **kwargs) -> None: """Instantiate the trainer based on the model parameters.""" if self._cache.requires_update(**kwargs) or self._trainer is None: self._cache.update(**kwargs) + # set up xpu device + if self._device.accelerator == DeviceType.xpu: + self._cache.update(strategy="xpu_single") + # add plugin for Automatic Mixed Precision on XPU + if self._cache.args.get("precision", 32) == 16: + self._cache.update(plugins=[MixedPrecisionXPUPlugin()]) + self._cache.args["precision"] = None + kwargs = self._cache.args self._trainer = Trainer(**kwargs) self._cache.is_trainer_args_identical = True diff --git a/src/otx/utils/utils.py b/src/otx/utils/utils.py index 89cf03a2c79..797274aa634 100644 --- a/src/otx/utils/utils.py +++ b/src/otx/utils/utils.py @@ -8,10 +8,19 @@ from decimal import Decimal from typing import TYPE_CHECKING, Any +import torch + if TYPE_CHECKING: from pathlib import Path +XPU_AVAILABLE = None +try: + import intel_extension_for_pytorch # noqa: F401 +except ImportError: + XPU_AVAILABLE = False + + def get_using_dot_delimited_key(key: str, target: Any) -> Any: # noqa: ANN401 """Get values of attribute in target object using dot delimited key. @@ -114,3 +123,11 @@ def remove_matched_files(directory: Path, pattern: str, file_to_leave: Path | No for weight in directory.rglob(pattern): if weight != file_to_leave: weight.unlink() + + +def is_xpu_available() -> bool: + """Checks if XPU device is available.""" + global XPU_AVAILABLE # noqa: PLW0603 + if XPU_AVAILABLE is None: + XPU_AVAILABLE = hasattr(torch, "xpu") and torch.xpu.is_available() + return XPU_AVAILABLE diff --git a/tests/unit/algo/accelerators/__init__.py b/tests/unit/algo/accelerators/__init__.py new file mode 100644 index 00000000000..9996ffc6523 --- /dev/null +++ b/tests/unit/algo/accelerators/__init__.py @@ -0,0 +1,4 @@ +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# +"""Unit tests of accelerators of OTX algo.""" diff --git a/tests/unit/algo/accelerators/test_xpu.py b/tests/unit/algo/accelerators/test_xpu.py new file mode 100644 index 00000000000..793bbe18331 --- /dev/null +++ b/tests/unit/algo/accelerators/test_xpu.py @@ -0,0 +1,62 @@ +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# +"""Test for otx.algo.accelerators.xpu""" + + +import pytest +import torch +from otx.algo.accelerators import XPUAccelerator +from otx.utils.utils import is_xpu_available + + +class TestXPUAccelerator: + @pytest.fixture() + def accelerator(self, mocker): + mock_torch = mocker.patch("otx.algo.accelerators.xpu.torch") + mocker.patch.object(XPUAccelerator, "patch_packages_xpu") + mocker.patch.object(XPUAccelerator, "teardown") + return XPUAccelerator(), mock_torch + + def test_setup_device(self, accelerator): + accelerator, mock_torch = accelerator + device = torch.device("xpu") + accelerator.setup_device(device) + assert mock_torch.xpu.set_device.called + + def test_parse_devices(self, accelerator): + accelerator, _ = accelerator + devices = [1, 2, 3] + parsed_devices = accelerator.parse_devices(devices) + assert isinstance(parsed_devices, list) + assert parsed_devices == devices + + def test_get_parallel_devices(self, accelerator, mocker): + accelerator, _ = accelerator + devices = [1, 2, 3] + parallel_devices = accelerator.get_parallel_devices(devices) + assert isinstance(parallel_devices, list) + for device in parallel_devices: + assert isinstance(device, mocker.MagicMock) + + def test_auto_device_count(self, accelerator, mocker): + accelerator, mock_torch = accelerator + count = accelerator.auto_device_count() + assert isinstance(count, mocker.MagicMock) + assert mock_torch.xpu.device_count.called + + def test_is_available(self, accelerator): + accelerator, _ = accelerator + available = accelerator.is_available() + assert isinstance(available, bool) + assert available == is_xpu_available() + + def test_get_device_stats(self, accelerator): + accelerator, _ = accelerator + device = torch.device("xpu") + stats = accelerator.get_device_stats(device) + assert isinstance(stats, dict) + + def test_teardown(self, accelerator): + accelerator, _ = accelerator + accelerator.teardown() diff --git a/tests/unit/algo/detection/utils/__init__.py b/tests/unit/algo/detection/utils/__init__.py new file mode 100644 index 00000000000..a3c91b9065c --- /dev/null +++ b/tests/unit/algo/detection/utils/__init__.py @@ -0,0 +1,4 @@ +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# +"""Test of utils for OTX Detection task.""" diff --git a/tests/unit/algo/detection/utils/test_mmcv_patched_ops.py b/tests/unit/algo/detection/utils/test_mmcv_patched_ops.py new file mode 100644 index 00000000000..09daa1b2cab --- /dev/null +++ b/tests/unit/algo/detection/utils/test_mmcv_patched_ops.py @@ -0,0 +1,139 @@ +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# +"""Test of mmcv_patched_ops.""" + +import pytest +import torch +from mmcv.ops import nms +from otx.algo.detection.utils.mmcv_patched_ops import monkey_patched_nms + + +class TestMonkeyPatchedNMS: + @pytest.fixture() + def setup(self): + self.ctx = None + self.bboxes = torch.tensor( + [[0.324, 0.422, 0.469, 0.123], [0.324, 0.422, 0.469, 0.123], [0.314, 0.423, 0.469, 0.123]], + ) + self.scores = torch.tensor([0.9, 0.2, 0.3]) + self.iou_threshold = 0.5 + self.offset = 0 + self.score_threshold = 0 + self.max_num = 0 + + def test_case1(self, setup): + # Testing when is_filtering_by_score is False + result = monkey_patched_nms( + self.ctx, + self.bboxes, + self.scores, + self.iou_threshold, + self.offset, + self.score_threshold, + self.max_num, + ) + assert torch.equal(result, torch.tensor([0, 2, 1])) + + def test_case2(self, setup): + # Testing when is_filtering_by_score is True + self.score_threshold = 0.8 + result = monkey_patched_nms( + self.ctx, + self.bboxes, + self.scores, + self.iou_threshold, + self.offset, + self.score_threshold, + self.max_num, + ) + assert torch.equal(result, torch.tensor([0])) + + def test_case3(self, setup): + # Testing when bboxes and scores have torch.bfloat16 dtype + self.bboxes = torch.tensor( + [[0.324, 0.422, 0.469, 0.123], [0.324, 0.422, 0.469, 0.123], [0.314, 0.423, 0.469, 0.123]], + dtype=torch.bfloat16, + ) + self.scores = torch.tensor([0.9, 0.2, 0.3], dtype=torch.bfloat16) + result1 = monkey_patched_nms( + self.ctx, + self.bboxes, + self.scores, + self.iou_threshold, + self.offset, + self.score_threshold, + self.max_num, + ) + assert torch.equal(result1, torch.tensor([0, 2, 1])) + + def test_case4(self, setup): + # Testing when offset is not 0 + self.offset = 1 + result = monkey_patched_nms( + self.ctx, + self.bboxes, + self.scores, + self.iou_threshold, + self.offset, + self.score_threshold, + self.max_num, + ) + assert torch.equal(result, torch.tensor([0])) + + def test_case5(self, setup): + # Testing when max_num is greater than 0 + self.max_num = 1 + result = monkey_patched_nms( + self.ctx, + self.bboxes, + self.scores, + self.iou_threshold, + self.offset, + self.score_threshold, + self.max_num, + ) + assert torch.equal(result, torch.tensor([0])) + + def test_case6(self, setup): + # Testing that monkey_patched_nms equals mmcv nms + self.score_threshold = 0.7 + result1 = monkey_patched_nms( + self.ctx, + self.bboxes, + self.scores, + self.iou_threshold, + self.offset, + self.score_threshold, + self.max_num, + ) + result2 = nms(self.bboxes, self.scores, self.iou_threshold, score_threshold=self.score_threshold) + assert torch.equal(result1, result2[1]) + # test random bboxes and scores + bboxes = torch.rand((100, 4)) + scores = torch.rand(100) + result1 = monkey_patched_nms( + self.ctx, + bboxes, + scores, + self.iou_threshold, + self.offset, + self.score_threshold, + self.max_num, + ) + result2 = nms(bboxes, scores, self.iou_threshold, score_threshold=self.score_threshold) + assert torch.equal(result1, result2[1]) + # no score threshold + self.iou_threshold = 0.7 + self.score_threshold = 0.0 + result1 = monkey_patched_nms( + self.ctx, + bboxes, + scores, + self.iou_threshold, + self.offset, + self.score_threshold, + self.max_num, + ) + result2 = nms(bboxes, scores, self.iou_threshold) + assert torch.equal(result1, result2[1]) diff --git a/tests/unit/algo/plugins/__init__.py b/tests/unit/algo/plugins/__init__.py new file mode 100644 index 00000000000..be7fe475146 --- /dev/null +++ b/tests/unit/algo/plugins/__init__.py @@ -0,0 +1,4 @@ +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# +"""Unit tests of plugins of OTX algo.""" diff --git a/tests/unit/algo/plugins/test_plugins.py b/tests/unit/algo/plugins/test_plugins.py new file mode 100644 index 00000000000..a84f4ec18d6 --- /dev/null +++ b/tests/unit/algo/plugins/test_plugins.py @@ -0,0 +1,56 @@ +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# +"""Test for otx.algo.plugins.xpu_precision""" + + +import pytest +import torch +from otx.algo.plugins.xpu_precision import MixedPrecisionXPUPlugin +from torch.optim import Optimizer + + +class TestMixedPrecisionXPUPlugin: + @pytest.fixture() + def plugin(self): + return MixedPrecisionXPUPlugin() + + def test_init(self, plugin): + assert plugin.scaler is None + + def test_pre_backward(self, plugin, mocker): + tensor = torch.zeros(1) + module = mocker.MagicMock() + output = plugin.pre_backward(tensor, module) + assert output == tensor + + def test_optimizer_step_no_scaler(self, plugin, mocker): + optimizer = mocker.MagicMock(Optimizer) + model = mocker.MagicMock() + closure = mocker.MagicMock() + kwargs = {} + mock_optimizer_step = mocker.patch( + "otx.algo.plugins.xpu_precision.Precision.optimizer_step", + ) + out = plugin.optimizer_step(optimizer, model, closure, **kwargs) + assert isinstance(out, mocker.MagicMock) + mock_optimizer_step.assert_called_once() + + def test_optimizer_step_with_scaler(self, plugin, mocker): + optimizer = mocker.MagicMock(Optimizer) + model = mocker.MagicMock() + closure = mocker.MagicMock() + plugin.scaler = mocker.MagicMock() + kwargs = {} + out = plugin.optimizer_step(optimizer, model, closure, **kwargs) + assert isinstance(out, mocker.MagicMock) + + def test_clip_gradients(self, plugin, mocker): + optimizer = mocker.MagicMock(Optimizer) + clip_val = 0.1 + gradient_clip_algorithm = "norm" + mock_clip_gradients = mocker.patch( + "otx.algo.plugins.xpu_precision.Precision.clip_gradients", + ) + plugin.clip_gradients(optimizer, clip_val, gradient_clip_algorithm) + mock_clip_gradients.assert_called_once() diff --git a/tests/unit/algo/strategies/__init__.py b/tests/unit/algo/strategies/__init__.py new file mode 100644 index 00000000000..8830174eb83 --- /dev/null +++ b/tests/unit/algo/strategies/__init__.py @@ -0,0 +1,4 @@ +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# +"""Unit tests of strategies of OTX algo.""" diff --git a/tests/unit/algo/strategies/test_strategies.py b/tests/unit/algo/strategies/test_strategies.py new file mode 100644 index 00000000000..0ef457351ff --- /dev/null +++ b/tests/unit/algo/strategies/test_strategies.py @@ -0,0 +1,54 @@ +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# +"""Tests the XPU strategy.""" + + +import pytest +import pytorch_lightning as pl +import torch +from lightning.pytorch.utilities.exceptions import MisconfigurationException +from otx.algo.strategies.xpu_single import SingleXPUStrategy + + +class TestSingleXPUStrategy: + def test_init(self, mocker): + with pytest.raises(MisconfigurationException): + strategy = SingleXPUStrategy(device="xpu:0") + mocked_is_xpu_available = mocker.patch( + "otx.algo.strategies.xpu_single.is_xpu_available", + return_value=True, + ) + strategy = SingleXPUStrategy(device="xpu:0") + assert mocked_is_xpu_available.call_count == 1 + assert strategy._root_device.type == "xpu" + assert strategy.accelerator is None + + @pytest.fixture() + def strategy(self, mocker): + mocker.patch( + "otx.algo.strategies.xpu_single.is_xpu_available", + return_value=True, + ) + return SingleXPUStrategy(device="xpu:0", accelerator="xpu") + + def test_is_distributed(self, strategy): + assert not strategy.is_distributed + + def test_setup_optimizers(self, strategy, mocker): + from otx.algo.strategies.xpu_single import SingleDeviceStrategy + + mocker.patch("otx.algo.strategies.xpu_single.torch") + mocker.patch( + "otx.algo.strategies.xpu_single.torch.xpu.optimize", + return_value=(mocker.MagicMock(), mocker.MagicMock()), + ) + mocker.patch.object(SingleDeviceStrategy, "setup_optimizers") + trainer = pl.Trainer() + trainer.task = "CLASSIFICATION" + # Create mock optimizers and models for testing + model = torch.nn.Linear(10, 2) + strategy._optimizers = [torch.optim.Adam(model.parameters(), lr=0.001)] + strategy._model = model + strategy.setup_optimizers(trainer) + assert len(strategy.optimizers) == 1 diff --git a/tests/unit/core/exporter/exportable_code/demo/demo_package/visualizers/test_vis_utils.py b/tests/unit/core/exporter/exportable_code/demo/demo_package/visualizers/test_vis_utils.py new file mode 100644 index 00000000000..04f0e6580c1 --- /dev/null +++ b/tests/unit/core/exporter/exportable_code/demo/demo_package/visualizers/test_vis_utils.py @@ -0,0 +1,127 @@ +from pathlib import Path +from unittest.mock import MagicMock, patch + +import cv2 +import numpy as np +import pytest +from numpy.random import PCG64, Generator + + +@pytest.fixture(scope="module", autouse=True) +def fxt_import_module(): + global ColorPalette, get_actmap, dump_frames # noqa: PLW0603 + from otx.core.exporter.exportable_code.demo.demo_package.visualizers.vis_utils import ( + ColorPalette as _ColorPalette, + ) + from otx.core.exporter.exportable_code.demo.demo_package.visualizers.vis_utils import ( + dump_frames as _dump_frames, + ) + from otx.core.exporter.exportable_code.demo.demo_package.visualizers.vis_utils import ( + get_actmap as _get_actmap, + ) + + ColorPalette = _ColorPalette + get_actmap = _get_actmap + dump_frames = _dump_frames + + +def test_activation_map_shape(): + generator = Generator(PCG64()) + saliency_map = (generator.random((100, 100)) * 255).astype(np.uint8) + output_res = (50, 50) + result = get_actmap(saliency_map, output_res) + assert result.shape == (50, 50, 3) + + +def test_no_saved_frames(): + output = "output" + input_path = "input" + capture = MagicMock() + saved_frames = [] + dump_frames(saved_frames, output, input_path, capture) + assert not Path(output).exists() + + +@patch("otx.core.exporter.exportable_code.demo.demo_package.visualizers.vis_utils.cv2.VideoWriter_fourcc") +@patch("otx.core.exporter.exportable_code.demo.demo_package.visualizers.vis_utils.cv2.VideoWriter") +@patch("otx.core.exporter.exportable_code.demo.demo_package.visualizers.vis_utils.get_input_names_list") +def test_video_input(mock_get_input_names_list, mock_video_writer, mock_video_writer_fourcc, tmp_path): + output = str(tmp_path / "output") + input_path = "input" + capture = MagicMock(spec=cv2.VideoCapture) + capture.get_type = lambda: "VIDEO" + capture.fps = lambda: 30 + saved_frames = [MagicMock(shape=(100, 100, 3))] + filenames = ["video.mp4"] + mock_get_input_names_list.return_value = filenames + mock_video_writer_fourcc.return_value = "mp4v" + dump_frames(saved_frames, output, input_path, capture) + mock_video_writer_fourcc.assert_called_once_with(*"mp4v") + mock_video_writer.assert_called_once() + + +@patch("otx.core.exporter.exportable_code.demo.demo_package.visualizers.vis_utils.cv2.imwrite") +@patch("otx.core.exporter.exportable_code.demo.demo_package.visualizers.vis_utils.get_input_names_list") +@patch("otx.core.exporter.exportable_code.demo.demo_package.visualizers.vis_utils.cv2.cvtColor") +def test_image_input(mock_imwrite, mock_get_input_names_list, mock_cvtcolor, tmp_path): + output = str(tmp_path / "output") + input_path = "input" + capture = MagicMock(spec=cv2.VideoCapture) + capture.get_type = lambda: "IMAGE" + saved_frames = [MagicMock(), MagicMock()] + filenames = ["image1.jpeg", "image2.jpeg"] + mock_get_input_names_list.return_value = filenames + dump_frames(saved_frames, output, input_path, capture) + assert mock_cvtcolor.call_count == 2 + assert mock_imwrite.call_count == 2 + + +class TestColorPalette: + def test_colorpalette_init_with_zero_classes(self): + expected_msg = "ColorPalette accepts only the positive number of colors" + with pytest.raises(ValueError, match=expected_msg): + ColorPalette(num_classes=0) + with pytest.raises(ValueError, match=expected_msg): + ColorPalette(num_classes=-5) + + def test_colorpalette_length(self): + num_classes = 5 + palette = ColorPalette(num_classes) + assert len(palette) == num_classes + + def test_colorpalette_getitem(self): + num_classes = 3 + palette = ColorPalette(num_classes) + color = palette[1] # assuming 0-based indexing + assert isinstance(color, tuple) + assert len(color) == 3 + + def test_colorpalette_getitem_out_of_range(self): + num_classes = 3 + palette = ColorPalette(num_classes) + color = palette[num_classes + 2] # out-of-range index + assert color == palette[2] # because it should wrap around + + def test_colorpalette_to_numpy_array(self): + num_classes = 2 + palette = ColorPalette(num_classes) + np_array = palette.to_numpy_array() + assert isinstance(np_array, np.ndarray) + assert np_array.shape == (num_classes, 3) + + def test_colorpalette_hsv2rgb_known_values(self): + h, s, v = 0.5, 1, 1 # Cyan in HSV + expected_rgb = (0, 255, 255) # Cyan in RGB + assert ColorPalette.hsv2rgb(h, s, v) == expected_rgb + + def test_dist_same_color(self): + # Colors that are the same should have a distance of 0 + color = (0.5, 0.5, 0.5) + assert ColorPalette._dist(color, color) == 0 + + def test_dist_different_colors(self): + # Test distance between two different colors + color1 = (0.1, 0.2, 0.3) + color2 = (0.4, 0.5, 0.6) + expected_distance = 0.54 + assert ColorPalette._dist(color1, color2) == expected_distance diff --git a/tests/unit/core/exporter/exportable_code/demo/demo_package/visualizers/test_visualizers.py b/tests/unit/core/exporter/exportable_code/demo/demo_package/visualizers/test_visualizers.py new file mode 100644 index 00000000000..331aa1cc3e1 --- /dev/null +++ b/tests/unit/core/exporter/exportable_code/demo/demo_package/visualizers/test_visualizers.py @@ -0,0 +1,272 @@ +from unittest.mock import Mock, patch + +import numpy as np +import pytest +from numpy.random import PCG64, Generator +from openvino.model_api.models.utils import ( + ClassificationResult, + Detection, + DetectionResult, + ImageResultWithSoftPrediction, + InstanceSegmentationResult, + SegmentedObject, +) + + +@pytest.fixture(scope="module", autouse=True) +def fxt_import_module(): + global BaseVisualizer, ClassificationVisualizer, InstanceSegmentationVisualizer, ObjectDetectionVisualizer, SemanticSegmentationVisualizer # noqa: PLW0603 + from otx.core.exporter.exportable_code.demo.demo_package import ( + BaseVisualizer as _BaseVisualizer, + ) + from otx.core.exporter.exportable_code.demo.demo_package import ( + ClassificationVisualizer as _ClassificationVisualizer, + ) + from otx.core.exporter.exportable_code.demo.demo_package import ( + InstanceSegmentationVisualizer as _InstanceSegmentationVisualizer, + ) + from otx.core.exporter.exportable_code.demo.demo_package import ( + ObjectDetectionVisualizer as _ObjectDetectionVisualizer, + ) + from otx.core.exporter.exportable_code.demo.demo_package import ( + SemanticSegmentationVisualizer as _SemanticSegmentationVisualizer, + ) + + BaseVisualizer = _BaseVisualizer + ClassificationVisualizer = _ClassificationVisualizer + InstanceSegmentationVisualizer = _InstanceSegmentationVisualizer + ObjectDetectionVisualizer = _ObjectDetectionVisualizer + SemanticSegmentationVisualizer = _SemanticSegmentationVisualizer + + +class TestBaseVisualizer: + def test_init(self): + visualizer = BaseVisualizer(window_name="TestWindow", no_show=True, delay=10, output="test_output") + assert visualizer.window_name == "TestWindow" + assert visualizer.no_show is True + assert visualizer.delay == 10 + assert visualizer.output == "test_output" + + # Test show method without displaying the window + @patch("cv2.imshow") + def test_show_no_display(self, mock_imshow): + visualizer = BaseVisualizer(no_show=True) + test_image = np.zeros((100, 100, 3), dtype=np.uint8) + visualizer.show(test_image) + mock_imshow.assert_not_called() + + # Test show method with displaying the window + @patch("cv2.imshow") + def test_show_display(self, mock_imshow): + visualizer = BaseVisualizer(no_show=False) + test_image = np.zeros((100, 100, 3), dtype=np.uint8) + visualizer.show(test_image) + mock_imshow.assert_called_once_with(visualizer.window_name, test_image) + + # Test is_quit method + @patch("cv2.waitKey", return_value=ord("q")) + def test_is_quit(self, mock_waitkey): + visualizer = BaseVisualizer(no_show=False) + assert visualizer.is_quit() is True + + # Test video_delay method + @patch("time.sleep") + def test_video_delay(self, mock_sleep): + streamer = Mock() + streamer.get_type.return_value = "VIDEO" + streamer.fps.return_value = 30 + visualizer = BaseVisualizer(no_show=False) + visualizer.video_delay(0.02, streamer) + mock_sleep.assert_called_once_with(1 / 30 - 0.02) + + +class TestClassificationVisualizer: + @pytest.fixture() + def visualizer(self): + return ClassificationVisualizer(window_name="TestWindow", no_show=True, delay=10, output="test_output") + + @pytest.fixture() + def frame(self): + return np.zeros((100, 100, 3), dtype=np.uint8) + + @pytest.fixture() + def predictions(self): + return ClassificationResult( + top_labels=[(0, "cat", 0.9)], + saliency_map=None, + feature_vector=None, + raw_scores=[0.9], + ) + + def test_draw_one_prediction(self, frame, predictions, visualizer): + # test one prediction + copied_frame = frame.copy() + output = visualizer.draw(frame, predictions) + assert output.shape == (100, 100, 3) + assert np.any(output != copied_frame) + + def test_draw_multiple_predictions(self, frame, predictions, visualizer): + # test multiple predictions + copied_frame = frame.copy() + predictions.top_labels.extend([(1, "dog", 0.8), (2, "bird", 0.7)]) + output = visualizer.draw(frame, predictions) + assert output.shape == (100, 100, 3) + assert np.any(output != copied_frame) + + def test_label_overflow(self, frame, predictions, visualizer): + # test multiple predictions + copied_frame = frame.copy() + predictions.top_labels.extend([(1, "dog", 0.8), (2, "bird", 0.7), (3, "cat", 0.6)]) + output = visualizer.draw(frame, predictions) + assert output.shape == (100, 100, 3) + assert np.any(output != copied_frame) + + def test_draw_no_predictions(self, frame, visualizer): + # test no predictions + copied_frame = frame.copy() + predictions = ClassificationResult(top_labels=[()], saliency_map=None, feature_vector=None, raw_scores=[]) + output = visualizer.draw(frame, predictions) + assert output.shape == (100, 100, 3) + assert np.equal(output, copied_frame).all() + + +class TestDetectionVisualizer: + @pytest.fixture() + def visualizer(self): + return ObjectDetectionVisualizer( + labels=["Pedestrian", "Car"], + window_name="TestWindow", + no_show=True, + delay=10, + output="test_output", + ) + + def test_draw_no_predictions(self, visualizer): + frame = np.zeros((100, 100, 3), dtype=np.uint8) + predictions = DetectionResult([], saliency_map=None, feature_vector=None) + output_frame = visualizer.draw(frame, predictions) + assert np.array_equal(frame, output_frame) + + def test_draw_with_predictions(self, visualizer): + frame = np.zeros((100, 100, 3), dtype=np.uint8) + predictions = DetectionResult( + [Detection(10, 40, 30, 80, 0.7, 2, "Car")], + saliency_map=None, + feature_vector=None, + ) + copied_frame = frame.copy() + output_frame = visualizer.draw(frame, predictions) + assert np.any(output_frame != copied_frame) + + +class TestInstanceSegmentationVisualizer: + @pytest.fixture() + def rand_generator(self): + return Generator(PCG64()) + + @pytest.fixture() + def visualizer(self): + return InstanceSegmentationVisualizer( + labels=["person", "car"], + window_name="TestWindow", + no_show=True, + delay=10, + output="test_output", + ) + + def test_draw_multiple_objects(self, visualizer, rand_generator): + # Create a frame + frame = np.zeros((100, 100, 3), dtype=np.uint8) + copied_frame = frame.copy() + + # Create instance segmentation results with multiple objects + predictions = InstanceSegmentationResult( + segmentedObjects=[ + SegmentedObject( + xmin=10, + ymin=10, + xmax=30, + ymax=30, + score=0.9, + id=0, + mask=rand_generator.integers(2, size=(100, 100), dtype=np.uint8), + str_label="person", + ), + SegmentedObject( + xmin=40, + ymin=40, + xmax=60, + ymax=60, + score=0.8, + id=1, + mask=rand_generator.integers(2, size=(100, 100), dtype=np.uint8), + str_label="car", + ), + ], + saliency_map=None, + feature_vector=None, + ) + + drawn_frame = visualizer.draw(frame, predictions) + assert np.any(drawn_frame != copied_frame) + + # Assertion checks for the drawn frame + + def test_draw_no_objects(self, visualizer): + # Create a frame + frame = np.zeros((100, 100, 3), dtype=np.uint8) + copied_frame = frame.copy() + + # Create instance segmentation results with no objects + predictions = InstanceSegmentationResult(segmentedObjects=[], saliency_map=None, feature_vector=None) + + drawn_frame = visualizer.draw(frame, predictions) + assert np.array_equal(drawn_frame, copied_frame) + + +class TestSemanticSegmentationVisualizer: + @pytest.fixture() + def labels(self): + return ["background", "object1", "object2"] + + @pytest.fixture() + def visualizer(self, labels): + return SemanticSegmentationVisualizer( + labels=labels, + window_name="TestWindow", + no_show=True, + delay=10, + output="test_output", + ) + + @pytest.fixture() + def rand_generator(self): + return Generator(PCG64()) + + def test_initialization(self, visualizer): + assert isinstance(visualizer.color_palette, np.ndarray) + assert visualizer.color_map.shape == (256, 1, 3) + assert visualizer.color_map.dtype == np.uint8 + + def test_create_color_map(self, visualizer): + color_map = visualizer._create_color_map() + assert color_map.shape == (256, 1, 3) + assert color_map.dtype == np.uint8 + + def test_apply_color_map(self, visualizer, labels, rand_generator): + input_2d_mask = rand_generator.integers(0, len(labels), size=(10, 10)) + colored_mask = visualizer._apply_color_map(input_2d_mask) + assert colored_mask.shape == (10, 10, 3) + + def test_draw(self, visualizer, rand_generator): + frame = rand_generator.integers(0, 255, size=(10, 10, 3), dtype=np.uint8) + copied_frame = frame.copy() + masks = ImageResultWithSoftPrediction( + resultImage=rand_generator.integers(0, 255, size=(10, 10), dtype=np.uint8), + soft_prediction=rand_generator.random((10, 10)), + saliency_map=None, + feature_vector=None, + ) + output_image = visualizer.draw(frame, masks) + assert output_image.shape == frame.shape + assert np.any(output_image != copied_frame) diff --git a/tests/unit/core/exporter/test_base.py b/tests/unit/core/exporter/test_base.py new file mode 100644 index 00000000000..e40922362ad --- /dev/null +++ b/tests/unit/core/exporter/test_base.py @@ -0,0 +1,97 @@ +from unittest.mock import MagicMock, patch + +import pytest +from onnx import ModelProto +from onnxconverter_common import float16 +from otx.core.exporter.base import OTXExportFormatType, OTXModelExporter, OTXPrecisionType, ZipFile +from otx.core.types.export import TaskLevelExportParameters + + +class MockModelExporter(OTXModelExporter): + def to_openvino(self, model, output_dir, base_model_name, precision): + return output_dir / f"{base_model_name}.xml" + + def to_onnx(self, model, output_dir, base_model_name, precision): + return output_dir / f"{base_model_name}.onnx" + + +@pytest.fixture() +def mock_model(): + return MagicMock() + + +@pytest.fixture() +def exporter(mocker): + ZipFile.write = MagicMock() + mocker.patch("otx.core.exporter.base.json") + return MockModelExporter( + task_level_export_parameters=MagicMock(TaskLevelExportParameters), + input_size=(224, 224), + mean=(0.485, 0.456, 0.406), + std=(0.229, 0.224, 0.225), + ) + + +class TestOTXModelExporter: + def test_to_openvino(self, mock_model, exporter, tmp_path): + output_dir = tmp_path + base_model_name = "test_model" + precision = OTXPrecisionType.FP32 + result = exporter.export(mock_model, output_dir, base_model_name, OTXExportFormatType.OPENVINO, precision) + assert result == output_dir / f"{base_model_name}.xml" + + def test_to_onnx(self, mock_model, exporter, tmp_path): + output_dir = tmp_path + base_model_name = "test_model" + precision = OTXPrecisionType.FP32 + result = exporter.export(mock_model, output_dir, base_model_name, OTXExportFormatType.ONNX, precision) + assert result == output_dir / f"{base_model_name}.onnx" + + def test_export_unsupported_format_raises(self, exporter, mock_model, tmp_path): + export_format = "unsupported_format" + with pytest.raises(ValueError, match=f"Unsupported export format: {export_format}"): + exporter.export(mock_model, tmp_path, export_format=export_format) + + def test_to_exportable_code(self, mock_model, exporter, tmp_path): + from otx.core.exporter.base import ZipFile + + ZipFile.writestr = MagicMock() + + base_model_name = "test_model" + output_dir = tmp_path / "exportable_code" + precision = OTXPrecisionType.FP32 + + with patch("builtins.open", new_callable=MagicMock): + exporter.to_openvino = MagicMock() + result = exporter.to_exportable_code(mock_model, output_dir, base_model_name, precision) + + assert result == output_dir / "exportable_code.zip" + + def test_postprocess_openvino_model(self, mock_model, exporter): + # test output names do not match exporter parameters + exporter.output_names = ["output1"] + with pytest.raises(RuntimeError): + exporter._postprocess_openvino_model(mock_model) + # test output names match exporter parameters + exporter.output_names = ["output1", "output2"] + mock_model.outputs = [] + for output_name in exporter.output_names: + output = MagicMock() + output.get_names.return_value = output_name + mock_model.outputs.append(output) + processed_model = exporter._postprocess_openvino_model(mock_model) + # Verify the processed model is returned and the names are set correctly + assert processed_model is mock_model + for output, name in zip(processed_model.outputs, exporter.output_names): + output.tensor.set_names.assert_called_once_with({name}) + + def test_embed_metadata_true_precision_fp16(self, exporter): + onnx_model = ModelProto() + exporter._embed_onnx_metadata = MagicMock(return_value=onnx_model) + convert_float_to_float16_mock = MagicMock(return_value=onnx_model) + with pytest.MonkeyPatch.context() as m: + m.setattr(float16, "convert_float_to_float16", convert_float_to_float16_mock) + result = exporter._postprocess_onnx_model(onnx_model, embed_metadata=True, precision=OTXPrecisionType.FP16) + exporter._embed_onnx_metadata.assert_called_once() + convert_float_to_float16_mock.assert_called_once_with(onnx_model) + assert result is onnx_model diff --git a/tests/unit/core/exporter/test_native.py b/tests/unit/core/exporter/test_native.py new file mode 100644 index 00000000000..4c9f481e63c --- /dev/null +++ b/tests/unit/core/exporter/test_native.py @@ -0,0 +1,75 @@ +import onnx +import pytest +import torch +from otx.core.exporter.native import OTXNativeModelExporter +from otx.core.types.export import TaskLevelExportParameters +from otx.core.types.precision import OTXPrecisionType + + +class TestOTXNativeModelExporter: + @pytest.fixture() + def exporter(self, mocker): + # Create an instance of OTXNativeModelExporter with default params + return OTXNativeModelExporter( + task_level_export_parameters=mocker.MagicMock(TaskLevelExportParameters), + input_size=(3, 224, 224), + ) + + @pytest.fixture() + def dummy_model(self): + # Define a simple dummy torch model for testing + return torch.nn.Sequential( + torch.nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1), + torch.nn.ReLU(), + ) + + def test_to_openvino_export(self, exporter, dummy_model, tmp_path): + # Use tmp_path provided by pytest for temporary file creation + output_dir = tmp_path / "model_export" + output_dir.mkdir() + + # Call the to_openvino method + exported_path = exporter.to_openvino( + model=dummy_model, + output_dir=output_dir, + base_model_name="test_model", + precision=OTXPrecisionType.FP32, + ) + + # Check that the exported files exist + assert exported_path.exists() + assert (output_dir / "test_model.xml").exists() + assert (output_dir / "test_model.bin").exists() + + exporter.via_onnx = True + exported_path = exporter.to_openvino( + model=dummy_model, + output_dir=output_dir, + base_model_name="test_model", + precision=OTXPrecisionType.FP32, + ) + + assert exported_path.exists() + assert (output_dir / "test_model.xml").exists() + assert (output_dir / "test_model.bin").exists() + + def test_to_onnx_export(self, exporter, dummy_model, tmp_path): + # Use tmp_path provided by pytest for temporary file creation + output_dir = tmp_path / "onnx_export" + output_dir.mkdir() + + # Call the to_onnx method + exported_path = exporter.to_onnx( + model=dummy_model, + output_dir=output_dir, + base_model_name="test_onnx_model", + precision=OTXPrecisionType.FP32, + ) + + # Check that the exported ONNX file exists + assert exported_path.exists() + assert (output_dir / "test_onnx_model.onnx").exists() + + # Load the model to verify it's a valid ONNX file + onnx_model = onnx.load(str(exported_path)) + onnx.checker.check_model(onnx_model)