Add functionality to store label colors in annotations

.github/workflows/builds.yml

@@ -22,6 +22,10 @@ jobs:
           # MacOS issue ref.: https://github.com/actions/setup-python/issues/855 & https://github.com/actions/setup-python/issues/865
           python-version: ${{ matrix.os == 'macos-latest' && matrix.python == '3.10' && '3.11' || matrix.python }}
           architecture: x64
+      # MacOS: use crnn_vgg16_bn as recognition architecture to avoid onnxruntime error
+      - name: Set recognition arch for macOS
+        if: matrix.os == 'macos-latest'
+        run: echo "RECOGNITION_ARCH=crnn_vgg16_bn" >> $GITHUB_ENV
       - name: Cache python modules
         uses: actions/cache@v5
         with:

.github/workflows/main.yml

@@ -21,6 +21,10 @@ jobs:
           # MacOS issue ref.: https://github.com/actions/setup-python/issues/855 & https://github.com/actions/setup-python/issues/865
           python-version: ${{ matrix.os == 'macos-latest' && matrix.python == '3.10' && '3.11' || matrix.python }}
           architecture: x64
+      # MacOS: use crnn_vgg16_bn as recognition architecture to avoid onnxruntime error
+      - name: Set recognition arch for macOS
+        if: matrix.os == 'macos-latest'
+        run: echo "RECOGNITION_ARCH=crnn_vgg16_bn" >> $GITHUB_ENV
       - name: Cache python modules
         uses: actions/cache@v5
         with:

labeler/views/canvas.py

@@ -192,9 +192,23 @@ def load_json(self):
             polygons = annotations.get("polygons", [])
             labels = annotations.get("labels", ["" for _ in polygons])
             types = annotations.get("types", [self.root.type_options[0]] * len(polygons))
+            type_color_mapping = annotations.get("type_color_mapping", {})
 
-            self.draw_polys(polygons, types, labels)
-            self.update_types(types)
+            if type_color_mapping:
+                self.root._update_color_palette(type_color_mapping)
+
+            colors = []
+            for poly_type in types:
+                if poly_type in self.root.type_options:
+                    type_idx = self.root.type_options.index(poly_type)
+                    colors.append(self.root.color_palette[type_idx])
+                else:
+                    logger.warning(
+                        f"Polygon type '{poly_type}' not found in type_options. Using gray as fallback color."
+                    )
+                    colors.append("#808080")
+
+            self.draw_polys(polygons, types, labels, colors)
 
             logger.info(f"Loaded {len(polygons)} polygons from annotations")
         self.drawing_polygon = False
@@ -232,13 +246,18 @@ def save_json(self) -> str:
                 logger.warning("Not saving JSON as no polygons are present")
                 return "--> Nothing to save, because no polygons are present"
 
+            type_color_mapping = {}
+            for type_name, color in zip(self.root.type_options, self.root.color_palette):
+                type_color_mapping[type_name] = color
+
             data = {
                 img_name: {
                     "img_dimensions": [self.img_height, self.img_width],
                     "img_hash": self._get_img_hash(),
                     "polygons": filtered_polygons,
                     "labels": labels,
                     "types": types,
+                    "type_color_mapping": type_color_mapping,
                 }
             }
 
@@ -268,14 +287,21 @@ def final_save(self):
             json.dump(data, fl, indent=4, ensure_ascii=False)  # type: ignore[arg-type]
         logger.info(f"Finally saved to {os.path.join(self.root_path, 'labels.json')}")
 
-    def draw_polys(self, coords: list[list[list[int]]], poly_types: list[str], poly_texts: list[str]):
-        """ "
+    def draw_polys(
+        self,
+        coords: list[list[list[int]]],
+        poly_types: list[str],
+        poly_texts: list[str],
+        colors: list[str] | None = None,
+    ):
+        """
         Draw polygons on the canvas
 
         Args:
             coords: list[list[list[int]]]: List of polygons to draw
             poly_types: list[str]: List of polygon types
             poly_texts: list[str]: List of polygon texts
