[DAR-1104][Internal] Fix colours missmatch between files on semanic m…

…ask export (#792)
v7labs · Mar 18, 2024 · ea40c79 · ea40c79
1 parent 423de1b
commit ea40c79
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 6 deletions.
diff --git a/darwin/exporter/formats/mask.py b/darwin/exporter/formats/mask.py
@@ -326,7 +326,6 @@ def render_raster(
     """
     errors: List[Exception] = []
 
-    mask_annotations: List[dt.AnnotationMask] = []
     raster_layer: dt.RasterLayer
 
     mask_colours: Dict[str, int] = {}
@@ -348,13 +347,11 @@ def render_raster(
                 errors.append(e)
                 continue
 
-            mask_annotations.append(new_mask)
-
             # Add the category to the list of categories
             if new_mask.name not in categories:
                 categories.append(new_mask.name)
 
-            colour_to_draw = categories.index(new_mask.name) + 1
+            colour_to_draw = categories.index(new_mask.name)
 
             if new_mask.id not in mask_colours:
                 mask_colours[new_mask.id] = colour_to_draw

diff --git a/tests/darwin/exporter/formats/export_mask_test.py b/tests/darwin/exporter/formats/export_mask_test.py
@@ -431,7 +431,7 @@ def test_render_raster() -> None:
     ]
     mask = np.zeros((100, 100), dtype=np.uint8)
     colours: dt.MaskTypes.ColoursDict = {}
-    categories: dt.MaskTypes.CategoryList = []
+    categories: dt.MaskTypes.CategoryList = ["__background__"]
     annotations: List[dt.AnnotationLike] = [
         dt.Annotation(
             dt.AnnotationClass("mask1", "mask"),
@@ -496,7 +496,7 @@ def test_render_raster() -> None:
 
         assert_array_equal(result_mask, np.array(rle_code, dtype=np.uint8).reshape((100, 100)))  # type: ignore
 
-        assert result_categories == ["mask1", "mask2", "mask3"]
+        assert result_categories == ["__background__", "mask1", "mask2", "mask3"]
         assert result_colours == {"mask1": 1, "mask2": 2, "mask3": 3}