EngineOutput to normalize engine outputs

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.

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: []

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: []

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: []

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: []

configs/yolov8n_inst_seg_config.yaml

@@ -27,7 +27,7 @@ pipeline:
 
   evaluators:
     - parser:
-        name: YOLOInstanceSegmentationParser
+        name: YOLOExtendedParser
         params:
           subtype: yolov8
           n_classes: 80

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 (

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,
         )

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:

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(

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",
 ]

luxonis_eval/engines/base_engine.py

@@ -1,29 +1,16 @@
+import inspect
 from abc import ABC, abstractmethod
-from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
 
 import numpy as np
 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."""
         ...