Fix Gray Norm Issue (#328)

Author: JSabadin

luxonis_ml/data/__main__.py

@@ -472,6 +472,8 @@ def export(
     ] = False,
     bucket_storage: BucketStorage = bucket_option,
 ):
+    """Export a Luxonis dataset to disk in native format, with optional
+    partitioning and zip compression for flexible storage."""
     save_dir = save_dir or dataset_name
     if delete_existing and Path(save_dir).exists():
         shutil.rmtree(save_dir)

luxonis_ml/data/augmentations/albumentations_engine.py

@@ -554,6 +554,8 @@ def postprocess(
         """
         out_labels = {}
         out_image = data.pop("image")
+        if out_image.ndim == 2:
+            out_image = np.expand_dims(out_image, axis=-1)
         image_height, image_width, _ = out_image.shape
 
         bboxes_indices = {}

luxonis_ml/data/loaders/luxonis_loader.py

@@ -249,10 +249,9 @@ def __getitem__(self, idx: int) -> LoaderOutput:
         if not self.exclude_empty_annotations:
             img, labels = self._add_empty_annotations(img, labels)
 
+        # Albumentations needs RGB
         if self.color_space == "BGR":
             img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-        elif self.color_space == "GRAY":
-            img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)[..., np.newaxis]
 
         return img, labels
 
@@ -316,7 +315,12 @@ def _load_data(self, idx: int) -> tuple[np.ndarray, Labels]:
 
         img_path = self.idx_to_img_path[idx]
 
-        img = cv2.cvtColor(cv2.imread(str(img_path)), cv2.COLOR_BGR2RGB)
+        if self.color_space == "GRAY":
+            img = cv2.cvtColor(cv2.imread(str(img_path)), cv2.COLOR_RGB2GRAY)[
+                ..., np.newaxis
+            ]
+        else:
+            img = cv2.cvtColor(cv2.imread(str(img_path)), cv2.COLOR_BGR2RGB)
 
         labels_by_task: dict[str, list[Annotation]] = defaultdict(list)
         class_ids_by_task: dict[str, list[int]] = defaultdict(list)

tests/test_data/test_loader.py

@@ -4,7 +4,6 @@
 from pathlib import Path
 from typing import Any
 
-import cv2
 import numpy as np
 
 from luxonis_ml.data import (
@@ -562,19 +561,40 @@ def round_nested_list(
 
 
 def test_colorspace(storage_url: str, tempdir: Path):
-    loader = create_loader(storage_url, tempdir)
+    norm_3d = [
+        {
+            "name": "Normalize",
+            "params": {
+                "mean": [0.5, 0.5, 0.5],
+                "std": [0.5, 0.5, 0.5],
+                "p": 1,
+            },
+        },
+    ]
+    norm_1d = [
+        {
+            "name": "Normalize",
+            "params": {
+                "mean": [0.5],
+                "std": [0.5],
+                "p": 1,
+            },
+        },
+    ]
+    loader = create_loader(storage_url, tempdir, augmentation_config=norm_3d)
     rgb_img, _ = next(iter(loader))
     assert len(rgb_img.shape) == 3
     assert rgb_img.shape[2] == 3
-    loader = create_loader(storage_url, tempdir, color_space="BGR")
+    loader = create_loader(
+        storage_url, tempdir, color_space="BGR", augmentation_config=norm_3d
+    )
     bgr_img, _ = next(iter(loader))
     assert len(bgr_img.shape) == 3
     assert bgr_img.shape[2] == 3
     assert np.array_equal(rgb_img, bgr_img[:, :, ::-1])
-    loader = create_loader(storage_url, tempdir, color_space="GRAY")
+    loader = create_loader(
+        storage_url, tempdir, color_space="GRAY", augmentation_config=norm_1d
+    )
     gray_img, _ = next(iter(loader))
     assert len(gray_img.shape) == 3
     assert gray_img.shape[2] == 1
-    assert np.array_equal(
-        cv2.cvtColor(rgb_img, cv2.COLOR_RGB2GRAY), gray_img[:, :, 0]
-    )