More changes and extensions to existing exporters and tests

Author: dtronmans

luxonis_ml/data/datasets/luxonis_dataset.py

@@ -1544,8 +1544,8 @@ def export(
                 SegmentationMaskDirectoryExporter, {}
             ),
             DatasetType.VOC: ExporterSpec(VOCExporter, {}),
-            DatasetType.CREATEML: CreateMLExporter,
-            DatasetType.TFCSV: TensorflowCSVExporter,
+            DatasetType.CREATEML: ExporterSpec(CreateMLExporter, {}),
+            DatasetType.TFCSV: ExporterSpec(TensorflowCSVExporter, {}),
         }
         spec = EXPORTER_MAP.get(dataset_type)
         if spec is None:

luxonis_ml/data/exporters/base_exporter.py

@@ -2,7 +2,7 @@
 
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from luxonis_ml.data.exporters.exporter_utils import PreparedLDF
@@ -30,9 +30,7 @@ def __init__(
         self.current_size = 0
 
     @abstractmethod
-    def transform(
-        self, prepared_ldf: PreparedLDF
-    ) -> dict[str, list[dict[str, Any]]]:
+    def transform(self, prepared_ldf: PreparedLDF) -> None:
         """Convert the prepared dataset into the exporter's format."""
         raise NotImplementedError
 

luxonis_ml/data/exporters/createml_exporter.py

@@ -11,26 +11,6 @@
 
 
 class CreateMLExporter(BaseExporter):
-    """Exports LDF to CreateML format.
-
-    Output structure:
-
-        output/
-        └── <dataset_identifier>[_partN]/
-            ├── train/
-            │   ├── 0.jpg
-            │   ├── 1.jpg
-            │   └── _annotations.createml.json
-            ├── valid/
-            │   └── ...
-            └── test/
-                └── ...
-
-    Notes:
-    - CreateML uses center-based pixel coordinates.
-    - This exporter converts from normalized TL (x,y,w,h) to pixel centers.
-    """
-
     def __init__(
         self,
         dataset_identifier: str,
@@ -74,44 +54,17 @@ def transform(self, prepared_ldf: PreparedLDF) -> None:
             with Image.open(file_path) as im:
                 width, height = im.size
 
-            per_image_anns: list[dict[str, Any]] = []
-            for row in group_df.iter_rows(named=True):
-                ttype = row.get("task_type")
-                ann_str = row.get("annotation")
-                cname = row.get("class_name")
-
-                if ttype != "boundingbox" or ann_str is None or not cname:
-                    continue
-
-                data = json.loads(ann_str)
-                x_tl = float(data.get("x", 0.0))
-                y_tl = float(data.get("y", 0.0))
-                w = float(data.get("w", 0.0))
-                h = float(data.get("h", 0.0))
-
-                x_px = x_tl * width
-                y_px = y_tl * height
-                w_px = w * width
-                h_px = h * height
-                cx_px = x_px + w_px / 2.0
-                cy_px = y_px + h_px / 2.0
-
-                per_image_anns.append(
-                    {
-                        "label": cname,
-                        "coordinates": {
-                            "x": cx_px,
-                            "y": cy_px,
-                            "width": w_px,
-                            "height": h_px,
-                        },
-                    }
-                )
+            per_image_anns = self._collect_bounding_box_annotations(
+                group_df=group_df, width=width, height=height
+            )
 
             anns_by_split[split_name][new_name] = per_image_anns
 
-            ann_size_est = sum(
-                64 + len(a.get("label", "")) for a in per_image_anns
+            per_image_anns = self._collect_bounding_box_annotations(
+                group_df, width, height
+            )
+            ann_size_est = self._estimate_annotation_bytes(
+                new_name, per_image_anns
             )
             img_size = file_path.stat().st_size
 
@@ -131,9 +84,60 @@ def transform(self, prepared_ldf: PreparedLDF) -> None:
                     dest_img.write_bytes(file_path.read_bytes())
                 self.current_size += img_size
 
-        # Final dump
         self._dump_annotations(anns_by_split, self.output_path, self.part)
 
+    @staticmethod
+    def _estimate_annotation_bytes(
+        img_name: str, anns: list[dict[str, Any]]
+    ) -> int:
+        payload = {"image": img_name, "annotations": anns}
+        return len(
+            (json.dumps(payload, ensure_ascii=False) + "\n").encode("utf-8")
+        )
+
+    def _collect_bounding_box_annotations(
+        self,
+        group_df: Any,
+        width: int,
+        height: int,
+    ) -> list[dict[str, Any]]:
+        per_image_anns: list[dict[str, Any]] = []
+
+        for row in group_df.iter_rows(named=True):
+            ttype = row.get("task_type")
+            ann_str = row.get("annotation")
+            cname = row.get("class_name")
+
+            if ttype != "boundingbox" or ann_str is None or not cname:
+                continue
+
+            data = json.loads(ann_str)
+            x_tl = float(data.get("x", 0.0))
+            y_tl = float(data.get("y", 0.0))
+            w = float(data.get("w", 0.0))
+            h = float(data.get("h", 0.0))
+
+            x_px = x_tl * width
+            y_px = y_tl * height
+            w_px = w * width
+            h_px = h * height
+            cx_px = x_px + w_px / 2.0
+            cy_px = y_px + h_px / 2.0
+
+            per_image_anns.append(
+                {
+                    "label": cname,
+                    "coordinates": {
+                        "x": cx_px,
+                        "y": cy_px,
+                        "width": w_px,
+                        "height": h_px,
+                    },
+                }
+            )
+
+        return per_image_anns
+
     def _maybe_roll_partition(
         self,
         anns_by_split: dict[str, dict[str, list[dict[str, Any]]]],

luxonis_ml/data/exporters/darknet_exporter.py

@@ -49,35 +49,7 @@ def transform(self, prepared_ldf: PreparedLDF) -> None:
             new_name = f"{idx}{file_path.suffix}"
             new_stem = Path(new_name).stem
 
-            label_lines: list[str] = []
-            for row in group_df.iter_rows(named=True):
-                ttype = row.get("task_type")
-                ann_str = row.get("annotation")
-                cname = row.get("class_name")
-
-                if ttype != "boundingbox" or ann_str is None:
-                    continue
-
-                if cname and cname not in self.class_to_id:
-                    self.class_to_id[cname] = len(self.class_to_id)
-                    self.class_names.append(cname)
-
-                if not cname or cname not in self.class_to_id:
-                    continue
-
-                data = json.loads(ann_str)
-                x_tl = float(data.get("x", 0.0))
-                y_tl = float(data.get("y", 0.0))
-                w = float(data.get("w", 0.0))
-                h = float(data.get("h", 0.0))
-
-                cx = x_tl + w / 2.0
-                cy = y_tl + h / 2.0
-
-                cid = self.class_to_id[cname]
-                label_lines.append(
-                    f"{cid} {cx:.12f} {cy:.12f} {w:.12f} {h:.12f}"
-                )
+            label_lines = self._collect_darknet_bounding_box_labels(group_df)
 
             labels_by_split[split_name][new_stem] = label_lines
 
@@ -101,6 +73,42 @@ def transform(self, prepared_ldf: PreparedLDF) -> None:
 
         self._dump_annotations(labels_by_split, self.output_path, self.part)
 
+    def _collect_darknet_bounding_box_labels(
+        self,
+        group_df: Any,
+    ) -> list[str]:
+        label_lines: list[str] = []
+
+        for row in group_df.iter_rows(named=True):
+            ttype = row.get("task_type")
+            ann_str = row.get("annotation")
+            cname = row.get("class_name")
+
+            if ttype != "boundingbox" or ann_str is None:
+                continue
+
+            # Register class if new
+            if cname and cname not in self.class_to_id:
+                self.class_to_id[cname] = len(self.class_to_id)
+                self.class_names.append(cname)
+
+            if not cname or cname not in self.class_to_id:
+                continue
+
+            data = json.loads(ann_str)
+            x_tl = float(data.get("x", 0.0))
+            y_tl = float(data.get("y", 0.0))
+            w = float(data.get("w", 0.0))
+            h = float(data.get("h", 0.0))
+
+            cx = x_tl + w / 2.0
+            cy = y_tl + h / 2.0
+
+            cid = self.class_to_id[cname]
+            label_lines.append(f"{cid} {cx:.12f} {cy:.12f} {w:.12f} {h:.12f}")
+
+        return label_lines
+
     def _maybe_roll_partition(
         self,
         labels_by_split: dict[str, dict[str, list[str]]],

luxonis_ml/data/exporters/exporter_utils.py

@@ -3,7 +3,9 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
+import cv2
 import polars as pl
+from pycocotools import mask
 
 if TYPE_CHECKING:
     from luxonis_ml.data.datasets.luxonis_dataset import LuxonisDataset
@@ -154,6 +156,28 @@ def get_single_skeleton(
         skeleton_1_based = [[a + 1, b + 1] for a, b in edges]
         return labels, skeleton_1_based
 
+    @staticmethod
+    def rle_to_yolo_polygon(rle: str, height: int, width: int) -> list:
+        # Decode RLE to binary mask
+        m = mask.decode({"size": [height, width], "counts": rle})
+
+        # Each contour = one polygon
+        contours, _ = cv2.findContours(
+            m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
+        )
+
+        polygons = []
+        for contour in contours:
+            contour = contour.squeeze()
+            if len(contour.shape) != 2:
+                continue
+            polygon = []
+            for x, y in contour:
+                polygon.extend([x / width, y / height])
+            polygons.append(polygon)
+
+        return polygons
+
     def _normalize(
         self, xs: list[float], ys: list[float], w: float, h: float
     ) -> list[float]:

luxonis_ml/data/exporters/segmentation_mask_directory_exporter.py

@@ -46,7 +46,7 @@ def get_split_names(self) -> dict[str, str]:
     def _class_id_for(self, split: str, class_name: str) -> int:
         cmap = self.split_class_maps[split]
         if class_name not in cmap:
-            cmap[class_name] = len(cmap) + 1  # start at 1 (0 is background)
+            cmap[class_name] = len(cmap) + 1
         return cmap[class_name]
 
     def _write_classes_csv(self, split: str, split_dir: Path) -> None:

luxonis_ml/data/exporters/tensorflow_csv_exporter.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import csv
+import json
 from pathlib import Path
 from typing import Any, cast
 
@@ -58,6 +59,7 @@ def transform(self, prepared_ldf: PreparedLDF) -> None:
                 if row.get("task_type") != "boundingbox":
                     continue
                 ann = row.get("annotation")
+                ann = json.loads(ann)
                 cname = row.get("class_name")
                 if ann is None or not cname:
                     continue