+            colors: list[str] | None: List of polygon colors (optional)
         """
         with self.polygons_mutex:
             # Scale coordinates and create polygons
@@ -284,6 +310,12 @@ def draw_polys(self, coords: list[list[list[int]]], poly_types: list[str], poly_
                 Polygon(self.root, self.canvas, poly, poly_type, poly_text)
                 for poly, poly_type, poly_text in zip(scaled_coords, poly_types, poly_texts)
             )
+            if colors:
+                start_idx = len(self.polygons) - len(colors)
+                for i, color in enumerate(colors):
+                    polygon = self.polygons[start_idx + i]
+                    polygon.update_color(color)
+
             # Update also the original coordinates for accurate scaling to the original image size
             for polygon in self.polygons:
                 polygon.original_coords = [
@@ -292,26 +324,6 @@ def draw_polys(self, coords: list[list[list[int]]], poly_types: list[str], poly_
 
         logger.info(f"Total Polygons drawn: {len(self.polygons)}")
 
-    def update_types(self, poly_types: list[str]):
-        """
-        Update the polygon type colors based on the selected type options.
-
-        Args:
-            poly_types: list[str]: List of polygon types
-        """
-        # Update colors for polygons with type annotations
-        # Check that all poly_type values are in the type options otherwise log a warning
-        missing_types = list(set(poly_types) - set(self.root.type_options))
-        if len(self.root.type_options) == 1:
-            self.root.type_options.extend(missing_types)
-            self.root.color_palette.extend(self.root._generate_color_palette(len(missing_types)))
-            self.root.label_type["values"] = self.root.type_options
-        if missing_types:
-            logger.warning(f"Polygon types {missing_types} not in type options {self.root.type_options}")
-        for polygon in self.polygons:
-            if polygon.poly_type != self.root.type_options[0] and polygon.poly_type in self.root.type_options:
-                polygon.update_color(self.root.color_palette[self.root.type_options.index(polygon.poly_type) - 1])
-
     def add_poly(self, pts: list[list[int]]):
         """
         Add a polygon to the canvas.

labeler/views/gui.py

@@ -5,7 +5,6 @@
 
 import colorsys
 import os
-import random
 import threading
 
 import darkdetect
@@ -298,18 +297,43 @@ def _generate_color_palette(self, num_colors: int):
         """
         palette = []
         golden_ratio_conjugate = 0.61803398875
-        h = random.random()
-
-        for _ in range(num_colors):
+        h = 0.1
+        for i in range(num_colors):
             h = (h + golden_ratio_conjugate) % 1
-            s = 0.5 + random.random() * 0.5
-            v = 0.7 + random.random() * 0.3
-            r, g, b = colorsys.hsv_to_rgb(h, s, v)
-            hex_color = "#{:02x}{:02x}{:02x}".format(int(r * 255), int(g * 255), int(b * 255))
+            # use hue range 0.1 (orange) to 0.9 (pink) to avoid red
+            # formula: h_without_red = min + h * (max - min)
+            h_without_red = 0.1 + h * 0.75
+
+            if i % 2 == 0:
+                s, v = 0.85, 0.90
+            else:
+                s, v = 0.60, 0.95
+
+            r, g, b = colorsys.hsv_to_rgb(h_without_red, s, v)
+            hex_color = "#{:02X}{:02X}{:02X}".format(int(r * 255), int(g * 255), int(b * 255))
             palette.append(hex_color)
 
         return palette
 
+    def _update_color_palette(self, type_color_mapping: dict[str, str]):
+        """
+        Update the color palette based on loaded type-color mapping.
+
+        Args:
+            type_color_mapping: Dictionary mapping type names to hex colors
+        """
+        num_new_colors = 0
+        for type_name in self.type_options:
+            if type_name in type_color_mapping:
+                type_idx = self.type_options.index(type_name)
+                self.color_palette[type_idx] = type_color_mapping[type_name]
+            else:
+                type_idx = self.type_options.index(type_name)
+                if type_idx >= len(self.color_palette):
+                    num_new_colors += 1
+        new_colors = self._generate_color_palette(num_new_colors)
+        self.color_palette.extend(new_colors)
+
     def _validate_numeric_input(self, new_value: str) -> bool:
         """
         Validation function to allow only numeric input.

tests/common/test_canvas.py

@@ -1,6 +1,6 @@
-import logging
+import json
 import os
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock, mock_open, patch
 
 import pytest
 
@@ -39,9 +39,10 @@ def test_current_state(image_on_canvas):
 
 
 def test_save_json(image_on_canvas):
-    mock_polygon = MagicMock(pt_coords=[[0, 0], [10, 0], [10, 10], [0, 10]], text="label", poly_type="type")
+    mock_polygon = MagicMock(pt_coords=[[0, 0], [10, 0], [10, 10], [0, 10]], text="label", poly_type="words")
     image_on_canvas.polygons = [mock_polygon]
