Working tests, pre-commit and pyright

Author: dtronmans

luxonis_ml/data/datasets/luxonis_dataset.py

@@ -28,8 +28,9 @@
     SegmentationMaskDirectoryExporter,
     TensorflowCSVExporter,
     VOCExporter,
-    YoloExporter,
-    YOLOFormat,
+    YoloV4Exporter,
+    YoloV6Exporter,
+    YoloV8Exporter,
 )
 from luxonis_ml.data.exporters.exporter_utils import (
     ExporterSpec,
@@ -1528,14 +1529,13 @@ def export(
                 },
             ),
             DatasetType.YOLOV8: ExporterSpec(
-                YoloExporter, {"version": YOLOFormat.V8}
-            ),
-            DatasetType.YOLOV6: ExporterSpec(
-                YoloExporter, {"version": YOLOFormat.V6}
-            ),
-            DatasetType.YOLOV4: ExporterSpec(
-                YoloExporter, {"version": YOLOFormat.V4}
+                YoloV8Exporter,
+                {
+                    "skeletons": getattr(self.metadata, "skeletons", None),
+                },
             ),
+            DatasetType.YOLOV6: ExporterSpec(YoloV6Exporter, {}),
+            DatasetType.YOLOV4: ExporterSpec(YoloV4Exporter, {}),
             DatasetType.DARKNET: ExporterSpec(DarknetExporter, {}),
             DatasetType.CLSDIR: ExporterSpec(
                 ClassificationDirectoryExporter, {}

luxonis_ml/data/exporters/__init__.py

@@ -10,7 +10,9 @@
 )
 from .tensorflow_csv_exporter import TensorflowCSVExporter
 from .voc_exporter import VOCExporter
-from .yolo_exporter import YoloExporter, YOLOFormat
+from .yolov4_exporter import YoloV4Exporter
+from .yolov6_exporter import YoloV6Exporter
+from .yolov8_exporter import YoloV8Exporter
 
 __all__ = [
     "BaseExporter",
@@ -23,6 +25,7 @@
     "SegmentationMaskDirectoryExporter",
     "TensorflowCSVExporter",
     "VOCExporter",
-    "YOLOFormat",
-    "YoloExporter",
+    "YoloV4Exporter",
+    "YoloV6Exporter",
+    "YoloV8Exporter",
 ]

luxonis_ml/data/exporters/base_exporter.py

@@ -44,3 +44,7 @@ def _dump_annotations(
         self, annotations: dict, output_path: Path, part: int | None = None
     ) -> None:
         raise NotImplementedError
+
+    @abstractmethod
+    def supported_ann_types(self) -> list[str]:
+        raise NotImplementedError

luxonis_ml/data/exporters/classification_directory_exporter.py

@@ -11,8 +11,14 @@ class ClassificationDirectoryExporter(BaseExporter):
     def get_split_names(self) -> dict[str, str]:
         return {"train": "train", "val": "valid", "test": "test"}
 
+    def supported_ann_types(self) -> list[str]:
+        return ["classification"]
+
     def transform(self, prepared_ldf: PreparedLDF) -> None:
         ExporterUtils.check_group_file_correspondence(prepared_ldf)
