diff --git a/README.md b/README.md index dd73b3f..1e902f8 100644 --- a/README.md +++ b/README.md @@ -315,7 +315,7 @@ pipeline: evaluators: - task_name: instance_segmentation parser: - name: YOLOInstanceSegmentationParser + name: YOLOExtendedParser params: subtype: yolov8 n_classes: 80 @@ -395,7 +395,7 @@ pipeline: evaluators: - task_name: instance_segmentation parser: - name: YOLOInstanceSegmentationParser + name: YOLOExtendedParser params: subtype: yolov8 n_classes: 80 @@ -471,10 +471,8 @@ Subclass [`BaseParser`](luxonis_eval/parsers/base_parser.py) and implement the s The parser bridges the gap between model-specific tensor layouts and the standardized message types that downstream metrics expect. The built-in parsers produce the following output types: - [**ClassificationParser**](luxonis_eval/parsers/classification.py) -> [depthai_nodes.Classifications](https://github.com/luxonis/depthai-nodes/tree/main/depthai_nodes/message#classifications) -- [**YOLODetectionParser**](luxonis_eval/parsers/detection.py) -> [dai.ImgDetections](https://docs.luxonis.com/software-v3/depthai/api/cpp/#classdai_1_1ImgDetections) -- [**YOLOInstanceSegmentationParser**](luxonis_eval/parsers/instance_seg.py) -> [dai.ImgDetections](https://docs.luxonis.com/software-v3/depthai/api/cpp/#classdai_1_1ImgDetections) -- [**YOLOKeypointDetectionParser**](luxonis_eval/parsers/keypoint_detection.py) -> [dai.ImgDetections](https://docs.luxonis.com/software-v3/depthai/api/cpp/#classdai_1_1ImgDetections) -- [**SemanticSegmentationParser**](luxonis_eval/parsers/semantic_seg.py) -> [depthai_nodes.SegmentationMask](https://github.com/luxonis/depthai-nodes/tree/main/depthai_nodes/message#segmentationmask) +- [**YOLOExtendedParser**](luxonis_eval/parsers/yolo.py) -> [dai.ImgDetections](https://docs.luxonis.com/software-v3/depthai/api/cpp/#classdai_1_1ImgDetections) +- [**SegmentationParser**](luxonis_eval/parsers/segmentation.py) -> [depthai_nodes.SegmentationMask](https://github.com/luxonis/depthai-nodes/tree/main/depthai_nodes/message#segmentationmask) > [!IMPORTANT] > The parser must produce outputs that the configured metrics can consume. For example, if a metric expects `dai.ImgDetections`, the parser must return that message type. diff --git a/configs/classification_config.yaml b/configs/classification_config.yaml new file mode 100644 index 0000000..465cc31 --- /dev/null +++ b/configs/classification_config.yaml @@ -0,0 +1,25 @@ +pipeline: + loader: + name: LuxonisLoader + params: + dataset_name: flowersclassification + view: [test] + + engine: + name: depthai + model_path: classification_light.rvc4.tar.xz + params: + device_ip: 10.12.143.168 + + evaluators: + - task_name: "" + parser: + name: ClassificationParser + params: + # Set to false if the NNArchive already outputs probabilities. + apply_softmax: true + metrics: + - name: TopKAccuracy + params: + topk: [1] + visualizers: [] diff --git a/configs/detection_config.yaml b/configs/detection_config.yaml new file mode 100644 index 0000000..ec8d9d3 --- /dev/null +++ b/configs/detection_config.yaml @@ -0,0 +1,34 @@ +pipeline: + loader: + name: LuxonisLoader + params: + dataset_name: coco-2017 + view: [test] + + engine: + name: onnx + model_path: path/to/yolov8n.onnx + params: + providers: [CPUExecutionProvider] + + # engine: + # name: depthai + # model_path: path/to/yolov8n.rvc4.tar.xz + # params: + # device_ip: 10.12.143.173 + + evaluators: + - task_name: "" + parser: + name: YOLOExtendedParser + params: + subtype: yolov8 + n_classes: 80 + conf_threshold: 0.25 + iou_threshold: 0.7 + max_det: 300 + metrics: + - name: BboxMeanAveragePrecision + params: + iou_type: bbox + visualizers: [] diff --git a/configs/pose_config.yaml b/configs/pose_config.yaml new file mode 100644 index 0000000..d6376f4 --- /dev/null +++ b/configs/pose_config.yaml @@ -0,0 +1,71 @@ +pipeline: + loader: + name: LuxonisLoader + params: + dataset_name: your_pose_dataset + view: [test] + class_mapping: + 0: person + + engine: + name: onnx + model_path: path/to/yolov8n-pose.onnx + params: + providers: [CPUExecutionProvider] + + # engine: + # name: depthai + # model_path: path/to/yolov8n-pose.rvc4.tar.xz + # params: + # device_ip: 10.12.143.173 + + evaluators: + - task_name: "" + parser: + name: YOLOExtendedParser + params: + subtype: yolov8 + n_classes: 1 + conf_threshold: 0.25 + iou_threshold: 0.7 + max_det: 300 + keypoint_label_names: + - nose + - left_eye + - right_eye + - left_ear + - right_ear + - left_shoulder + - right_shoulder + - left_elbow + - right_elbow + - left_wrist + - right_wrist + - left_hip + - right_hip + - left_knee + - right_knee + - left_ankle + - right_ankle + keypoint_edges: + - [0, 1] + - [0, 2] + - [1, 3] + - [2, 4] + - [5, 6] + - [5, 7] + - [7, 9] + - [6, 8] + - [8, 10] + - [5, 11] + - [6, 12] + - [11, 12] + - [11, 13] + - [13, 15] + - [12, 14] + - [14, 16] + metrics: + - name: KeypointMeanAveragePrecision + params: + compute_area_from_keypoints: true + visualizers: [] diff --git a/configs/semantic_segmentation_config.yaml b/configs/semantic_segmentation_config.yaml new file mode 100644 index 0000000..b53d7e9 --- /dev/null +++ b/configs/semantic_segmentation_config.yaml @@ -0,0 +1,40 @@ +pipeline: + loader: + name: LuxonisLoader + params: + dataset_name: your_segmentation_dataset + view: [test] + + engine: + name: onnx + model_path: path/to/segmentation_model.onnx + params: + providers: [CPUExecutionProvider] + + # engine: + # name: depthai + # model_path: path/to/segmentation_model.rvc4.tar.xz + # params: + # device_ip: 10.12.143.173 + + evaluators: + - task_name: "" + parser: + name: SegmentationParser + params: + # Set to true for binary / single-logit outputs that encode one class in one channel. + classes_in_one_layer: false + metrics: + - name: MIoU + params: + num_classes: 21 + include_background: false + per_class: true + input_format: index + - name: DiceCoefficient + params: + num_classes: 21 + include_background: false + average: micro + input_format: index + visualizers: [] diff --git a/configs/yolov8n_inst_seg_config.yaml b/configs/yolov8n_inst_seg_config.yaml index 4342a68..a029c95 100644 --- a/configs/yolov8n_inst_seg_config.yaml +++ b/configs/yolov8n_inst_seg_config.yaml @@ -27,7 +27,7 @@ pipeline: evaluators: - parser: - name: YOLOInstanceSegmentationParser + name: YOLOExtendedParser params: subtype: yolov8 n_classes: 80 diff --git a/luxonis_eval/config/resolver.py b/luxonis_eval/config/resolver.py index ccb0126..bf51614 100644 --- a/luxonis_eval/config/resolver.py +++ b/luxonis_eval/config/resolver.py @@ -292,9 +292,9 @@ def _map_archive_parser_name(self, parser_name: str) -> str: mapping = { "Classification": "ClassificationParser", "ClassificationParser": "ClassificationParser", - "SemanticSegmentation": "SemanticSegmentationParser", - "SemanticSegmentationParser": "SemanticSegmentationParser", - "Segmentation": "SemanticSegmentationParser", + "SemanticSegmentation": "SegmentationParser", + "SemanticSegmentationParser": "SegmentationParser", + "Segmentation": "SegmentationParser", } resolved_name = mapping.get(parser_name, parser_name) if resolved_name not in PARSERS_REGISTRY: @@ -306,12 +306,7 @@ def _map_archive_parser_name(self, parser_name: str) -> str: def resolve_yolo_archive_parser(head: Any) -> ParserConfig: metadata = resolve_head_metadata(head) - if "n_keypoints" in metadata: - parser_name = "YOLOKeypointDetectionParser" - elif "n_prototypes" in metadata or "mask_outputs" in metadata: - parser_name = "YOLOInstanceSegmentationParser" - else: - parser_name = "YOLODetectionParser" + parser_name = "YOLOExtendedParser" params = {} for key in ( diff --git a/luxonis_eval/core/core.py b/luxonis_eval/core/core.py index 4ae4624..1d7cd65 100644 --- a/luxonis_eval/core/core.py +++ b/luxonis_eval/core/core.py @@ -146,6 +146,7 @@ def _run_evaluators(self) -> dict[str, Any]: assert self.parser is not None assert self.throughput_metric is not None assert self.evaluator_cfg is not None + assert self.model_spec is not None with self._progress( f"Running {engine_name.upper()} inference ({model_name})...", @@ -169,6 +170,7 @@ def _run_evaluators(self) -> dict[str, Any]: raw_output, self.evaluator_cfg.outputs, ), + model_spec=self.model_spec, class_map=self.class_map, **self.evaluator_cfg.parser.params, ) @@ -261,6 +263,7 @@ def _sanity_check_pipeline(self) -> None: assert self.engine is not None assert self.parser is not None assert self.evaluator_cfg is not None + assert self.model_spec is not None if len(self.loader) == 0: raise ValueError( @@ -279,6 +282,7 @@ def _sanity_check_pipeline(self) -> None: raw_output = self.engine.infer_once(img) predictions = self.parser.parse( select_evaluator_outputs(raw_output, self.evaluator_cfg.outputs), + model_spec=self.model_spec, class_map=self.class_map, **self.evaluator_cfg.parser.params, ) diff --git a/luxonis_eval/core/factories.py b/luxonis_eval/core/factories.py index aca4f02..010e7b6 100644 --- a/luxonis_eval/core/factories.py +++ b/luxonis_eval/core/factories.py @@ -25,6 +25,7 @@ def create_engine(cfg: EvalConfig) -> BaseEngine: ENGINES_REGISTRY, cfg.pipeline.engine.name, cfg.pipeline.engine.model_path, + nn_archive_cfg=cfg.nn_archive_cfg, **cfg.pipeline.engine.params, ) except KeyError as e: diff --git a/luxonis_eval/core/runtime.py b/luxonis_eval/core/runtime.py index 8fba0e4..7b7148d 100644 --- a/luxonis_eval/core/runtime.py +++ b/luxonis_eval/core/runtime.py @@ -10,6 +10,7 @@ from luxonis_eval.config import EvaluatorConfig from luxonis_eval.engines.base_engine import ModelSpec +from luxonis_eval.engines.io import EngineOutput from luxonis_eval.loaders.base_loader import BaseEvalLoader from luxonis_eval.metrics.metrics_utils import normalized_xywh_to_coco_xywh @@ -82,19 +83,13 @@ def get_metric_ctx(base_ctx: dict[str, Any], **kwargs: Any) -> dict[str, Any]: def select_evaluator_outputs( - raw_output: Any, + raw_output: EngineOutput, outputs: list[str] | None, -) -> Any: +) -> EngineOutput: if not outputs: return raw_output - # TODO: Implement filtering - - logger.warning( - "Evaluator outputs filtering is not implemented for this engine " - "output type yet. Using all raw outputs." - ) - return raw_output + return raw_output.select(outputs) def resolve_luxonis_task_name( diff --git a/luxonis_eval/engines/__init__.py b/luxonis_eval/engines/__init__.py index 6774d29..8efacae 100644 --- a/luxonis_eval/engines/__init__.py +++ b/luxonis_eval/engines/__init__.py @@ -1,10 +1,15 @@ from .base_engine import BaseEngine, ModelSpec -from .depthai_engine import DepthAIEngine -from .onnx_engine import OnnxEngine +from .depthai_engine import DepthAIEngine, DepthAIEngineOutput +from .io import EngineOutput, TensorSpec +from .onnx_engine import ONNXEngineOutput, OnnxEngine __all__ = [ "BaseEngine", + "DepthAIEngineOutput", "DepthAIEngine", + "EngineOutput", "ModelSpec", "OnnxEngine", + "ONNXEngineOutput", + "TensorSpec", ] diff --git a/luxonis_eval/engines/base_engine.py b/luxonis_eval/engines/base_engine.py index 3b3805a..28a0379 100644 --- a/luxonis_eval/engines/base_engine.py +++ b/luxonis_eval/engines/base_engine.py @@ -1,5 +1,5 @@ +import inspect from abc import ABC, abstractmethod -from dataclasses import dataclass from pathlib import Path from typing import Any @@ -7,23 +7,10 @@ from luxonis_ml.typing import PathType from luxonis_ml.utils.registry import AutoRegisterMeta +from luxonis_eval.engines.io import EngineOutput, ModelSpec from luxonis_eval.registry import ENGINES_REGISTRY -@dataclass(frozen=True) -class ModelSpec: - """Validated model input specification.""" - - width: int - height: int - - def __post_init__(self) -> None: - if self.width <= 0 or self.height <= 0: - raise ValueError( - "ModelSpec width and height must be positive integers." - ) - - class BaseEngine( ABC, metaclass=AutoRegisterMeta, @@ -32,6 +19,21 @@ class BaseEngine( ): """Abstract base class for inference engines.""" + output_type: type[EngineOutput] | None = None + + def __init_subclass__(cls, **kwargs: Any) -> None: + super().__init_subclass__(**kwargs) + if inspect.isabstract(cls): + return + if "output_type" not in cls.__dict__ or cls.output_type is None: + raise TypeError( + f"{cls.__name__} must define `output_type` as its EngineOutput class." + ) + if not issubclass(cls.output_type, EngineOutput): + raise TypeError( + f"{cls.__name__}.output_type must be a subclass of EngineOutput." + ) + def __init__(self, model_path: PathType, **kwargs: Any) -> None: """Initialize the engine. @@ -66,7 +68,7 @@ def _get_model_spec(self) -> ModelSpec: return self.model_spec @abstractmethod - def infer_once(self, img: np.ndarray) -> Any: + def infer_once(self, img: np.ndarray) -> EngineOutput: """Run inference on a single image.""" ... diff --git a/luxonis_eval/engines/depthai_engine.py b/luxonis_eval/engines/depthai_engine.py index 3486c8e..ce33667 100644 --- a/luxonis_eval/engines/depthai_engine.py +++ b/luxonis_eval/engines/depthai_engine.py @@ -1,21 +1,84 @@ -from typing import Any +from collections.abc import Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any import depthai as dai import numpy as np -from depthai import ADatatype from loguru import logger from luxonis_ml.typing import PathType from luxonis_eval.engines.base_engine import BaseEngine, ModelSpec +from luxonis_eval.engines.io import EngineOutput, TensorLayout, TensorSpec + +if TYPE_CHECKING: + from luxonis_ml.nn_archive.config import Config as NNArchiveConfig + + +@dataclass(frozen=True, slots=True) +class DepthAIEngineOutput(EngineOutput): + raw_output: dai.NNData + _selected_names: tuple[str, ...] | None = None + + def names(self) -> tuple[str, ...]: + available = tuple(self.raw_output.getAllLayerNames()) + if self._selected_names is None: + return available + + missing = [name for name in self._selected_names if name not in available] + if missing: + raise ValueError( + f"Requested outputs are missing from engine result: {missing}" + ) + return self._selected_names + + def get( + self, + name: str, + *, + layout: TensorLayout | None = None, + ) -> np.ndarray: + if name not in self.names(): + raise ValueError( + f"Requested output tensor {name!r} is not available. " + f"Available tensors: {list(self.names())}." + ) + + get_tensor_kwargs: dict[str, object] = {"dequantize": True} + if layout == "NCHW": + get_tensor_kwargs["storageOrder"] = ( + dai.TensorInfo.StorageOrder.NCHW + ) + elif layout == "NHWC": + get_tensor_kwargs["storageOrder"] = ( + dai.TensorInfo.StorageOrder.NHWC + ) + + tensor = self.raw_output.getTensor(name, **get_tensor_kwargs) + return np.asarray(tensor) + + def select(self, names: Sequence[str] | None) -> "DepthAIEngineOutput": + if names is None: + return self + + requested_names = tuple(names) + missing = [name for name in requested_names if name not in self.names()] + if missing: + raise ValueError( + f"Requested outputs are missing from engine result: {missing}" + ) + return DepthAIEngineOutput(self.raw_output, requested_names) class DepthAIEngine(BaseEngine, register_name="depthai"): """DepthAI inference engine.""" + output_type = DepthAIEngineOutput + def __init__( self, model_path: PathType, device_ip: str | None = None, + nn_archive_cfg: "NNArchiveConfig | None" = None, **kwargs: Any, ) -> None: """Initialize the DepthAI inference engine. @@ -31,11 +94,13 @@ def __init__( """ super().__init__(model_path=model_path, **kwargs) self.device_ip = device_ip + self.nn_archive_cfg = nn_archive_cfg self._pipeline: dai.Pipeline | None = None self.device: dai.Device | None = None self.device_platform: str | None = None self.nn_archive: dai.NNArchive | None = None - self.input_info: dict[str, Any] = {} + self.input_spec: TensorSpec | None = None + self.output_specs: tuple[TensorSpec, ...] = () self.model_platform: str | None = None self._input_queue: Any = None self._output_queue: Any = None @@ -45,7 +110,12 @@ def setup(self) -> ModelSpec: """Set up the DepthAI pipeline and resolve model spec.""" if self._pipeline is None: self.device, self.device_platform = self._setup_device() - self.nn_archive, self.input_info, self.model_platform = ( + ( + self.nn_archive, + self.input_spec, + self.output_specs, + self.model_platform, + ) = ( self._load_nn_archive() ) @@ -85,35 +155,27 @@ def _resolve_platform_name(self) -> str: def _resolve_model_spec(self) -> ModelSpec: """Resolve model input dimensions from the loaded archive.""" - if not self.input_info or "shape" not in self.input_info: + if self.input_spec is None or self.input_spec.shape is None: raise ValueError("Invalid input shape information.") - shape = self.input_info["shape"] + shape = self.input_spec.shape platform_name = self._resolve_platform_name() if platform_name == "RVC2": # RVC2 uses NCHW format: [batch, channels, height, width] - if len(shape) != 4: + if len(shape) != 4 or self.input_spec.layout != "NCHW": raise ValueError( f"Unexpected input shape for RVC2: {shape}. Expected input shape in NCHW format." ) - height, width = shape[2], shape[3] else: # RVC4 uses NHWC format: [batch, height, width, channels] - if len(shape) != 4: + if len(shape) != 4 or self.input_spec.layout != "NHWC": raise ValueError( f"Unexpected input shape for {platform_name}: {shape}. Expected input shape in NHWC format." ) - height, width = shape[1], shape[2] + return ModelSpec(input=self.input_spec, outputs=self.output_specs) - if not isinstance(width, int) or not isinstance(height, int): - raise TypeError( - f"DepthAI input shape must be statically defined. Got {shape}." - ) - - return ModelSpec(width=width, height=height) - - def infer_once(self, img: np.ndarray) -> ADatatype: + def infer_once(self, img: np.ndarray) -> DepthAIEngineOutput: """Run inference on a single image using DepthAI. Parameters @@ -160,7 +222,7 @@ def infer_once(self, img: np.ndarray) -> ADatatype: new_input.setType(img_frame_type) self._input_queue.send(new_input) - return self._output_queue.get() + return DepthAIEngineOutput(self._output_queue.get()) def vis_frame(self) -> np.ndarray: """Get visualization frame from passthrough. @@ -189,7 +251,8 @@ def close(self) -> None: self.device = None self.device_platform = None self.nn_archive = None - self.input_info = {} + self.input_spec = None + self.output_specs = () self.model_platform = None self._input_queue = None self._output_queue = None @@ -227,7 +290,14 @@ def _setup_device(self) -> tuple[dai.Device, str]: return device, platform_name - def _load_nn_archive(self) -> tuple[dai.NNArchive, dict, str | None]: + def _load_nn_archive( + self, + ) -> tuple[ + dai.NNArchive, + TensorSpec, + tuple[TensorSpec, ...], + str | None, + ]: """Load the model from an NNArchive. Returns @@ -247,24 +317,71 @@ def _load_nn_archive(self) -> tuple[dai.NNArchive, dict, str | None]: logger.error(f"Failed to load model: {e}") raise - input_info = {} + input_spec: TensorSpec | None = None + output_specs: tuple[TensorSpec, ...] = () infered_platform = None try: inputs = nn_archive.getConfig().model.inputs - if inputs: - input_shape = inputs[0].shape - input_info = { - "shape": input_shape, - "name": inputs[0].name - if hasattr(inputs[0], "name") - else "input", - } - logger.info(f"Model input shape: {input_shape}") - if inputs[0].layout and inputs[0].layout == "NHWC": - infered_platform = "RVC4" - elif inputs[0].layout and inputs[0].layout == "NCHW": - infered_platform = "RVC2" + if len(inputs) != 1: + raise NotImplementedError( + "Only single-input NNArchive models are supported in luxonis-eval." + ) + + input_meta = inputs[0] + input_shape = tuple(input_meta.shape) + input_layout = getattr(input_meta, "layout", None) + if input_layout is not None and not isinstance(input_layout, str): + input_layout = ( + getattr(input_layout, "name", None) + or getattr(input_layout, "value", None) + ) + input_spec = TensorSpec( + name=getattr(input_meta, "name", "input"), + shape=input_shape, + dtype=str(getattr(input_meta, "dtype", None)) + if getattr(input_meta, "dtype", None) is not None + else None, + layout=input_layout, + ) + logger.info(f"Model input shape: {input_shape}") + if input_layout == "NHWC": + infered_platform = "RVC4" + elif input_layout == "NCHW": + infered_platform = "RVC2" + + outputs = getattr(nn_archive.getConfig().model, "outputs", []) + output_specs = tuple( + TensorSpec( + name=getattr(output_meta, "name", f"output_{idx}"), + shape=tuple(output_meta.shape) + if getattr(output_meta, "shape", None) is not None + else None, + dtype=str(getattr(output_meta, "dtype", None)) + if getattr(output_meta, "dtype", None) is not None + else None, + layout=( + getattr(output_meta, "layout", None) + if isinstance(getattr(output_meta, "layout", None), str) + or getattr(output_meta, "layout", None) is None + else getattr( + getattr(output_meta, "layout", None), "name", None + ) + or getattr( + getattr(output_meta, "layout", None), "value", None + ) + ), + ) + for idx, output_meta in enumerate(outputs) + ) except AttributeError: logger.warning("Could not extract input shape from model") - return nn_archive, input_info, infered_platform + if input_spec is None: + raise ValueError("Could not extract model input spec from NNArchive.") + + return ( + nn_archive, + input_spec, + output_specs, + infered_platform, + ) diff --git a/luxonis_eval/engines/io.py b/luxonis_eval/engines/io.py new file mode 100644 index 0000000..61a25c6 --- /dev/null +++ b/luxonis_eval/engines/io.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Literal + +import numpy as np + +TensorLayout = Literal["NCHW", "NHWC", "NC", "CHW", "HWC", "HW"] + + +@dataclass(frozen=True, slots=True) +class TensorSpec: + name: str + shape: tuple[int, ...] | None = None + dtype: str | None = None + layout: TensorLayout | None = None + + +@dataclass(frozen=True, slots=True) +class ModelSpec: + input: TensorSpec + outputs: tuple[TensorSpec, ...] + + def __post_init__(self) -> None: + if self.input.shape is None or len(self.input.shape) != 4: + raise ValueError( + "ModelSpec input shape must be a statically defined 4D tensor." + ) + if self.input.layout not in {"NCHW", "NHWC"}: + raise ValueError( + "ModelSpec input layout must be either 'NCHW' or 'NHWC'." + ) + if not self.outputs: + raise ValueError("ModelSpec must define at least one output.") + + @property + def width(self) -> int: + width_idx = 3 if self.input.layout == "NCHW" else 2 + width = self.input.shape[width_idx] + if not isinstance(width, int) or width <= 0: + raise ValueError( + f"ModelSpec width must be a positive integer, got {width!r}." + ) + return width + + @property + def height(self) -> int: + height_idx = 2 if self.input.layout == "NCHW" else 1 + height = self.input.shape[height_idx] + if not isinstance(height, int) or height <= 0: + raise ValueError( + f"ModelSpec height must be a positive integer, got {height!r}." + ) + return height + + +class EngineOutput(ABC): + """Parser-facing abstraction over backend-specific inference outputs. + + Implemented engines wrap their native output type behind this interface + """ + + @abstractmethod + def names(self) -> tuple[str, ...]: + """Return output tensor names in engine-defined order.""" + + @abstractmethod + def get( + self, + name: str, + *, + layout: TensorLayout | None = None, + ) -> np.ndarray: + """Return one named tensor as a NumPy array. + """ + + @abstractmethod + def select(self, names: Sequence[str] | None) -> "EngineOutput": + """Return a filtered view over this output. + """ + + def first(self) -> tuple[str, np.ndarray]: + """Convenience helper for single-output models.""" + names = self.names() + if len(names) != 1: + raise ValueError( + f"Expected exactly one output tensor, got {list(names)}." + ) + name = names[0] + return name, self.get(name) diff --git a/luxonis_eval/engines/onnx_engine.py b/luxonis_eval/engines/onnx_engine.py index b429ceb..6a4900c 100644 --- a/luxonis_eval/engines/onnx_engine.py +++ b/luxonis_eval/engines/onnx_engine.py @@ -1,3 +1,5 @@ +from collections.abc import Sequence +from dataclasses import dataclass from typing import Any import numpy as np @@ -9,11 +11,51 @@ load_onnx_bytes_from_nn_archive, ) from luxonis_eval.engines.base_engine import BaseEngine, ModelSpec +from luxonis_eval.engines.io import EngineOutput, TensorLayout, TensorSpec + + +@dataclass(frozen=True, slots=True) +class ONNXEngineOutput(EngineOutput): + tensors: dict[str, np.ndarray] + + def names(self) -> tuple[str, ...]: + return tuple(self.tensors.keys()) + + def get( + self, + name: str, + *, + layout: TensorLayout | None = None, + ) -> np.ndarray: + del layout + try: + return self.tensors[name] + except KeyError as err: + raise ValueError( + f"Requested output tensor {name!r} is not available. " + f"Available tensors: {list(self.tensors)}." + ) from err + + def select(self, names: Sequence[str] | None) -> "ONNXEngineOutput": + if names is None: + return self + + missing = [name for name in names if name not in self.tensors] + if missing: + raise ValueError( + f"Requested outputs are missing from engine result: {missing}" + ) + + return ONNXEngineOutput( + tensors={name: self.tensors[name] for name in names} + ) class OnnxEngine(BaseEngine, register_name="onnx"): """ONNX Runtime inference engine.""" + output_type = ONNXEngineOutput + def __init__( self, model_path: PathType, @@ -35,6 +77,7 @@ def __init__( self.providers = providers or ["CPUExecutionProvider"] self._session: ort.InferenceSession | None = None self._input_name: str | None = None + self._output_names: tuple[str, ...] = () self._visualization_frame: np.ndarray | None = None def setup(self) -> ModelSpec: @@ -52,27 +95,54 @@ def setup(self) -> ModelSpec: self._session = ort.InferenceSession( session_source, providers=self.providers ) - self._input_name = self._session.get_inputs()[0].name + session_inputs = self._session.get_inputs() + if len(session_inputs) != 1: + raise NotImplementedError( + "Only single-input ONNX models are supported in luxonis-eval." + ) + self._input_name = session_inputs[0].name + self._output_names = tuple( + output_meta.name for output_meta in self._session.get_outputs() + ) if self.model_spec is not None: return self.model_spec - shape = self._session.get_inputs()[0].shape + input_meta = self._session.get_inputs()[0] + shape = input_meta.shape if len(shape) != 4: raise ValueError( f"Unexpected input shape for ONNX: {shape}. Expected input shape in NCHW format." ) - height = shape[2] - width = shape[3] - if not isinstance(width, int) or not isinstance(height, int): + if not all(isinstance(dim, int) for dim in shape): raise TypeError( f"ONNX input shape must be statically defined. Got {shape}." ) - return self._set_model_spec(ModelSpec(width=width, height=height)) + output_specs = tuple( + TensorSpec( + name=output_meta.name, + shape=tuple(output_meta.shape) + if all(isinstance(dim, int) for dim in output_meta.shape) + else None, + dtype=output_meta.type, + ) + for output_meta in self._session.get_outputs() + ) + return self._set_model_spec( + ModelSpec( + input=TensorSpec( + name=input_meta.name, + shape=tuple(shape), + dtype=input_meta.type, + layout="NCHW", + ), + outputs=output_specs, + ) + ) - def infer_once(self, img: np.ndarray) -> Any: + def infer_once(self, img: np.ndarray) -> ONNXEngineOutput: """Run inference on a single image using ONNX Runtime. Parameters @@ -106,7 +176,10 @@ def infer_once(self, img: np.ndarray) -> Any: "Expected `HW` or `HWC` loader output." ) - return self._session.run(None, {self._input_name: x}) + output_values = self._session.run(None, {self._input_name: x}) + return ONNXEngineOutput( + tensors=dict(zip(self._output_names, output_values, strict=True)) + ) def vis_frame(self) -> np.ndarray: """Return the visualization frame. @@ -124,5 +197,6 @@ def close(self) -> None: """Release ONNX Runtime resources.""" self._session = None self._input_name = None + self._output_names = () self._visualization_frame = None self.model_spec = None diff --git a/luxonis_eval/metrics/__init__.py b/luxonis_eval/metrics/__init__.py index 61f34cd..86b193a 100755 --- a/luxonis_eval/metrics/__init__.py +++ b/luxonis_eval/metrics/__init__.py @@ -1,9 +1,9 @@ from .base_metric import BaseMetric from .bbox_map import BboxMeanAveragePrecision -from .dice_coef import DiceCoefficient +from .dice_coef import DiceCoefficient, F1Score from .keypoint_map import KeypointMeanAveragePrecision from .mask_map import MaskMeanAveragePrecision -from .mIoU import MIoU +from .mIoU import JaccardIndex, MIoU from .throughput import ThroughputMetric from .topk_accuracy import TopKAccuracy @@ -11,6 +11,8 @@ "BaseMetric", "BboxMeanAveragePrecision", "DiceCoefficient", + "F1Score", + "JaccardIndex", "KeypointMeanAveragePrecision", "MIoU", "MaskMeanAveragePrecision", diff --git a/luxonis_eval/metrics/dice_coef.py b/luxonis_eval/metrics/dice_coef.py index a9a412c..96fcdad 100644 --- a/luxonis_eval/metrics/dice_coef.py +++ b/luxonis_eval/metrics/dice_coef.py @@ -1,15 +1,19 @@ from typing import Any, Literal +import depthai as dai import numpy as np import torch -from depthai_nodes import SegmentationMask from torchmetrics.segmentation import DiceScore from luxonis_eval.metrics.base_metric import BaseMetric from luxonis_eval.metrics.metrics_utils import ( + binary_segmentation_confusion, mask_ignore_pixels, + normalize_prediction_segmentation_mask, remap_prediction_mask, + target_segmentation_to_index_mask, ) +from luxonis_eval.utils.depthai_nodes import extract_segmentation_mask class DiceCoefficient(BaseMetric): @@ -66,7 +70,7 @@ def reset(self) -> None: def update( self, - predictions: SegmentationMask, + predictions: dai.SegmentationMask, target: dict[str, np.ndarray], **kwargs: Any, ) -> None: @@ -86,13 +90,22 @@ def update( target_bg = kwargs.get("target_bg") class_index_map = kwargs.get("class_index_map") - pred_mask: np.ndarray = predictions.mask - target_mask = np.argmax(target[self.required_target_keys()[0]], axis=0) + target_mask, binary_target = target_segmentation_to_index_mask( + target[self.required_target_keys()[0]] + ) + pred_mask = normalize_prediction_segmentation_mask( + extract_segmentation_mask(predictions), + binary_target=binary_target, + ) - if class_index_map is not None: + if class_index_map is not None and not binary_target: pred_mask = remap_prediction_mask(pred_mask, class_index_map) - if not self.include_background and target_bg is not None: + if ( + not binary_target + and not self.include_background + and target_bg is not None + ): pred_mask, target_mask = mask_ignore_pixels( pred_mask, target_mask, ignore_index=target_bg ) @@ -111,3 +124,73 @@ def compute(self) -> dict[str, float]: Computed Dice coefficient results. """ return {"Dice Score": float(self.metric.compute())} + + +class F1Score(BaseMetric): + """LuxonisTrain-compatible F1 score for binary semantic segmentation.""" + + def __init__( + self, + num_classes: int | None = None, + include_background: bool = False, + average: Literal["micro", "macro", "weighted", "none"] + | None = "micro", + input_format: Literal["one-hot", "index"] = "index", + **kwargs: Any, + ) -> None: + self.num_classes = num_classes + self.include_background = include_background + self.average = average + self.input_format = input_format + self.target_class_map = None + super().__init__(**kwargs) + + def required_target_keys(self) -> list[str]: + return ["/segmentation"] + + def reset(self) -> None: + self.true_positives = 0 + self.false_positives = 0 + self.false_negatives = 0 + + def update( + self, + predictions: dai.SegmentationMask, + target: dict[str, np.ndarray], + **kwargs: Any, + ) -> None: + if self.target_class_map is None: + self.target_class_map = kwargs.get("target_class_map", {}) + class_index_map = kwargs.get("class_index_map") + target_bg = kwargs.get("target_bg") + + target_mask, binary_target = target_segmentation_to_index_mask( + target[self.required_target_keys()[0]] + ) + pred_mask = normalize_prediction_segmentation_mask( + extract_segmentation_mask(predictions), + binary_target=binary_target, + ) + + if binary_target: + tp, fp, fn = binary_segmentation_confusion(pred_mask, target_mask) + self.true_positives += tp + self.false_positives += fp + self.false_negatives += fn + return + + del class_index_map, target_bg, pred_mask, target_mask + raise NotImplementedError( + "luxonis-eval `F1Score` currently mirrors LuxonisTrain " + "behavior only for binary semantic segmentation. Use " + "`DiceCoefficient` for non-binary semantic segmentation." + ) + + def compute(self) -> dict[str, float]: + denom = ( + 2 * self.true_positives + + self.false_positives + + self.false_negatives + ) + score = 0.0 if denom == 0 else (2 * self.true_positives) / denom + return {"F1Score": float(score)} diff --git a/luxonis_eval/metrics/mIoU.py b/luxonis_eval/metrics/mIoU.py index 304d05d..626a00a 100644 --- a/luxonis_eval/metrics/mIoU.py +++ b/luxonis_eval/metrics/mIoU.py @@ -1,15 +1,19 @@ from typing import Any, Literal +import depthai as dai import numpy as np import torch -from depthai_nodes import SegmentationMask from torchmetrics.segmentation import MeanIoU from luxonis_eval.metrics.base_metric import BaseMetric from luxonis_eval.metrics.metrics_utils import ( + binary_segmentation_confusion, mask_ignore_pixels, + normalize_prediction_segmentation_mask, remap_prediction_mask, + target_segmentation_to_index_mask, ) +from luxonis_eval.utils.depthai_nodes import extract_segmentation_mask class MIoU(BaseMetric): @@ -66,7 +70,7 @@ def reset(self) -> None: def update( self, - predictions: SegmentationMask, + predictions: dai.SegmentationMask, target: dict[str, np.ndarray], **kwargs: Any, ) -> None: @@ -87,13 +91,22 @@ def update( target_bg = kwargs.get("target_bg") class_index_map = kwargs.get("class_index_map") - pred_mask: np.ndarray = predictions.mask - target_mask = np.argmax(target[self.required_target_keys()[0]], axis=0) + target_mask, binary_target = target_segmentation_to_index_mask( + target[self.required_target_keys()[0]] + ) + pred_mask = normalize_prediction_segmentation_mask( + extract_segmentation_mask(predictions), + binary_target=binary_target, + ) - if class_index_map is not None: + if class_index_map is not None and not binary_target: pred_mask = remap_prediction_mask(pred_mask, class_index_map) - if not self.include_background and target_bg is not None: + if ( + not binary_target + and not self.include_background + and target_bg is not None + ): pred_mask, target_mask = mask_ignore_pixels( pred_mask, target_mask, ignore_index=target_bg ) @@ -127,3 +140,72 @@ def compute(self) -> dict[str, float]: f"mIoU ({name})": float(r) for name, r in zip(class_names, results, strict=True) } + + +class JaccardIndex(BaseMetric): + """LuxonisTrain-compatible Jaccard index for binary segmentation.""" + + def __init__( + self, + num_classes: int | None = None, + include_background: bool = False, + per_class: bool = False, + input_format: Literal["one-hot", "index", "mixed"] = "index", + **kwargs: Any, + ) -> None: + self.num_classes = num_classes + self.include_background = include_background + self.per_class = per_class + self.input_format = input_format + self.target_class_map = None + super().__init__(**kwargs) + + def required_target_keys(self) -> list[str]: + return ["/segmentation"] + + def reset(self) -> None: + self.true_positives = 0 + self.false_positives = 0 + self.false_negatives = 0 + + def update( + self, + predictions: dai.SegmentationMask, + target: dict[str, np.ndarray], + **kwargs: Any, + ) -> None: + if self.target_class_map is None: + self.target_class_map = kwargs.get("target_class_map", {}) + class_index_map = kwargs.get("class_index_map") + target_bg = kwargs.get("target_bg") + + target_mask, binary_target = target_segmentation_to_index_mask( + target[self.required_target_keys()[0]] + ) + pred_mask = normalize_prediction_segmentation_mask( + extract_segmentation_mask(predictions), + binary_target=binary_target, + ) + + if binary_target: + tp, fp, fn = binary_segmentation_confusion(pred_mask, target_mask) + self.true_positives += tp + self.false_positives += fp + self.false_negatives += fn + return + + del class_index_map, target_bg, pred_mask, target_mask + raise NotImplementedError( + "luxonis-eval `JaccardIndex` currently mirrors LuxonisTrain " + "behavior only for binary semantic segmentation. Use `MIoU` " + "for non-binary semantic segmentation." + ) + + def compute(self) -> dict[str, float]: + denom = ( + self.true_positives + + self.false_positives + + self.false_negatives + ) + score = 0.0 if denom == 0 else self.true_positives / denom + return {"JaccardIndex": float(score)} diff --git a/luxonis_eval/metrics/metrics_utils.py b/luxonis_eval/metrics/metrics_utils.py index b854b72..ac9eb92 100644 --- a/luxonis_eval/metrics/metrics_utils.py +++ b/luxonis_eval/metrics/metrics_utils.py @@ -117,6 +117,42 @@ def remap_prediction_mask( return remapped +def target_segmentation_to_index_mask( + target_mask: np.ndarray, +) -> tuple[np.ndarray, bool]: + """Convert a segmentation target tensor to an index mask.""" + mask = np.asarray(target_mask) + + if mask.ndim == 2: + return mask.astype(np.int64), False + + if mask.ndim != 3: + raise ValueError( + f"Expected segmentation target with 2 or 3 dimensions, got {mask.ndim}." + ) + + if mask.shape[0] == 1: + return mask[0].astype(np.int64), True + + return np.argmax(mask, axis=0).astype(np.int64), False + + +def normalize_prediction_segmentation_mask( + pred_mask: np.ndarray, + *, + binary_target: bool, +) -> np.ndarray: + """Normalize a predicted segmentation mask to class indices.""" + mask = np.asarray(pred_mask).astype(np.int64) + + if binary_target: + # Current depthai-nodes binary semantic parsing emits 255 for + # background/unassigned pixels and 0 for the only foreground class. + return (mask == 0).astype(np.int64) + + return mask + + def mask_ignore_pixels( pred_mask: np.ndarray, target_mask: np.ndarray, ignore_index: int = 0 ) -> tuple[np.ndarray, np.ndarray]: @@ -147,6 +183,20 @@ def mask_ignore_pixels( return pred_mask, target_mask +def binary_segmentation_confusion( + pred_mask: np.ndarray, + target_mask: np.ndarray, +) -> tuple[int, int, int]: + """Compute TP, FP and FN for binary segmentation masks.""" + pred_fg = np.asarray(pred_mask) > 0 + target_fg = np.asarray(target_mask) > 0 + + tp = int(np.logical_and(pred_fg, target_fg).sum()) + fp = int(np.logical_and(pred_fg, np.logical_not(target_fg)).sum()) + fn = int(np.logical_and(np.logical_not(pred_fg), target_fg).sum()) + return tp, fp, fn + + def to_coco_kpts_flat(kpts: np.ndarray) -> list[float]: """Convert keypoints to COCO flat list [x,y,v,...]. diff --git a/luxonis_eval/parsers/__init__.py b/luxonis_eval/parsers/__init__.py index ff26c41..2332084 100755 --- a/luxonis_eval/parsers/__init__.py +++ b/luxonis_eval/parsers/__init__.py @@ -1,15 +1,11 @@ from .base_parser import BaseParser from .classification import ClassificationParser -from .detection import YOLODetectionParser -from .instance_seg import YOLOInstanceSegmentationParser -from .keypoint_detection import YOLOKeypointDetectionParser -from .semantic_seg import SemanticSegmentationParser +from .segmentation import SegmentationParser +from .yolo import YOLOExtendedParser __all__ = [ "BaseParser", "ClassificationParser", - "SemanticSegmentationParser", - "YOLODetectionParser", - "YOLOInstanceSegmentationParser", - "YOLOKeypointDetectionParser", + "SegmentationParser", + "YOLOExtendedParser", ] diff --git a/luxonis_eval/parsers/base_parser.py b/luxonis_eval/parsers/base_parser.py index 8b7953f..9208851 100644 --- a/luxonis_eval/parsers/base_parser.py +++ b/luxonis_eval/parsers/base_parser.py @@ -3,6 +3,8 @@ from luxonis_ml.utils.registry import AutoRegisterMeta +from luxonis_eval.engines.base_engine import ModelSpec +from luxonis_eval.engines.io import EngineOutput from luxonis_eval.registry import PARSERS_REGISTRY @@ -24,6 +26,11 @@ def __init__(self, **kwargs: Any) -> None: """ @abstractmethod - def parse(self, raw_output: Any, **kwargs: Any) -> Any: + def parse( + self, + output: EngineOutput, + model_spec: ModelSpec, + **kwargs: Any, + ) -> Any: """Parse raw backend output into predictions.""" ... diff --git a/luxonis_eval/parsers/classification.py b/luxonis_eval/parsers/classification.py index 39227d3..87e34fe 100644 --- a/luxonis_eval/parsers/classification.py +++ b/luxonis_eval/parsers/classification.py @@ -1,12 +1,16 @@ from typing import Any -import depthai as dai import numpy as np from depthai_nodes import Classifications from depthai_nodes.message.creators import create_classification_message -from loguru import logger +from depthai_nodes.node.parsers.classification import ( + ClassificationParser as DepthAINodesClassificationParser, +) +from luxonis_eval.engines.base_engine import ModelSpec +from luxonis_eval.engines.io import EngineOutput from luxonis_eval.parsers.base_parser import BaseParser +from luxonis_eval.utils.depthai_nodes import ordered_class_names class ClassificationParser(BaseParser): @@ -18,7 +22,8 @@ def __init__(self, **kwargs: Any) -> None: def parse( self, - raw_output: dai.NNData | list[np.ndarray], + output: EngineOutput, + model_spec: ModelSpec, *, class_map: dict[int, str], apply_softmax: bool = False, @@ -28,8 +33,10 @@ def parse( Parameters ---------- - raw_output : dai.NNData | list[np.ndarray] - Backend inference output. + output : EngineOutput + Engine-normalized inference output. + model_spec : ModelSpec + Resolved model IO metadata. apply_softmax : bool, default=False Whether to apply softmax to the output scores. **kwargs : Any @@ -40,44 +47,13 @@ def parse( Classifications Classification scores. """ - classes = list(class_map.values()) - if isinstance(raw_output, dai.NNData): - layer_names = raw_output.getAllLayerNames() - logger.debug(f"Processing output with layers: {layer_names}") - output_name = layer_names[0] - scores = raw_output.getTensor(output_name, dequantize=True) - elif isinstance(raw_output, list): - scores = raw_output[0] - else: - raise TypeError( - f"Unsupported raw_output type: {type(raw_output)}. Expected dai.NNData or list[np.ndarray]." - ) - - scores = np.array(scores).flatten() - - if apply_softmax: - scores = self._softmax(scores) + del model_spec, kwargs + classes = ordered_class_names(class_map) + _, scores = output.first() + scores = np.asarray(scores).flatten() + scores = DepthAINodesClassificationParser.compute( + scores, + is_softmax=not apply_softmax, + ) return create_classification_message(classes=classes, scores=scores) - - def _softmax( - self, x: np.ndarray, axis: int | None = None, keep_dims: bool = False - ) -> np.ndarray: - """Apply softmax to an array. - - Parameters - ---------- - x : np.ndarray - Input array. - axis : int | None, optional - Axis over which to apply softmax. - keep_dims : bool, default=False - Whether to keep reduced dimensions. - - Returns - ------- - np.ndarray - Softmax-normalized array. - """ - ex = np.exp(x) - return ex / np.sum(ex, axis=axis, keepdims=keep_dims) diff --git a/luxonis_eval/parsers/detection.py b/luxonis_eval/parsers/detection.py deleted file mode 100644 index 1981c51..0000000 --- a/luxonis_eval/parsers/detection.py +++ /dev/null @@ -1,158 +0,0 @@ -from typing import Any - -import depthai as dai -import numpy as np -from depthai_nodes.message.creators import create_detection_message -from depthai_nodes.node.parsers.utils import normalize_bboxes, xyxy_to_xywh -from depthai_nodes.node.parsers.utils.yolo import ( - YOLOSubtype, - decode_yolo_output, -) -from loguru import logger - -from .base_parser import BaseParser - - -class YOLODetectionParser(BaseParser): - """Parser for YOLO-based detection model outputs.""" - - def __init__(self, **kwargs: Any) -> None: - """Initialize the YOLO detection parser.""" - super().__init__(**kwargs) - - def parse( - self, - raw_output: dai.NNData | list[np.ndarray], - *, - class_map: dict[int, str], - subtype: str, - n_classes: int | None = None, - anchors: list[list[list[float]]] | None = None, - conf_threshold: float = 0.001, - iou_threshold: float = 0.7, - max_det: int = 300, - **kwargs: Any, - ) -> dai.ImgDetections: - """Parse backend output into detection predictions. - - Parameters - ---------- - raw_output : dai.NNData | list[np.ndarray] - Backend inference output. - class_map : dict[int, str] - Mapping from class indices to class names. - subtype : str - YOLO model subtype. - n_classes : int | None, optional - Number of classes. - anchors : list[list[list[float]]] | None, optional - Anchor boxes. - conf_threshold : float, default=0.001 - Confidence threshold. - iou_threshold : float, default=0.7 - IoU threshold. - max_det : int, default=300 - Maximum detections. - **kwargs : Any - Additional parser arguments. - - Returns - ------- - dai.ImgDetections - Detection results including boxes, scores, classes, and metadata. - """ - try: - subtype = YOLOSubtype(subtype.lower()) - except ValueError as err: - raise ValueError( - f"Invalid YOLO subtype {subtype}. Supported YOLO subtypes are {[e.value for e in YOLOSubtype][:-1]}." - ) from err - - if isinstance(raw_output, dai.NNData): - layer_names = raw_output.getAllLayerNames() - logger.debug(f"Processing output with layers: {layer_names}") - - outputs_names = sorted( - [n for n in layer_names if "_yolo" in n or "yolo-" in n] - ) - outputs_values = [ - raw_output.getTensor( - o, - dequantize=True, - storageOrder=dai.TensorInfo.StorageOrder.NCHW, - ).astype(np.float32) # type: ignore - for o in outputs_names - ] - elif isinstance(raw_output, list): - outputs_names = [f"output_{i}" for i in range(len(raw_output))] - outputs_values = raw_output - else: - raise TypeError( - f"Unsupported raw_output type: {type(raw_output)}. Expected dai.NNData or list[np.ndarray]." - ) - - strides = ( - [8, 16, 32] - if subtype - not in [YOLOSubtype.V3UT, YOLOSubtype.V3T, YOLOSubtype.V4T] - else [16, 32] - ) - input_shape = tuple( - dim * strides[0] for dim in outputs_values[0].shape[2:4] - ) - final_anchors: np.ndarray | None = ( - np.array(anchors).reshape(len(strides), -1) if anchors else None - ) - inferred_n_classes = ( - outputs_values[0].shape[1] - 5 - if not final_anchors - else (outputs_values[0].shape[1] // final_anchors.shape[0]) - 5 - ) - if n_classes and inferred_n_classes != n_classes: - raise ValueError( - f"The provided number of classes {n_classes} does not match the model's {inferred_n_classes}." - ) - - results = decode_yolo_output( - yolo_outputs=outputs_values, - strides=strides, - anchors=final_anchors, - kpts=None, - conf_thres=conf_threshold, - iou_thres=iou_threshold, - num_classes=inferred_n_classes, - det_mode=True, - subtype=subtype, - max_nms=max_det, - ) - - bboxes, labels, label_names, scores, additional_output = ( - [], - [], - [], - [], - [], - ) - for i in range(results.shape[0]): - bbox, conf, label, other = ( - results[i, :4], - results[i, 4], - results[i, 5].astype(int), - results[i, 6:], - ) - bbox = xyxy_to_xywh(bbox.reshape(1, 4)) - bbox = normalize_bboxes( - bbox, height=input_shape[0], width=input_shape[1] - )[0] - bboxes.append(bbox) - scores.append(float(conf)) - labels.append(int(label)) - label_names.append(class_map[int(label)]) - additional_output.append(other) - - return create_detection_message( - bboxes=np.array(bboxes), - scores=np.array(scores), - labels=np.array(labels), - label_names=label_names, - ) diff --git a/luxonis_eval/parsers/instance_seg.py b/luxonis_eval/parsers/instance_seg.py deleted file mode 100644 index 7915bae..0000000 --- a/luxonis_eval/parsers/instance_seg.py +++ /dev/null @@ -1,199 +0,0 @@ -from typing import Any - -import cv2 -import depthai as dai -import numpy as np -from depthai_nodes.message.creators import create_detection_message -from depthai_nodes.node.parsers.utils import normalize_bboxes, xyxy_to_xywh -from depthai_nodes.node.parsers.utils.masks_utils import ( - get_segmentation_outputs, - process_single_mask, -) -from depthai_nodes.node.parsers.utils.yolo import ( - YOLOSubtype, - decode_yolo_output, -) -from loguru import logger - -from .base_parser import BaseParser - - -class YOLOInstanceSegmentationParser(BaseParser): - """Parser for YOLO-based instance segmentation model outputs.""" - - def __init__(self, **kwargs: Any) -> None: - """Initialize the YOLO instance segmentation parser.""" - super().__init__(**kwargs) - - def parse( - self, - raw_output: dai.NNData | list[np.ndarray], - *, - class_map: dict[int, str], - subtype: str, - n_classes: int | None = None, - anchors: list[list[list[float]]] | None = None, - conf_threshold: float = 0.5, - iou_threshold: float = 0.5, - mask_conf: float = 0.5, - max_det: int = 300, - **kwargs: Any, - ) -> dai.ImgDetections: - """Parse backend output into detection predictions. - - Parameters - ---------- - raw_output : dai.NNData | list[np.ndarray] - Backend inference output. - class_map : dict[int, str] - Mapping from class indices to class names. - subtype : str - YOLO model subtype. - n_classes : int | None, optional - Number of classes. - anchors : list[list[list[float]]] | None, optional - Anchor boxes. - conf_threshold : float, default=0.5 - Confidence threshold. - iou_threshold : float, default=0.5 - IoU threshold. - mask_conf : float, default=0.5 - Mask confidence threshold. - max_det : int, default=300 - Maximum detections. - **kwargs : Any - Additional parser arguments. - - Returns - ------- - dai.ImgDetections - Detection results including boxes, scores, classes, and metadata. - """ - try: - subtype = YOLOSubtype(subtype.lower()) - except ValueError as err: - raise ValueError( - f"Invalid YOLO subtype {subtype}. Supported YOLO subtypes are {[e.value for e in YOLOSubtype][:-1]}." - ) from err - - if isinstance(raw_output, dai.NNData): - layer_names = raw_output.getAllLayerNames() - logger.debug(f"Processing output with layers: {layer_names}") - - outputs_names = sorted( - [n for n in layer_names if "_yolo" in n or "yolo-" in n] - ) - outputs_values = [ - raw_output.getTensor( - o, - dequantize=True, - storageOrder=dai.TensorInfo.StorageOrder.NCHW, - ).astype(np.float32) # type: ignore - for o in outputs_names - ] - ( - masks_outputs_values, - protos_output, - protos_len, - ) = get_segmentation_outputs(raw_output) - elif isinstance(raw_output, list): - outputs_names = [f"output_{i}" for i in range(len(raw_output))] - outputs_values = raw_output[:3] - masks_outputs_values = raw_output[3:-1] - protos_output = raw_output[-1] - protos_len = protos_output.shape[1] - else: - raise TypeError( - f"Unsupported raw_output type: {type(raw_output)}. Expected dai.NNData or list[np.ndarray]." - ) - - strides = ( - [8, 16, 32] - if subtype - not in [YOLOSubtype.V3UT, YOLOSubtype.V3T, YOLOSubtype.V4T] - else [16, 32] - ) - input_shape = tuple( - dim * strides[0] for dim in outputs_values[0].shape[2:4] - ) - final_anchors: np.ndarray | None = ( - np.array(anchors).reshape(len(strides), -1) if anchors else None - ) - inferred_n_classes = ( - outputs_values[0].shape[1] - 5 - if not final_anchors - else (outputs_values[0].shape[1] // final_anchors.shape[0]) - 5 - ) - if n_classes and inferred_n_classes != n_classes: - raise ValueError( - f"The provided number of classes {n_classes} does not match the model's {inferred_n_classes}." - ) - - results = decode_yolo_output( - yolo_outputs=outputs_values, - strides=strides, - anchors=final_anchors, - kpts=None, - conf_thres=conf_threshold, - iou_thres=iou_threshold, - num_classes=inferred_n_classes, - det_mode=False, - subtype=subtype, - max_nms=max_det, - ) - - bboxes, labels, label_names, scores, additional_output = ( - [], - [], - [], - [], - [], - ) - instance_masks: list[np.ndarray] = [] - for i in range(results.shape[0]): - bbox, conf, label, other = ( - results[i, :4], - results[i, 4], - results[i, 5].astype(int), - results[i, 6:], - ) - bbox = xyxy_to_xywh(bbox.reshape(1, 4)) - bbox = normalize_bboxes( - bbox, height=input_shape[0], width=input_shape[1] - )[0] - bboxes.append(bbox) - scores.append(float(conf)) - labels.append(int(label)) - label_names.append(class_map[int(label)]) - additional_output.append(other) - - seg_coeff = other.astype(int) - hi, ai, xi, yi = seg_coeff - mask_coeff = masks_outputs_values[hi][ - 0, ai * protos_len : (ai + 1) * protos_len, yi, xi - ] - mask = process_single_mask( - protos_output[0], mask_coeff, mask_conf, bbox - ) - - resized_mask = cv2.resize( - mask, - (input_shape[1], input_shape[0]), - interpolation=cv2.INTER_NEAREST, - ) - - bin_mask = resized_mask > 0 - instance_masks.append(bin_mask) - - final_mask = np.asarray(instance_masks.copy()) - if final_mask.size != 0: - # Flatten (N, H, W) to (N*H, W) since dai.ImgDetections expects a 2D mask. - final_mask = final_mask.reshape(-1, final_mask.shape[-1]) - - return create_detection_message( - bboxes=np.array(bboxes), - scores=np.array(scores), - labels=np.array(labels), - label_names=label_names, - masks=final_mask, - ) diff --git a/luxonis_eval/parsers/keypoint_detection.py b/luxonis_eval/parsers/keypoint_detection.py deleted file mode 100644 index 45ff8e1..0000000 --- a/luxonis_eval/parsers/keypoint_detection.py +++ /dev/null @@ -1,195 +0,0 @@ -from typing import Any - -import depthai as dai -import numpy as np -from depthai_nodes.message.creators import create_detection_message -from depthai_nodes.node.parsers.utils import normalize_bboxes, xyxy_to_xywh -from depthai_nodes.node.parsers.utils.yolo import ( - YOLOSubtype, - decode_yolo_output, - parse_kpts, -) -from loguru import logger - -from .base_parser import BaseParser - - -class YOLOKeypointDetectionParser(BaseParser): - """Parser for YOLO-based keypoint detection model outputs.""" - - def __init__(self, **kwargs: Any) -> None: - """Initialize the YOLO keypoint detection parser.""" - super().__init__(**kwargs) - - def parse( - self, - raw_output: dai.NNData | list[np.ndarray], - *, - class_map: dict[int, str], - subtype: str, - n_classes: int | None = None, - anchors: list[list[list[float]]] | None = None, - conf_threshold: float = 0.5, - iou_threshold: float = 0.5, - max_det: int = 300, - keypoint_label_names: list[str] | None = None, - keypoint_edges: list[tuple[int, int]] | None = None, - **kwargs: Any, - ) -> dai.ImgDetections: - """Parse backend output into detection predictions. - - Parameters - ---------- - raw_output : dai.NNData | list[np.ndarray] - Backend inference output. - class_map : dict[int, str] - Mapping from class indices to class names. - subtype : str - YOLO model subtype. - n_classes : int | None, optional - Number of classes. - anchors : list[list[list[float]]] | None, optional - Anchor boxes. - conf_threshold : float, default=0.5 - Confidence threshold. - iou_threshold : float, default=0.5 - IoU threshold. - max_det : int, default=300 - Maximum detections. - keypoint_label_names : list[str] | None, optional - Names of keypoint labels. - keypoint_edges : list[tuple[int, int]] | None, optional - Edges connecting keypoints. - **kwargs : Any - Additional parser arguments. - - Returns - ------- - dai.ImgDetections - Detection results including boxes, scores, classes, and metadata. - """ - try: - subtype = YOLOSubtype(subtype.lower()) - except ValueError as err: - raise ValueError( - f"Invalid YOLO subtype {subtype}. Supported YOLO subtypes are {[e.value for e in YOLOSubtype][:-1]}." - ) from err - - if isinstance(raw_output, dai.NNData): - layer_names = raw_output.getAllLayerNames() - logger.debug(f"Processing output with layers: {layer_names}") - - outputs_names = sorted( - [n for n in layer_names if "_yolo" in n or "yolo-" in n] - ) - outputs_values = [ - raw_output.getTensor( - o, - dequantize=True, - storageOrder=dai.TensorInfo.StorageOrder.NCHW, - ).astype(np.float32) # type: ignore - for o in outputs_names - ] - kpts_output_names = sorted( - [name for name in layer_names if "kpt_output" in name] - ) - kpts_outputs = [ - raw_output.getTensor( - o, - dequantize=True, - ).astype(np.float32) # type: ignore - for o in kpts_output_names - ] - elif isinstance(raw_output, list): - outputs_names = [f"output_{i}" for i in range(len(raw_output))] - outputs_values = raw_output[:3] - kpts_outputs = raw_output[3:] - else: - raise TypeError( - f"Unsupported raw_output type: {type(raw_output)}. Expected dai.NNData or list[np.ndarray]." - ) - - strides = ( - [8, 16, 32] - if subtype - not in [YOLOSubtype.V3UT, YOLOSubtype.V3T, YOLOSubtype.V4T] - else [16, 32] - ) - input_shape = tuple( - dim * strides[0] for dim in outputs_values[0].shape[2:4] - ) - final_anchors: np.ndarray | None = ( - np.array(anchors).reshape(len(strides), -1) if anchors else None - ) - inferred_n_classes = ( - outputs_values[0].shape[1] - 5 - if not final_anchors - else (outputs_values[0].shape[1] // final_anchors.shape[0]) - 5 - ) - if n_classes and inferred_n_classes != n_classes: - raise ValueError( - f"The provided number of classes {n_classes} does not match the model's {inferred_n_classes}." - ) - num_keypoints = kpts_outputs[0].shape[1] // 3 - - results = decode_yolo_output( - yolo_outputs=outputs_values, - strides=strides, - anchors=final_anchors, - kpts=kpts_outputs, - conf_thres=conf_threshold, - iou_thres=iou_threshold, - num_classes=inferred_n_classes, - det_mode=False, - subtype=subtype, - max_nms=max_det, - ) - - bboxes, labels, label_names, scores, additional_output = ( - [], - [], - [], - [], - [], - ) - for i in range(results.shape[0]): - bbox, conf, label, other = ( - results[i, :4], - results[i, 4], - results[i, 5].astype(int), - results[i, 6:], - ) - bbox = xyxy_to_xywh(bbox.reshape(1, 4)) - bbox = normalize_bboxes( - bbox, height=input_shape[0], width=input_shape[1] - )[0] - bboxes.append(bbox) - scores.append(float(conf)) - labels.append(int(label)) - label_names.append(class_map[int(label)]) - - kpts = parse_kpts(other, num_keypoints, input_shape) # type: ignore - additional_output.append(kpts) - - additional_output = np.array(additional_output) - keypoints = ( - additional_output[:, :, :2] - if additional_output.size > 0 - else np.array([]) - ) - keypoints_scores = ( - additional_output[:, :, 2] - if additional_output.size > 0 - else np.array([]) - ) - - return create_detection_message( - bboxes=np.array(bboxes), - scores=np.array(scores), - labels=np.array(labels), - label_names=label_names, - keypoints=keypoints, - keypoints_scores=keypoints_scores, - keypoint_label_names=keypoint_label_names, - keypoint_edges=keypoint_edges, - ) diff --git a/luxonis_eval/parsers/segmentation.py b/luxonis_eval/parsers/segmentation.py new file mode 100644 index 0000000..837fd45 --- /dev/null +++ b/luxonis_eval/parsers/segmentation.py @@ -0,0 +1,39 @@ +from typing import Any + +import depthai as dai +import numpy as np +from depthai_nodes.message.creators import create_segmentation_message +from depthai_nodes.node.parsers.segmentation import ( + SegmentationParser as DepthAINodesSegmentationParser, +) + +from luxonis_eval.engines.base_engine import ModelSpec +from luxonis_eval.engines.io import EngineOutput + +from .base_parser import BaseParser + + +class SegmentationParser(BaseParser): + """Parser for semantic segmentation model outputs.""" + + def __init__(self, **kwargs: Any) -> None: + """Initialize the segmentation parser.""" + super().__init__(**kwargs) + + def parse( + self, + output: EngineOutput, + model_spec: ModelSpec, + *, + classes_in_one_layer: bool = False, + **kwargs: Any, + ) -> dai.SegmentationMask: + """Parse backend output into segmentation predictions.""" + del model_spec, kwargs + _, segmentation_mask = output.first() + class_map = DepthAINodesSegmentationParser.compute( + np.asarray(segmentation_mask), + classes_in_one_layer=classes_in_one_layer, + ) + + return create_segmentation_message(class_map) diff --git a/luxonis_eval/parsers/semantic_seg.py b/luxonis_eval/parsers/semantic_seg.py deleted file mode 100755 index 7c4640c..0000000 --- a/luxonis_eval/parsers/semantic_seg.py +++ /dev/null @@ -1,98 +0,0 @@ -from typing import Any - -import depthai as dai -import numpy as np -from depthai_nodes import SegmentationMask -from depthai_nodes.message.creators import create_segmentation_message -from loguru import logger - -from .base_parser import BaseParser - - -class SemanticSegmentationParser(BaseParser): - """Parser for semantic segmentation model outputs.""" - - def __init__(self, **kwargs: Any) -> None: - """Initialize the semantic segmentation parser.""" - super().__init__(**kwargs) - - def parse( - self, - raw_output: dai.NNData | list[np.ndarray], - *, - classes_in_one_layer: bool = False, - **kwargs: Any, - ) -> SegmentationMask: - """Parse backend output into detection predictions. - - Parameters - ---------- - raw_output : dai.NNData | list[np.ndarray] - Backend inference output. - classes_in_one_layer : bool, default=False - Whether the model outputs classes in a single layer. - **kwargs : Any - Additional parser arguments. - - Returns - ------- - SegmentationMask - Detection results including boxes, scores, classes, and metadata. - """ - if isinstance(raw_output, dai.NNData): - layer_names = raw_output.getAllLayerNames() - logger.debug(f"Processing output with layers: {layer_names}") - output_name = layer_names[0] - segmentation_mask: np.ndarray = raw_output.getTensor( - output_name, dequantize=True - ) # type: ignore - elif isinstance(raw_output, list): - segmentation_mask: np.ndarray = raw_output[0] - else: - raise TypeError( - f"Unsupported raw_output type: {type(raw_output)}. Expected dai.NNData or list[np.ndarray]." - ) - - if len(segmentation_mask.shape) == 4: - segmentation_mask = segmentation_mask[0] - - if len(segmentation_mask.shape) != 3: - raise ValueError( - f"Expected 3D output tensor, got {len(segmentation_mask.shape)}D." - ) - - np_function = np.argmax - mask_shape = segmentation_mask.shape - min_dim = np.argmin(mask_shape) - if min_dim == len(mask_shape) - 1: - segmentation_mask = segmentation_mask.transpose(2, 0, 1) - adding_unassigned_class = False - if segmentation_mask.shape[0] == 1: # shape is (1, H, W) - if classes_in_one_layer: - np_function = np.max - else: - # If there is only one class, add an unassigned class - adding_unassigned_class = True - segmentation_mask = np.vstack( - ( - np.zeros( - ( - 1, - segmentation_mask.shape[1], - segmentation_mask.shape[2], - ), - dtype=np.float32, - ), - segmentation_mask, - ) - ) - - class_map = ( - np_function(segmentation_mask, axis=0) - .reshape(segmentation_mask.shape[1], segmentation_mask.shape[2]) - .astype(np.int16) - ) - if adding_unassigned_class: - class_map = class_map - 1 - - return create_segmentation_message(class_map) diff --git a/luxonis_eval/parsers/yolo.py b/luxonis_eval/parsers/yolo.py new file mode 100644 index 0000000..fa5d4b5 --- /dev/null +++ b/luxonis_eval/parsers/yolo.py @@ -0,0 +1,116 @@ +from typing import Any + +import depthai as dai +from depthai_nodes.message.creators import create_detection_message +from depthai_nodes.node.parsers.yolo import ( + YOLOComputeInputs, + YOLOExtendedParser as DepthAINodesYOLOExtendedParser, +) + +from luxonis_eval.engines.base_engine import ModelSpec +from luxonis_eval.engines.io import EngineOutput +from luxonis_eval.utils.depthai_nodes import build_yolo_compute_kwargs + +from .base_parser import BaseParser + + +class YOLOExtendedParser(BaseParser): + """Parser for YOLO-based detection, segmentation, and pose outputs.""" + + _DET_MODE = 0 + _KPTS_MODE = 1 + _SEG_MODE = 2 + + def __init__(self, **kwargs: Any) -> None: + """Initialize the YOLO parser.""" + super().__init__(**kwargs) + + def parse( + self, + output: EngineOutput, + model_spec: ModelSpec, + *, + class_map: dict[int, str], + subtype: str, + n_classes: int | None = None, + anchors: list[list[list[float]]] | None = None, + conf_threshold: float = 0.5, + iou_threshold: float = 0.5, + mask_conf: float = 0.5, + max_det: int = 300, + keypoint_label_names: list[str] | None = None, + keypoint_edges: list[tuple[int, int]] | None = None, + **kwargs: Any, + ) -> dai.ImgDetections: + """Parse backend output into YOLO predictions.""" + del kwargs + payload = DepthAINodesYOLOExtendedParser.compute( + YOLOComputeInputs( + **build_yolo_compute_kwargs( + output, + model_spec=model_spec, + class_map=class_map, + subtype=subtype, + n_classes=n_classes, + anchors=anchors, + conf_threshold=conf_threshold, + iou_threshold=iou_threshold, + max_det=max_det, + mask_conf=mask_conf, + keypoint_label_names=keypoint_label_names, + keypoint_edges=keypoint_edges, + ) + ) + ) + + mode = self._resolve_mode(payload) + if mode == self._KPTS_MODE: + return create_detection_message( + bboxes=payload["bboxes"], + scores=payload["scores"], + labels=payload["labels"], + label_names=payload["label_names"], + keypoints=payload["keypoints"], + keypoints_scores=payload["keypoints_scores"], + keypoint_label_names=payload.get( + "keypoint_label_names", + keypoint_label_names, + ), + keypoint_edges=payload.get( + "keypoint_edges", + keypoint_edges, + ), + ) + + if mode == self._SEG_MODE: + return create_detection_message( + bboxes=payload["bboxes"], + scores=payload["scores"], + labels=payload["labels"], + label_names=payload["label_names"], + masks=payload["masks"], + ) + + return create_detection_message( + bboxes=payload["bboxes"], + scores=payload["scores"], + labels=payload["labels"], + label_names=payload["label_names"], + ) + + def _resolve_mode(self, payload: dict[str, Any]) -> int: + mode = payload.get("mode") + if mode is not None: + return int(mode) + keypoints = payload.get("keypoints") + if keypoints is not None: + keypoints_size = ( + int(keypoints.size) + if hasattr(keypoints, "size") + else len(keypoints) + ) + if keypoints_size > 0: + return self._KPTS_MODE + if payload.get("masks") is not None: + return self._SEG_MODE + return self._DET_MODE diff --git a/luxonis_eval/utils/depthai_nodes.py b/luxonis_eval/utils/depthai_nodes.py new file mode 100644 index 0000000..009e08b --- /dev/null +++ b/luxonis_eval/utils/depthai_nodes.py @@ -0,0 +1,207 @@ +from __future__ import annotations + +from typing import Any + +import numpy as np +from depthai_nodes.node.parsers.utils.yolo import YOLOSubtype + +from luxonis_eval.engines.base_engine import ModelSpec +from luxonis_eval.engines.io import EngineOutput + + +def ordered_class_names(class_map: dict[int, str]) -> list[str]: + """Return class names ordered by class index.""" + if not class_map: + return [] + + ordered_indices = sorted(class_map) + expected_indices = list(range(len(ordered_indices))) + if ordered_indices != expected_indices: + raise ValueError( + "class_map must contain contiguous zero-based indices, got " + f"{ordered_indices}." + ) + + return [class_map[index] for index in ordered_indices] + + +def extract_segmentation_mask(predictions: Any) -> np.ndarray: + """Extract a semantic-segmentation mask from a DepthAI message.""" + if hasattr(predictions, "getCvMask"): + mask = predictions.getCvMask() + elif hasattr(predictions, "getCvSegmentationMask"): + mask = predictions.getCvSegmentationMask() + else: + raise TypeError( + "Unsupported segmentation prediction type " + f"{type(predictions)!r}: expected a DepthAI SegmentationMask " + "message." + ) + + if mask is None: + raise ValueError("Segmentation prediction does not contain a mask.") + + return np.asarray(mask) + + +def build_yolo_compute_kwargs( + output: EngineOutput, + model_spec: ModelSpec, + *, + class_map: dict[int, str], + subtype: str, + n_classes: int | None = None, + anchors: list[list[list[float]]] | None = None, + conf_threshold: float, + iou_threshold: float, + max_det: int, + mask_conf: float = 0.5, + keypoint_label_names: list[str] | None = None, + keypoint_edges: list[tuple[int, int]] | None = None, +) -> dict[str, Any]: + """Adapter that converts EngineOutput + ModelSpec into the field mapping + required to construct ``depthai_nodes`` ``YOLOComputeInputs``.""" + try: + subtype_enum = YOLOSubtype(subtype.lower()) + except ValueError as err: + raise ValueError( + f"Invalid YOLO subtype {subtype}. Supported YOLO subtypes are " + f"{[e.value for e in YOLOSubtype][:-1]}." + ) from err + + layer_names = list(output.names()) + outputs_values: list[np.ndarray] + kpts_outputs: list[np.ndarray] | None = None + masks_outputs_values: list[np.ndarray] | None = None + protos_output: np.ndarray | None = None + protos_len: int | None = None + v26_mask_coeffs: np.ndarray | None = None + v26_protos: np.ndarray | None = None + v26_pose_kpts: np.ndarray | None = None + + if subtype_enum == YOLOSubtype.V26: + if any("output_masks" in name for name in layer_names): + outputs_values = [ + output.get("output_yolo26").astype(np.float32, copy=False) + ] + mask_name = next( + name for name in layer_names if "output_masks" in name + ) + protos_name = next( + ( + name + for name in layer_names + if "protos" in name + ), + "protos_output", + ) + v26_mask_coeffs = output.get(mask_name).astype( + np.float32, copy=False + ) + v26_protos = output.get( + protos_name, layout="NCHW" + ).astype(np.float32, copy=False)[0] + elif any("kpt_output" in name for name in layer_names): + outputs_values = [ + output.get("output_yolo26").astype(np.float32, copy=False) + ] + kpt_name = next( + name for name in layer_names if "kpt_output" in name + ) + v26_pose_kpts = output.get(kpt_name).astype( + np.float32, copy=False + ) + else: + outputs_values = [ + output.get(name).astype(np.float32, copy=False) + for name in layer_names + ] + resolved_n_classes = n_classes or len(class_map) + else: + outputs_names = sorted( + [name for name in layer_names if "_yolo" in name or "yolo-" in name] + ) or list(layer_names) + outputs_values = [ + output.get(name, layout="NCHW").astype(np.float32, copy=False) + for name in outputs_names + ] + + if ( + any("kpt_output" in name for name in layer_names) + and subtype_enum != YOLOSubtype.P + ): + kpts_output_names = sorted( + [name for name in layer_names if "kpt_output" in name] + ) or layer_names[len(outputs_names) :] + kpts_outputs = [ + output.get(name).astype(np.float32, copy=False) + for name in kpts_output_names + ] + elif ( + any("mask" in name for name in layer_names) + and subtype_enum != YOLOSubtype.P + ): + protos_name = next( + (name for name in layer_names if "protos" in name), + layer_names[-1], + ) + mask_output_names = sorted( + [ + name + for name in layer_names + if "mask" in name and "proto" not in name + ] + ) or layer_names[len(outputs_names) : -1] + masks_outputs_values = [ + output.get(name, layout="NCHW").astype(np.float32, copy=False) + for name in mask_output_names + ] + protos_output = output.get(protos_name, layout="NCHW").astype( + np.float32, copy=False + ) + protos_len = protos_output.shape[1] + + strides = ( + [8, 16, 32] + if subtype_enum + not in [YOLOSubtype.V3UT, YOLOSubtype.V3T, YOLOSubtype.V4T] + else [16, 32] + ) + final_anchors: np.ndarray | None = ( + np.array(anchors).reshape(len(strides), -1) if anchors else None + ) + inferred_n_classes = ( + outputs_values[0].shape[1] - 5 + if final_anchors is None + else (outputs_values[0].shape[1] // final_anchors.shape[0]) - 5 + ) + if n_classes is not None and inferred_n_classes != n_classes: + raise ValueError( + f"The provided number of classes {n_classes} does not match the " + f"model's {inferred_n_classes}." + ) + resolved_n_classes = inferred_n_classes + + return { + "subtype": subtype_enum, + "layer_names": layer_names, + "outputs_values": outputs_values, + "conf_threshold": conf_threshold, + "n_classes": resolved_n_classes, + "iou_threshold": iou_threshold, + "max_det": max_det, + "anchors": anchors, + "n_keypoints": kpts_outputs[0].shape[1] // 3 if kpts_outputs else 17, + "label_names": ordered_class_names(class_map), + "keypoint_label_names": keypoint_label_names, + "keypoint_edges": keypoint_edges, + "input_shape": (model_spec.height, model_spec.width), + "kpts_outputs": kpts_outputs, + "masks_outputs_values": masks_outputs_values, + "protos_output": protos_output, + "protos_len": protos_len, + "mask_conf": mask_conf, + "v26_mask_coeffs": v26_mask_coeffs, + "v26_protos": v26_protos, + "v26_pose_kpts": v26_pose_kpts, + }