-
+    image_on_canvas.root.type_options = ["words"]
+    image_on_canvas.root.color_palette = ["#FF0000"]
     with (
         patch("os.makedirs") as mock_makedirs,
         patch("hashlib.sha256", return_value=MagicMock(hexdigest=lambda: "mockhash")),
@@ -98,24 +99,38 @@ def test_save_json_empty_polygons_returns_message(image_on_canvas):
     assert result.startswith("--> Nothing to save")
 
 
-def test_update_types_updates_colors_and_logs(image_on_canvas, caplog):
-    image_on_canvas.root.type_options = ["default"]
-    image_on_canvas.root.color_palette = ["#000000"]
-    image_on_canvas.root._generate_color_palette = MagicMock(return_value=["#ff0000", "#00ff00"])
-    image_on_canvas.root.label_type = {"values": image_on_canvas.root.type_options}
-
-    polygon1 = MagicMock(poly_type="invoice")
-    polygon2 = MagicMock(poly_type="receipt")
-    polygon3 = MagicMock(poly_type="default")  # should be ignored for update_color
-    image_on_canvas.polygons = [polygon1, polygon2, polygon3]
+def test_load_json_invalid_type_fallback_color(image_on_canvas, mock_annotation_data):
+    json_content = json.dumps(mock_annotation_data)
+    image_on_canvas.image_path = "/some/path/mock_image.jpg"
+    image_on_canvas.root.type_options = ["unknown_type"]
+    with (
+        patch("os.path.exists", return_value=True),
+        patch("builtins.open", mock_open(read_data=json_content)),
+        patch.object(image_on_canvas, "draw_polys") as mock_draw,
+    ):
+        image_on_canvas.load_json()
+        mock_draw.assert_called_once()
+        passed_colors = mock_draw.call_args[0][3]
+        assert passed_colors == ["#808080"]
+    assert image_on_canvas.drawing_polygon is False
 
-    with caplog.at_level(logging.WARNING):
-        image_on_canvas.update_types(["invoice", "receipt"])
 
-    assert "invoice" in image_on_canvas.root.type_options
-    assert "receipt" in image_on_canvas.root.type_options
-    assert len(image_on_canvas.root.color_palette) == 3
+def test_load_json_color_mapping(image_on_canvas, mock_annotation_data):
+    image_on_canvas.image_path = "mock_image.jpg"
+    image_on_canvas.root.type_options = ["words", "document_type", "invoice_id"]
+    image_on_canvas.root.color_palette = ["#FF0000", "#00FF00", "#0000FF"]
+    mock_annotation_data["mock_image.jpg"]["types"] = ["document_type", "invoice_id"]
+    json_content = json.dumps(mock_annotation_data)
 
-    polygon1.update_color.assert_called_once()
-    polygon2.update_color.assert_called_once()
-    polygon3.update_color.assert_not_called()
+    with (
+        patch("os.path.exists", return_value=True),
+        patch("builtins.open", mock_open(read_data=json_content)),
+        patch.object(image_on_canvas, "draw_polys") as mock_draw_polys,
+    ):
+        image_on_canvas.load_json()
+        mock_draw_polys.assert_called_once()
+        actual_types = mock_draw_polys.call_args[0][1]
+        assert actual_types == ["document_type", "invoice_id"]
+        actual_colors = mock_draw_polys.call_args[0][3]
+        assert actual_colors == ["#00FF00", "#0000FF"]
+        assert len(actual_colors) == len(actual_types)

tests/common/test_gui.py

@@ -416,3 +416,23 @@ def test_discard_tight(gui_app):
 
     mock_tight_box.discard_tight_box.assert_called_once()
     gui_app.save_image_button.configure.assert_called_once_with(state="normal")
+
+
+def test_generate_color_palette_alternating_values(gui_app):
+    num_colors = 2
+    palette = gui_app._generate_color_palette(num_colors)
+    assert len(palette) == num_colors
+    assert all(c.startswith("#") for c in palette)
+    assert palette[0] != palette[1]
+
+
+def test_update_color_palette_with_mapping_and_new_types(gui_app):
+    gui_app.type_options = ["words", "header", "footer"]
+    gui_app.color_palette = ["#FF0000", "#000000", "#FFFFFF"]
+    mapping = {"header": "#123456"}
+    gui_app.type_options.append("new_type")
+    gui_app._update_color_palette(mapping)
+
+    assert gui_app.color_palette[1] == "#123456"
+    assert len(gui_app.color_palette) == 4
+    assert gui_app.color_palette[3].startswith("#")

tests/conftest.py

@@ -110,3 +110,15 @@ def mock_mixed_data_folder(tmpdir_factory):
     with open(img_fn, "wb") as f:
         f.write(file.getbuffer())
     return str(temp_folder)
+
+
+@pytest.fixture(scope="session")
+def mock_annotation_data():
+    return {
+        "mock_image.jpg": {
+            "polygons": [[[10, 10], [100, 10], [100, 50], [10, 50]]],
+            "labels": ["INVOICE"],
+            "types": ["document_type"],
+            "type_color_mapping": {"document_type": "#FF5733"},
+        }
+    }