+        ExporterUtils.exporter_specific_annotation_warning(
+            prepared_ldf, self.supported_ann_types()
+        )
 
         grouped = prepared_ldf.processed_df.group_by(
             ["file", "group_id"], maintain_order=True

luxonis_ml/data/exporters/coco_exporter.py

@@ -24,7 +24,7 @@ def __init__(
         max_partition_size_gb: float | None,
         format: COCOFormat = COCOFormat.ROBOFLOW,
         *,
-        skeletons: dict[str, Any] | list[dict[str, Any]] | None = None,
+        skeletons: dict[str, Any] | None = None,
     ):
         super().__init__(
             dataset_identifier, output_path, max_partition_size_gb
@@ -56,14 +56,21 @@ def get_split_names(self) -> dict[str, str]:
             return {"train": "train", "val": "valid", "test": "test"}
         return {"train": "train", "val": "validation", "test": "test"}
 
+    def supported_ann_types(self) -> list[str]:
+        return ["boundingbox", "instance_segmentation", "keypoints"]
+
     def transform(self, prepared_ldf: PreparedLDF) -> None:
+        ExporterUtils.check_group_file_correspondence(prepared_ldf)
+        ExporterUtils.exporter_specific_annotation_warning(
+            prepared_ldf, self.supported_ann_types()
+        )
+
         splits = self.get_split_names()
         annotation_splits: dict[str, dict[str, Any]] = {
             s: {"images": [], "categories": [], "annotations": []}
             for s in splits
         }
         ann_id_counter: dict[str, int] = {s: 1 for s in splits}
-        ExporterUtils.check_group_file_correspondence(prepared_ldf)
 
         grouped = prepared_ldf.processed_df.group_by(
             ["file", "instance_id", "group_id"], maintain_order=True
@@ -262,8 +269,8 @@ def _fill_instance_segmentation(
                 "size": [H, W],
                 "counts": counts.encode("utf-8"),
             }
-            area = float(maskUtils.area(rle_runtime))
-            bbox = maskUtils.toBbox(rle_runtime).tolist()  # [x,y,w,h]
+            area = float(maskUtils.area(rle_runtime))  # type: ignore[arg-type]
+            bbox = maskUtils.toBbox(rle_runtime).tolist()  # type: ignore[arg-type]
             ann["area"] = area
             if not ann.get("bbox"):
                 ann["bbox"] = bbox

luxonis_ml/data/exporters/createml_exporter.py

@@ -25,8 +25,14 @@ def __init__(
     def get_split_names() -> dict[str, str]:
         return {"train": "train", "val": "valid", "test": "test"}
 
+    def supported_ann_types(self) -> list[str]:
+        return ["boundingbox"]
+
     def transform(self, prepared_ldf: PreparedLDF) -> None:
         ExporterUtils.check_group_file_correspondence(prepared_ldf)
+        ExporterUtils.exporter_specific_annotation_warning(
+            prepared_ldf, self.supported_ann_types()
+        )
 
         anns_by_split: dict[str, dict[str, list[dict[str, Any]]]] = {
             v: {} for v in self.get_split_names().values()

luxonis_ml/data/exporters/darknet_exporter.py

@@ -25,8 +25,14 @@ def __init__(
     def get_split_names() -> dict[str, str]:
         return {"train": "train", "val": "valid", "test": "test"}
 
+    def supported_ann_types(self) -> list[str]:
+        return ["boundingbox"]
+
     def transform(self, prepared_ldf: PreparedLDF) -> None:
         ExporterUtils.check_group_file_correspondence(prepared_ldf)
+        ExporterUtils.exporter_specific_annotation_warning(
+            prepared_ldf, self.supported_ann_types()
+        )
 
         labels_by_split: dict[str, dict[str, list[str]]] = {
             v: {} for v in self.get_split_names().values()

luxonis_ml/data/exporters/exporter_utils.py

@@ -4,7 +4,9 @@
 from typing import TYPE_CHECKING, Any
 
 import cv2
+import numpy as np
 import polars as pl
+from loguru import logger
 from pycocotools import mask
 
 if TYPE_CHECKING:
@@ -99,6 +101,29 @@ def check_group_file_correspondence(prepared_ldf: PreparedLDF) -> None:
             f"To export multiple files (e.g. RGB, depth) per instance, export to Native format"
         )
 
+    @staticmethod
+    def exporter_specific_annotation_warning(
+        prepared_ldf: PreparedLDF, supported_ann_types: list[str]
+    ) -> None:
+        df = prepared_ldf.processed_df
+
+        present_task_types = (
+            df.select(pl.col("task_type").drop_nans().drop_nulls().unique())
+            .to_series()
+            .to_list()
+        )
+
+        unsupported = [
+            name
+            for name in present_task_types
+            if name not in supported_ann_types
+        ]
+
+        for name in unsupported:
+            logger.warning(
+                f"Found unsupported annotation type '{name}'; skipping this annotation type."
+            )
+
     @staticmethod
     def split_of_group(prepared_ldf: PreparedLDF, group_id: Any) -> str:
         split = next(
@@ -137,7 +162,7 @@ def create_zip_output(
 
     @staticmethod
     def get_single_skeleton(
-        allow_keypoints: bool, skeletons: dict
+        allow_keypoints: bool, skeletons: dict[str, Any] | None = None
     ) -> tuple[list[str], list[list[int]]]:
         """Returns (labels, skeleton_edges_1_based) for the single
         skeleton.
@@ -178,6 +203,45 @@ def rle_to_yolo_polygon(rle: str, height: int, width: int) -> list:
 
         return polygons
 
+    @staticmethod
+    def _bbox_from_poly(
+        coords: list[float],
+    ) -> tuple[float, float, float, float]:
+        xs = coords[0::2]
+        ys = coords[1::2]
+        x_min, x_max = min(xs), max(xs)
+        y_min, y_max = min(ys), max(ys)
+        return x_min, y_min, (x_max - x_min), (y_max - y_min)
+
+    @staticmethod
+    def _iou_xywh(
+        a: tuple[float, float, float, float],
+        b: tuple[float, float, float, float],
+    ) -> float:
+        ax, ay, aw, ah = a
+        bx, by, bw, bh = b
+        ax2, ay2 = ax + aw, ay + ah
+        bx2, by2 = bx + bw, by + bh
+        inter_w = max(0.0, min(ax2, bx2) - max(ax, bx))
+        inter_h = max(0.0, min(ay2, by2) - max(ay, by))
+        inter = inter_w * inter_h
+        if inter <= 0.0:
+            return 0.0
+        union = aw * ah + bw * bh - inter
+        return inter / union if union > 0.0 else 0.0
+
+    @staticmethod
+    def decode_rle_with_pycoco(ann: dict[str, Any]) -> np.ndarray:
+        h = int(ann["height"])
+        w = int(ann["width"])
+        counts = ann["counts"]
+
+        # pycocotools expects an RLE object with 'size' and 'counts'
+        rle = {"size": [h, w], "counts": counts.encode("utf-8")}
+
+        m = mask.decode(rle)  # type: ignore[arg-type]
+        return np.array(m, dtype=np.uint8, order="C")
+
     def _normalize(
         self, xs: list[float], ys: list[float], w: float, h: float
     ) -> list[float]:

luxonis_ml/data/exporters/native_exporter.py

@@ -19,14 +19,19 @@ class NativeExporter(BaseExporter):
     def get_split_names() -> dict[str, str]:
         return {"train": "train", "val": "val", "test": "test"}
 
-    SUPPORTED_ANN_TYPES: set[str] = {
-        "boundingbox",
-        "segmentation",
-        "keypoints",
-        "instance_segmentation",
-    }
+    def supported_ann_types(self) -> list[str]:
+        return [
+            "boundingbox",
+            "segmentation",
+            "keypoints",
+            "instance_segmentation",
+        ]
 
     def transform(self, prepared_ldf: PreparedLDF) -> None:
+        ExporterUtils.exporter_specific_annotation_warning(
+            prepared_ldf, self.supported_ann_types()
+        )
+
         annotation_splits: dict[str, list[dict[str, Any]]] = {
             k: [] for k in self.get_split_names()
         }
@@ -125,7 +130,7 @@ def _process_row(
                 "instance_id": instance_id,
                 "class": class_name,
             }
-            if task_type in self.SUPPORTED_ANN_TYPES:
+            if task_type in self.supported_ann_types():
                 ann[task_type] = data
             elif task_type.startswith("metadata/"):
                 ann["metadata"] = {task_type[9:]: data}