diff --git a/README.md b/README.md index 1e902f8..61081f6 100644 --- a/README.md +++ b/README.md @@ -297,7 +297,7 @@ pipeline: mean: [0.485, 0.456, 0.406] std: [0.229, 0.224, 0.225] color_space: RGB # RGB | BGR | GRAY - keep_aspect_ratio: false # Preserve aspect ratio during resize + keep_aspect_ratio: true # Preserve aspect ratio during resize ``` > [!NOTE] diff --git a/luxonis_eval/config/resolved.py b/luxonis_eval/config/resolved.py index 29eaa73..2d35eeb 100644 --- a/luxonis_eval/config/resolved.py +++ b/luxonis_eval/config/resolved.py @@ -44,7 +44,7 @@ class PreProcessingConfig(BaseModelExtraForbid): default_factory=NormalizeAugmentationConfig ) color_space: Literal["RGB", "BGR", "GRAY"] = "RGB" - keep_aspect_ratio: bool = False + keep_aspect_ratio: bool = True class DataLoaderConfig(BaseModelExtraForbid): diff --git a/luxonis_eval/config/resolver.py b/luxonis_eval/config/resolver.py index 7c36af0..d186fa3 100644 --- a/luxonis_eval/config/resolver.py +++ b/luxonis_eval/config/resolver.py @@ -76,6 +76,7 @@ def _resolve_loader( prefer_archive=prefer_archive, default="RGB", ) + keep_aspect_ratio = self._resolve_keep_aspect_ratio(source_loader) normalize = self._resolve_normalize_config( loader_name=source_loader.name, source_normalize=source_loader.preprocessing.normalize, @@ -88,10 +89,22 @@ def _resolve_loader( preprocessing=PreProcessingConfig( normalize=normalize, color_space=color_space, - keep_aspect_ratio=source_loader.preprocessing.keep_aspect_ratio, + keep_aspect_ratio=keep_aspect_ratio, ), ) + def _resolve_keep_aspect_ratio( + self, source_loader: SourceDataLoaderConfig + ) -> bool: + keep_aspect_ratio = source_loader.preprocessing.keep_aspect_ratio + if keep_aspect_ratio is not None: + return keep_aspect_ratio + + if source_loader.name == "LuxonisLoader": + return True + + return False + def _resolve_normalize_config( self, *, @@ -177,17 +190,8 @@ def _resolve_archive_normalize_state( "NNArchive preprocessing normalization must define both " "`mean` and `scale`, or neither." ) - if self._is_identity_normalization(mean, scale): - return False, None return True, {"mean": list(mean), "std": list(scale)} - def _is_identity_normalization( - self, mean: list[float], scale: list[float] - ) -> bool: - return all(value == 0 for value in mean) and all( - value == 1 for value in scale - ) - def _resolve_evaluators( self, source_evaluators: list[SourceEvaluatorConfig] | None, @@ -315,6 +319,7 @@ def construct_yolo_parser_config(head: Any) -> ParserConfig: "anchors", "conf_threshold", "iou_threshold", + "mask_conf", "max_det", ): if key in metadata: diff --git a/luxonis_eval/config/source.py b/luxonis_eval/config/source.py index 6a8a531..d82af57 100644 --- a/luxonis_eval/config/source.py +++ b/luxonis_eval/config/source.py @@ -32,7 +32,7 @@ def validate_params(cls, value: Params | None) -> Params | None: class SourcePreProcessingConfig(BaseModelExtraForbid): normalize: SourceNormalizeAugmentationConfig | None = None color_space: Literal["RGB", "BGR", "GRAY"] | None = None - keep_aspect_ratio: bool = False + keep_aspect_ratio: bool | None = None class SourceDataLoaderConfig(BaseModelExtraForbid): diff --git a/luxonis_eval/core/core.py b/luxonis_eval/core/core.py index 1d7cd65..eee21a3 100644 --- a/luxonis_eval/core/core.py +++ b/luxonis_eval/core/core.py @@ -36,6 +36,7 @@ from luxonis_eval.metrics import ThroughputMetric from luxonis_eval.metrics.base_metric import BaseMetric from luxonis_eval.parsers.base_parser import BaseParser +from luxonis_eval.parsers.yolo import clear_prediction_metadata from luxonis_eval.visualizers.base_visualizer import BaseVisualizer @@ -176,39 +177,43 @@ def _run_evaluators(self) -> dict[str, Any]: ) parsing_elapsed = time.perf_counter() - parsing_t0 - metric_update_t0 = time.perf_counter() - for metric, metric_ctx in zip( - self.metrics, self.metric_contexts, strict=True - ): - metric.update( - predictions=predictions, - target=target, - **metric_ctx, + try: + metric_update_t0 = time.perf_counter() + for metric, metric_ctx in zip( + self.metrics, self.metric_contexts, strict=True + ): + metric.update( + predictions=predictions, + target=target, + **metric_ctx, + ) + metric_update_elapsed = ( + time.perf_counter() - metric_update_t0 ) - metric_update_elapsed = time.perf_counter() - metric_update_t0 - self.throughput_metric.update( - inference=inference_elapsed, - parsing=parsing_elapsed, - metric_update=metric_update_elapsed, - ) - - active_visualizer_cfgs = [ - visualizer_cfg - for visualizer_cfg in self.evaluator_cfg.visualizers - if visualizer_cfg.active - ] - for visualizer, visualizer_cfg in zip( - self.visualizers, - active_visualizer_cfgs, - strict=True, - ): - visualizer.visualize( - predictions, - self.engine.vis_frame(), - **visualizer_cfg.params, + self.throughput_metric.update( + inference=inference_elapsed, + parsing=parsing_elapsed, + metric_update=metric_update_elapsed, ) + active_visualizer_cfgs = [ + visualizer_cfg + for visualizer_cfg in self.evaluator_cfg.visualizers + if visualizer_cfg.active + ] + for visualizer, visualizer_cfg in zip( + self.visualizers, + active_visualizer_cfgs, + strict=True, + ): + visualizer.visualize( + predictions, + self.engine.vis_frame(), + **visualizer_cfg.params, + ) + finally: + clear_prediction_metadata(predictions) progress.update(advance=1) metric_compute_t0 = time.perf_counter() @@ -287,24 +292,27 @@ def _sanity_check_pipeline(self) -> None: **self.evaluator_cfg.parser.params, ) - for metric, metric_ctx in zip( - self.metrics, self.metric_contexts, strict=True - ): - missing = set(metric.required_target_keys()) - set(target) - if missing: - raise ValueError( - "Target is missing required keys for " - f"{metric.__class__.__name__}: {sorted(missing)}. " - f"Got keys: {sorted(target.keys())}." - ) + try: + for metric, metric_ctx in zip( + self.metrics, self.metric_contexts, strict=True + ): + missing = set(metric.required_target_keys()) - set(target) + if missing: + raise ValueError( + "Target is missing required keys for " + f"{metric.__class__.__name__}: {sorted(missing)}. " + f"Got keys: {sorted(target.keys())}." + ) - metric.update( - predictions=predictions, - target=target, - **metric_ctx, - ) - metric.compute() - metric.reset() + metric.update( + predictions=predictions, + target=target, + **metric_ctx, + ) + metric.compute() + metric.reset() + finally: + clear_prediction_metadata(predictions) def _clear_runtime_fields(self) -> None: self.engine: BaseEngine | None = None diff --git a/luxonis_eval/core/factories.py b/luxonis_eval/core/factories.py index 010e7b6..b2f96e0 100644 --- a/luxonis_eval/core/factories.py +++ b/luxonis_eval/core/factories.py @@ -190,13 +190,24 @@ def _create_luxonis_loader( loader_params["filter_task_names"] = [loader_task_name] augmentation_config = [] - if cfg.pipeline.loader.preprocessing.normalize.active: + if ( + cfg.pipeline.loader.preprocessing.normalize.active + and cfg.pipeline.engine.name != "depthai" + ): augmentation_config.append( { "name": "Normalize", "params": cfg.pipeline.loader.preprocessing.normalize.params, } ) + elif ( + cfg.pipeline.loader.preprocessing.normalize.active + and cfg.pipeline.engine.name == "depthai" + ): + logger.info( + "Skipping host-side loader normalization for the DepthAI backend. " + "NNArchive preprocessing is expected to run on-device." + ) dataloader = LuxonisLoader( dataset, diff --git a/luxonis_eval/engines/onnx_engine.py b/luxonis_eval/engines/onnx_engine.py index 1a82088..142da62 100644 --- a/luxonis_eval/engines/onnx_engine.py +++ b/luxonis_eval/engines/onnx_engine.py @@ -28,7 +28,7 @@ def get( ) -> np.ndarray: del layout try: - return self.tensors[name] + return self.tensors[name].copy() except KeyError as err: raise ValueError( f"Requested output tensor {name!r} is not available. " diff --git a/luxonis_eval/metrics/__init__.py b/luxonis_eval/metrics/__init__.py index 61f34cd..410be8d 100755 --- a/luxonis_eval/metrics/__init__.py +++ b/luxonis_eval/metrics/__init__.py @@ -1,6 +1,8 @@ from .base_metric import BaseMetric from .bbox_map import BboxMeanAveragePrecision from .dice_coef import DiceCoefficient +from .f1_score import F1Score +from .jaccard_index import JaccardIndex from .keypoint_map import KeypointMeanAveragePrecision from .mask_map import MaskMeanAveragePrecision from .mIoU import MIoU @@ -11,6 +13,8 @@ "BaseMetric", "BboxMeanAveragePrecision", "DiceCoefficient", + "F1Score", + "JaccardIndex", "KeypointMeanAveragePrecision", "MIoU", "MaskMeanAveragePrecision", diff --git a/luxonis_eval/metrics/f1_score.py b/luxonis_eval/metrics/f1_score.py new file mode 100644 index 0000000..9249b21 --- /dev/null +++ b/luxonis_eval/metrics/f1_score.py @@ -0,0 +1,82 @@ +from typing import Any, Literal + +import depthai as dai +import numpy as np + +from luxonis_eval.metrics.base_metric import BaseMetric +from luxonis_eval.metrics.metrics_utils import ( + binary_segmentation_confusion, + normalize_prediction_segmentation_mask, + target_segmentation_to_index_mask, +) +from luxonis_eval.utils.depthai_nodes import extract_segmentation_mask + + +class F1Score(BaseMetric): + """F1 score for binary semantic segmentation.""" + + def __init__( + self, + num_classes: int | None = None, + include_background: bool = True, + 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( + "`F1Score` behavior is only implemented 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/jaccard_index.py b/luxonis_eval/metrics/jaccard_index.py new file mode 100644 index 0000000..342c692 --- /dev/null +++ b/luxonis_eval/metrics/jaccard_index.py @@ -0,0 +1,81 @@ +from typing import Any, Literal + +import depthai as dai +import numpy as np + +from luxonis_eval.metrics.base_metric import BaseMetric +from luxonis_eval.metrics.metrics_utils import ( + binary_segmentation_confusion, + normalize_prediction_segmentation_mask, + target_segmentation_to_index_mask, +) +from luxonis_eval.utils.depthai_nodes import extract_segmentation_mask + + +class JaccardIndex(BaseMetric): + """Jaccard index for binary segmentation.""" + + def __init__( + self, + num_classes: int | None = None, + include_background: bool = True, + 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( + "`JaccardIndex` behavior is only implemeted 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/mask_map.py b/luxonis_eval/metrics/mask_map.py index 35ad8b6..d7adb4b 100644 --- a/luxonis_eval/metrics/mask_map.py +++ b/luxonis_eval/metrics/mask_map.py @@ -3,14 +3,14 @@ import depthai as dai import numpy as np +import torch +from torchmetrics.detection import MeanAveragePrecision from luxonis_eval.metrics.base_metric import BaseMetric from luxonis_eval.metrics.metrics_utils import ( - area_from_rle, - bbox_from_rle, - binary_mask_to_rle, + detection_to_coco_xywh, ) -from luxonis_eval.utils.coco_utils import COCOStore +from luxonis_eval.parsers.yolo import get_prediction_instance_masks class MaskMeanAveragePrecision(BaseMetric): @@ -29,7 +29,11 @@ def __init__(self, iou_type: str = "segm", **kwargs: Any) -> None: **kwargs : Any Additional metric configuration. """ - self._store = COCOStore(iou_type=iou_type) + if iou_type != "segm": + raise ValueError( + "MaskMeanAveragePrecision is fixed to instance-segmentation " + "evaluation and only supports iou_type='segm'." + ) super().__init__(**kwargs) def required_target_keys(self) -> list[str]: @@ -44,7 +48,10 @@ def required_target_keys(self) -> list[str]: def reset(self) -> None: """Reset internal metric state.""" - self._store.reset() + self.metric = MeanAveragePrecision( + iou_type=("bbox", "segm"), + backend="faster_coco_eval", + ) def update( self, @@ -69,83 +76,78 @@ def update( width = int(kwargs["width"]) height = int(kwargs["height"]) - class_map: dict[int, str] = kwargs.get("class_map", {}) category_ids: Sequence[int] | None = kwargs.get("category_ids") class_index_map = kwargs.get("class_index_map") + target_converter = kwargs.get("target_converter") - self._store.init_categories_once( - class_map=class_map, category_ids=category_ids + target_classes, target_boxes_xywh = target_converter( + target_boxes, width, height ) - img_id = self._store.new_image(width=width, height=height) - - # --- GT --- - target_classes = target_boxes[:, 0].astype(np.int64) - - for mask, cls in zip(target_masks, target_classes, strict=True): - cls = int(cls) - if class_index_map is not None: - cls = int(class_index_map[cls]) - if ( - self._store.category_ids_set is not None - and cls not in self._store.category_ids_set - ): - continue - - rle = binary_mask_to_rle(mask.astype(bool)) - area = area_from_rle(rle) - coco_bbox = bbox_from_rle(rle) - - self._store.add_gt( - { - "image_id": img_id, - "category_id": cls, - "segmentation": rle, - "bbox": coco_bbox, - "area": area, - "iscrowd": 0, - } + if class_index_map is not None: + target_classes = np.array( + [class_index_map[int(cls)] for cls in target_classes], + dtype=np.int64, ) + category_ids_set = ( + {int(category_id) for category_id in category_ids} + if category_ids is not None + else None + ) - # --- Predictions --- detections = predictions.detections - scores = [det.confidence for det in detections] - classes = [det.label for det in detections] - masks: np.ndarray = predictions.getCvSegmentationMask() # type: ignore - # Unflatten (N*H, W) → (N, H, W) - masks = ( - masks.reshape(len(detections), -1, masks.shape[-1]) - if detections - else np.zeros((0, 0, 0), dtype=np.uint8) + masks = self._resolve_prediction_instance_masks( + get_prediction_instance_masks(predictions), + n_detections=len(detections), + height=height, + width=width, ) - if masks.size > 0 and masks.shape[1:] != (height, width): - raise ValueError( - f"Mask dimensions {masks.shape[1:]} do not match image dimensions ({height}, {width}). " - f"Ensure masks in 'dai.ImgDetections' are stored as (N*H, W)." - ) - for mask, score, cls in zip(masks, scores, classes, strict=True): - cls = int(cls) - if ( - self._store.category_ids_set is not None - and cls not in self._store.category_ids_set - ): + pred_boxes_xyxy: list[list[float]] = [] + pred_scores: list[float] = [] + pred_classes: list[int] = [] + pred_masks: list[np.ndarray] = [] + for det, mask in zip(detections, masks, strict=True): + cls = int(det.label) + if category_ids_set is not None and cls not in category_ids_set: continue if mask.sum() == 0: continue - rle = binary_mask_to_rle(mask.astype(bool)) - coco_bbox = bbox_from_rle(rle) - - self._store.add_pred( - { - "image_id": img_id, - "category_id": cls, - "segmentation": rle, - "bbox": coco_bbox, - "score": float(score), - } + x, y, w, h = detection_to_coco_xywh(det, width, height) + pred_boxes_xyxy.append([x, y, x + w, y + h]) + pred_scores.append(float(det.confidence)) + pred_classes.append(cls) + pred_masks.append(mask.astype(bool, copy=False)) + + target_boxes_xyxy = self._xywh_to_xyxy(target_boxes_xywh) + if category_ids_set is not None: + keep = np.array( + [int(cls) in category_ids_set for cls in target_classes], + dtype=bool, ) + target_classes = target_classes[keep] + target_boxes_xyxy = target_boxes_xyxy[keep] + target_masks = target_masks[keep] + + preds = [ + { + "boxes": self._as_boxes_tensor(pred_boxes_xyxy), + "scores": self._as_scores_tensor(pred_scores), + "labels": self._as_labels_tensor(pred_classes), + "masks": self._as_masks_tensor(pred_masks, height, width), + } + ] + targets = [ + { + "boxes": self._as_boxes_tensor(target_boxes_xyxy.tolist()), + "labels": self._as_labels_tensor(target_classes.tolist()), + "masks": self._as_target_masks_tensor( + target_masks, height, width + ), + } + ] + self.metric.update(preds, targets) def compute(self) -> dict[str, float]: """Compute final mAP metrics. @@ -155,4 +157,91 @@ def compute(self) -> dict[str, float]: dict[str, float] Computed mAP results. """ - return self._store.evaluate() + metrics = self.metric.compute() + return { + "AP": float(metrics["segm_map"]), + "AP50": float(metrics["segm_map_50"]), + } + + @staticmethod + def _resolve_prediction_instance_masks( + raw_masks: np.ndarray | None, + *, + n_detections: int, + height: int, + width: int, + ) -> np.ndarray: + if raw_masks is None: + raise ValueError( + "MaskMeanAveragePrecision requires raw per-instance masks " + "for the prediction message." + ) + + masks = np.asarray(raw_masks) + if masks.size == 0: + if n_detections != 0: + raise ValueError( + "MaskMeanAveragePrecision received no raw instance masks " + f"for {n_detections} detections." + ) + return np.zeros((0, height, width), dtype=np.uint8) + + if masks.ndim != 3: + raise ValueError( + f"Unsupported raw instance mask rank {masks.ndim}. " + "Expected shape (N, H, W)." + ) + + if masks.shape != (n_detections, height, width): + raise ValueError( + "Raw instance masks do not align with detections. Expected " + f"({n_detections}, {height}, {width}), got {masks.shape}." + ) + + return masks.astype(np.uint8, copy=False) + + @staticmethod + def _xywh_to_xyxy(boxes_xywh: np.ndarray) -> np.ndarray: + boxes_xywh = np.asarray(boxes_xywh, dtype=np.float32) + if boxes_xywh.size == 0: + return np.zeros((0, 4), dtype=np.float32) + + boxes_xyxy = boxes_xywh.copy() + boxes_xyxy[:, 2] = boxes_xyxy[:, 0] + boxes_xyxy[:, 2] + boxes_xyxy[:, 3] = boxes_xyxy[:, 1] + boxes_xyxy[:, 3] + return boxes_xyxy + + @staticmethod + def _as_boxes_tensor(boxes: list[list[float]]) -> torch.Tensor: + if not boxes: + return torch.zeros((0, 4), dtype=torch.float32) + return torch.tensor(boxes, dtype=torch.float32) + + @staticmethod + def _as_scores_tensor(scores: list[float]) -> torch.Tensor: + if not scores: + return torch.zeros((0,), dtype=torch.float32) + return torch.tensor(scores, dtype=torch.float32) + + @staticmethod + def _as_labels_tensor(labels: list[int]) -> torch.Tensor: + if not labels: + return torch.zeros((0,), dtype=torch.int64) + return torch.tensor(labels, dtype=torch.int64) + + @staticmethod + def _as_masks_tensor( + masks: list[np.ndarray], height: int, width: int + ) -> torch.Tensor: + if not masks: + return torch.zeros((0, height, width), dtype=torch.bool) + return torch.from_numpy(np.stack(masks, axis=0)).to(torch.bool) + + @staticmethod + def _as_target_masks_tensor( + masks: np.ndarray, height: int, width: int + ) -> torch.Tensor: + masks = np.asarray(masks) + if masks.size == 0: + return torch.zeros((0, height, width), dtype=torch.bool) + return torch.from_numpy(masks.astype(bool, copy=False)) diff --git a/luxonis_eval/metrics/metrics_utils.py b/luxonis_eval/metrics/metrics_utils.py index 5e3c35d..f2b75aa 100644 --- a/luxonis_eval/metrics/metrics_utils.py +++ b/luxonis_eval/metrics/metrics_utils.py @@ -146,8 +146,6 @@ def normalize_prediction_segmentation_mask( 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 @@ -183,6 +181,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/yolo.py b/luxonis_eval/parsers/yolo.py index bc9a34d..7d3f989 100644 --- a/luxonis_eval/parsers/yolo.py +++ b/luxonis_eval/parsers/yolo.py @@ -1,21 +1,41 @@ 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.yolo import ( - YOLOComputeInputs, -) from depthai_nodes.node.parsers.yolo import ( 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 luxonis_eval.utils.depthai_nodes import ( + build_yolo_compute_inputs, + build_yolo_instance_masks, +) from .base_parser import BaseParser +prediction_instance_masks_by_id: dict[int, np.ndarray] = {} + + +def store_prediction_instance_masks( + predictions: dai.ImgDetections, instance_masks: np.ndarray +) -> None: + prediction_instance_masks_by_id[id(predictions)] = instance_masks + + +def get_prediction_instance_masks( + predictions: dai.ImgDetections, +) -> np.ndarray | None: + return prediction_instance_masks_by_id.get(id(predictions)) + + +def clear_prediction_metadata(predictions: Any) -> None: + prediction_instance_masks_by_id.pop(id(predictions), None) + + class YOLOExtendedParser(BaseParser): """Parser for YOLO-based detection, segmentation, and pose outputs.""" @@ -37,6 +57,7 @@ def parse( subtype: str, n_classes: int | None = None, anchors: list[list[list[float]]] | None = None, + strides: list[int] | None = None, conf_threshold: float = 0.5, iou_threshold: float = 0.5, mask_conf: float = 0.5, @@ -47,26 +68,24 @@ def parse( ) -> 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, - ) - ) + compute_inputs = build_yolo_compute_inputs( + output, + model_spec=model_spec, + class_map=class_map, + subtype=subtype, + n_classes=n_classes, + anchors=anchors, + strides=strides, + 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, ) + payload = DepthAINodesYOLOExtendedParser.compute(compute_inputs) - mode = self._resolve_mode(payload) + mode = int(payload["mode"]) if mode == self._KPTS_MODE: return create_detection_message( bboxes=payload["bboxes"], @@ -75,24 +94,43 @@ def parse( 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, - ), + keypoint_label_names=payload["keypoint_label_names"], + keypoint_edges=payload["keypoint_edges"], ) if mode == self._SEG_MODE: - return create_detection_message( + message = create_detection_message( bboxes=payload["bboxes"], scores=payload["scores"], labels=payload["labels"], label_names=payload["label_names"], masks=payload["masks"], ) + instance_masks = build_yolo_instance_masks( + build_yolo_compute_inputs( + output, + model_spec=model_spec, + class_map=class_map, + subtype=subtype, + n_classes=n_classes, + anchors=anchors, + strides=strides, + 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, + ) + ) + if instance_masks.shape[0] != len(message.detections): + raise ValueError( + "YOLOExtendedParser received mismatched segmentation outputs: " + f"{len(message.detections)} detections but " + f"{instance_masks.shape[0]} rebuilt instance masks." + ) + store_prediction_instance_masks(message, instance_masks) + return message return create_detection_message( bboxes=payload["bboxes"], @@ -100,20 +138,3 @@ def parse( 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 index 9daa7be..fb3f88b 100644 --- a/luxonis_eval/utils/depthai_nodes.py +++ b/luxonis_eval/utils/depthai_nodes.py @@ -3,7 +3,14 @@ from typing import Any import numpy as np -from depthai_nodes.node.parsers.utils.yolo import YOLOSubtype +import torch +import torch.nn.functional as F +from depthai_nodes.node.parsers.yolo import YOLOComputeInputs +from depthai_nodes.node.parsers.utils.yolo import ( + YOLOSubtype, + decode_yolo26, + decode_yolo_output, +) from luxonis_eval.engines.base_engine import ModelSpec from luxonis_eval.engines.io import EngineOutput @@ -44,7 +51,7 @@ def extract_segmentation_mask(predictions: Any) -> np.ndarray: return np.asarray(mask) -def build_yolo_compute_kwargs( +def build_yolo_compute_inputs( output: EngineOutput, model_spec: ModelSpec, *, @@ -52,13 +59,14 @@ def build_yolo_compute_kwargs( subtype: str, n_classes: int | None = None, anchors: list[list[list[float]]] | None = None, + strides: list[int] | 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]: +) -> YOLOComputeInputs: """Adapter that converts EngineOutput + ModelSpec into the field mapping required to construct ``depthai_nodes`` ``YOLOComputeInputs``.""" @@ -164,14 +172,16 @@ def build_yolo_compute_kwargs( ) protos_len = protos_output.shape[1] - strides = ( + resolved_strides = strides or ( [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 + np.array(anchors).reshape(len(resolved_strides), -1) + if anchors + else None ) inferred_n_classes = ( outputs_values[0].shape[1] - 5 @@ -185,26 +195,197 @@ def build_yolo_compute_kwargs( ) 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, - } + return YOLOComputeInputs( + subtype=subtype_enum, + layer_names=layer_names, + outputs_values=outputs_values, + strides=strides, + 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, + ) + + +def build_yolo_instance_masks( + compute_inputs: YOLOComputeInputs, +) -> np.ndarray: + """Rebuild per-instance masks with refinement. + + Detection selection still comes from the DepthAI YOLO decode path. + This helper only regenerates instance masks from the kept detections. + + Mirrored LuxonisTrain's + ``luxonis_train.nodes.heads.precision_seg_bbox_head.refine_and_apply_masks()``: + Eval uses the same prototype-combination, bbox-cropping, and + upsampling behavior. + """ + + subtype = compute_inputs.subtype + input_shape = compute_inputs.input_shape + if input_shape is None: + raise ValueError("YOLO mask rebuilding requires an input shape.") + height, width = input_shape + + if subtype == YOLOSubtype.V26: + results, mask_coeffs = decode_yolo26( + compute_inputs.outputs_values[0], + compute_inputs.conf_threshold, + compute_inputs.max_det, + extra_raw=compute_inputs.v26_mask_coeffs, + ) + if mask_coeffs is None: + raise ValueError( + "YOLO26 instance segmentation requires mask coefficients." + ) + mask_prototypes = compute_inputs.v26_protos + if mask_prototypes is None: + raise ValueError( + "YOLO26 instance segmentation requires prototype masks." + ) + return _refine_instance_masks( + mask_prototypes=mask_prototypes, + mask_coefficients=mask_coeffs, + bounding_boxes=results[:, :4], + height=height, + width=width, + ) + + resolved_strides = compute_inputs.strides + if resolved_strides is None: + resolved_strides = ( + [8, 16, 32] + if subtype not in [YOLOSubtype.V3UT, YOLOSubtype.V3T, YOLOSubtype.V4T] + else [16, 32] + ) + + anchors = compute_inputs.anchors + anchors_array = ( + np.array(anchors).reshape(len(resolved_strides), -1) + if anchors is not None + else None + ) + + results = decode_yolo_output( + compute_inputs.outputs_values, + resolved_strides, + anchors_array, + conf_thres=compute_inputs.conf_threshold, + iou_thres=compute_inputs.iou_threshold, + num_classes=compute_inputs.n_classes, + det_mode=False, + subtype=subtype, + ) + + protos_output = compute_inputs.protos_output + masks_outputs_values = compute_inputs.masks_outputs_values + protos_len = compute_inputs.protos_len + if protos_output is None or masks_outputs_values is None or protos_len is None: + raise ValueError( + "YOLO instance segmentation requires prototype and mask outputs." + ) + + mask_coefficients = [] + for other in results[:, 6:]: + hi, ai, xi, yi = other.astype(int) + mask_coefficients.append( + masks_outputs_values[hi][0, ai * protos_len : (ai + 1) * protos_len, yi, xi] + ) + + if not mask_coefficients: + return np.zeros((0, height, width), dtype=np.uint8) + + return _refine_instance_masks( + mask_prototypes=protos_output[0], + mask_coefficients=np.stack(mask_coefficients, axis=0), + bounding_boxes=results[:, :4], + height=height, + width=width, + ) + + +def _refine_instance_masks( + *, + mask_prototypes: np.ndarray, + mask_coefficients: np.ndarray, + bounding_boxes: np.ndarray, + height: int, + width: int, +) -> np.ndarray: + """NumPy/Torch port of LuxonisTrain's mask refinement step. + + This mirrors + ``luxonis_train.nodes.heads.precision_seg_bbox_head.refine_and_apply_masks()`` + because DepthAI returns a merged instance-id mask, while evaluation needs + the per-instance masks before that merge. + """ + if mask_coefficients.shape[0] == 0 or bounding_boxes.shape[0] == 0: + return np.zeros((0, height, width), dtype=np.uint8) + + prototypes_tensor = torch.as_tensor(mask_prototypes, dtype=torch.float32) + coefficients_tensor = torch.as_tensor( + mask_coefficients, dtype=torch.float32 + ) + boxes_tensor = torch.as_tensor(bounding_boxes, dtype=torch.float32) + + channels, proto_h, proto_w = prototypes_tensor.shape + masks_combined = ( + coefficients_tensor @ prototypes_tensor.view(channels, -1) + ).view(-1, proto_h, proto_w) + + scaled_boxes = boxes_tensor.clone() + scaled_boxes[:, [0, 2]] *= proto_w / width + scaled_boxes[:, [1, 3]] *= proto_h / height + + cropped_masks = _apply_bounding_box_to_masks(masks_combined, scaled_boxes) + upsampled_masks = F.interpolate( + cropped_masks.unsqueeze(0), + size=(height, width), + mode="bilinear", + align_corners=False, + ).squeeze(0) + + return (upsampled_masks > 0).to(torch.uint8).cpu().numpy() + + +def _apply_bounding_box_to_masks( + masks: torch.Tensor, + bounding_boxes: torch.Tensor, +) -> torch.Tensor: + """Mirror LuxonisTrain's bbox mask cropping helper. + + This matches + ``luxonis_train.utils.boundingbox.apply_bounding_box_to_masks()`` so the + rebuilt masks follow the same crop semantics. + """ + _, mask_height, mask_width = masks.shape + left, top, right, bottom = torch.split( + bounding_boxes[:, :, None], 1, dim=1 + ) + width_indices = torch.arange( + mask_width, device=masks.device, dtype=left.dtype + )[None, None, :] + height_indices = torch.arange( + mask_height, device=masks.device, dtype=left.dtype + )[None, :, None] + + return masks * ( + (width_indices >= left) + & (width_indices < right) + & (height_indices >= top) + & (height_indices < bottom) + )