diff --git a/.github/workflows/builds.yml b/.github/workflows/builds.yml index 211500d..9b42d53 100644 --- a/.github/workflows/builds.yml +++ b/.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: diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index f62abc0..307862b 100644 --- a/.github/workflows/main.yml +++ b/.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: diff --git a/labeler/views/canvas.py b/labeler/views/canvas.py index 02432a0..f933652 100644 --- a/labeler/views/canvas.py +++ b/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,6 +246,10 @@ 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], @@ -239,6 +257,7 @@ def save_json(self) -> str: "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. diff --git a/labeler/views/gui.py b/labeler/views/gui.py index 1654a04..b048e9f 100644 --- a/labeler/views/gui.py +++ b/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. diff --git a/tests/common/test_canvas.py b/tests/common/test_canvas.py index 6ddb396..a69ff77 100644 --- a/tests/common/test_canvas.py +++ b/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) diff --git a/tests/common/test_gui.py b/tests/common/test_gui.py index 79fb8b7..bf18bc1 100644 --- a/tests/common/test_gui.py +++ b/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("#") diff --git a/tests/conftest.py b/tests/conftest.py index 826f7a3..df95b19 100644 --- a/tests/conftest.py +++ b/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"}, + } + }