From 1c48a53c19ea6a877b7a0ba45313d10a9ed76d84 Mon Sep 17 00:00:00 2001 From: lukas Date: Sun, 22 Mar 2026 20:09:05 +0100 Subject: [PATCH] Implement background exclusion and refactor folder structure - Added configurable background exclusion (#1f2937) with tolerance - Implemented alpha thresholding (>= 128) to eliminate edge artifacts - Refactored folder structure into analyses/[slug]/images, settings, and results - Updated pattern puller to skip existing images and handle network errors - Updated .gitignore and automated tests for path integrity --- .gitignore | 2 +- app/lang/de.toml | 1 + app/lang/en.toml | 1 + app/logic/__init__.py | 4 ++ app/logic/constants.py | 15 +++++++- app/qt/app.py | 4 +- app/qt/image_processor.py | 76 +++++++++++++++++++++++++++++++++----- app/qt/main_window.py | 77 ++++++++++++++++++++++++++++++++------- app/qt/pattern_puller.py | 21 +++++++++-- config.toml | 4 ++ tests/test_file_paths.py | 72 ++++++++++++++++++++++++++++++++++++ 11 files changed, 249 insertions(+), 28 deletions(-) create mode 100644 tests/test_file_paths.py diff --git a/.gitignore b/.gitignore index 03ad978..4a423bd 100644 --- a/.gitignore +++ b/.gitignore @@ -158,4 +158,4 @@ uv/ .git/worktrees/ # ICRA specific -images/ +analyses/ diff --git a/app/lang/de.toml b/app/lang/de.toml index 78b6401..37d5938 100644 --- a/app/lang/de.toml +++ b/app/lang/de.toml @@ -12,6 +12,7 @@ "toolbar.toggle_theme" = "Theme umschalten" "toolbar.open_app_folder" = "Programmordner öffnen" "toolbar.prefer_dark" = "Dunkelheit bevorzugen" +"toolbar.exclude_bg" = "Hintergrund ausblenden ({color})" "status.no_file" = "Keine Datei geladen." "status.defaults_restored" = "Standardwerte aktiv." "status.free_draw_enabled" = "Freihand-Ausschluss aktiviert." diff --git a/app/lang/en.toml b/app/lang/en.toml index 259c5ae..0d3541e 100644 --- a/app/lang/en.toml +++ b/app/lang/en.toml @@ -12,6 +12,7 @@ "toolbar.toggle_theme" = "Toggle theme" "toolbar.open_app_folder" = "Open application folder" "toolbar.prefer_dark" = "Prefer darkness" +"toolbar.exclude_bg" = "Exclude Background ({color})" "status.no_file" = "No file loaded." "status.defaults_restored" = "Defaults restored." "status.free_draw_enabled" = "Free-draw exclusion mode enabled." diff --git a/app/logic/__init__.py b/app/logic/__init__.py index f6b1fe2..d4ca1eb 100644 --- a/app/logic/__init__.py +++ b/app/logic/__init__.py @@ -6,6 +6,8 @@ from .constants import ( IMAGES_DIR, LANGUAGE, OVERLAY_COLOR, + EXCLUDE_BG_COLOR, + EXCLUDE_BG_TOLERANCE, PREVIEW_MAX_SIZE, RESET_EXCLUSIONS_ON_IMAGE_CHANGE, SUPPORTED_IMAGE_EXTENSIONS, @@ -17,6 +19,8 @@ __all__ = [ "IMAGES_DIR", "LANGUAGE", "OVERLAY_COLOR", + "EXCLUDE_BG_COLOR", + "EXCLUDE_BG_TOLERANCE", "PREVIEW_MAX_SIZE", "RESET_EXCLUSIONS_ON_IMAGE_CHANGE", "SUPPORTED_IMAGE_EXTENSIONS", diff --git a/app/logic/constants.py b/app/logic/constants.py index 6e403a4..392d514 100644 --- a/app/logic/constants.py +++ b/app/logic/constants.py @@ -94,7 +94,12 @@ def _extract_language(data: dict[str, Any]) -> str: _CONFIG_DATA = _load_config_data() -_OPTION_DEFAULTS = {"reset_exclusions_on_image_change": False, "overlay_color": "#ff0000"} +_OPTION_DEFAULTS = { + "reset_exclusions_on_image_change": False, + "overlay_color": "#ff0000", + "exclude_bg_color": "#1f2937", + "exclude_bg_tolerance": 5, +} def _extract_options(data: dict[str, Any]) -> dict[str, Any]: @@ -108,6 +113,12 @@ def _extract_options(data: dict[str, Any]) -> dict[str, Any]: color = section.get("overlay_color") if isinstance(color, str) and color.startswith("#") and len(color) in (7, 9): result["overlay_color"] = color + exclude_bg = section.get("exclude_bg_color") + if isinstance(exclude_bg, str) and exclude_bg.startswith("#") and len(exclude_bg) in (7, 9): + result["exclude_bg_color"] = exclude_bg + tolerance = section.get("exclude_bg_tolerance") + if isinstance(tolerance, int): + result["exclude_bg_tolerance"] = max(0, min(255, tolerance)) return result @@ -116,3 +127,5 @@ LANGUAGE = _extract_language(_CONFIG_DATA) OPTIONS = {**_OPTION_DEFAULTS, **_extract_options(_CONFIG_DATA)} RESET_EXCLUSIONS_ON_IMAGE_CHANGE = OPTIONS["reset_exclusions_on_image_change"] OVERLAY_COLOR = OPTIONS["overlay_color"] +EXCLUDE_BG_COLOR = OPTIONS["exclude_bg_color"] +EXCLUDE_BG_TOLERANCE = OPTIONS["exclude_bg_tolerance"] diff --git a/app/qt/app.py b/app/qt/app.py index a857eee..48ce3ad 100644 --- a/app/qt/app.py +++ b/app/qt/app.py @@ -46,12 +46,14 @@ def create_application() -> QtWidgets.QApplication: def run() -> int: """Run the PySide6 GUI.""" app = create_application() - from app.logic import OVERLAY_COLOR + from app.logic import OVERLAY_COLOR, EXCLUDE_BG_COLOR, EXCLUDE_BG_TOLERANCE window = MainWindow( language=LANGUAGE, defaults=DEFAULTS.copy(), reset_exclusions=RESET_EXCLUSIONS_ON_IMAGE_CHANGE, overlay_color=OVERLAY_COLOR, + exclude_bg_color=EXCLUDE_BG_COLOR, + exclude_bg_tolerance=EXCLUDE_BG_TOLERANCE, ) # Respect saved geometry from QSettings; fall back to maximised on first launch diff --git a/app/qt/image_processor.py b/app/qt/image_processor.py index 39d207f..2b4c35a 100644 --- a/app/qt/image_processor.py +++ b/app/qt/image_processor.py @@ -127,6 +127,9 @@ class QtImageProcessor: self._cached_mask_size: Tuple[int, int] | None = None self.exclude_ref_size: Tuple[int, int] | None = None self.prefer_dark: bool = False + self.exclude_bg: bool = True + self.exclude_bg_rgb: Tuple[int, int, int] = (31, 41, 55) + self.exclude_bg_tolerance: int = 5 def set_defaults(self, defaults: dict) -> None: for key in self.defaults: @@ -184,13 +187,29 @@ class QtImageProcessor: if self.orig_img is None: self.preview_img = None return - width, height = self.orig_img.size + + img_to_process = self.orig_img.convert("RGBA") + if self.exclude_bg: + # Mask the background color with tolerance on the original image before resizing + # this prevents interpolation artifacts from leaving a background 'halo' + arr = np.array(img_to_process) + r_bg, g_bg, b_bg = self.exclude_bg_rgb + tol = self.exclude_bg_tolerance + bg_mask = ( + (np.abs(arr[..., 0].astype(np.int16) - r_bg) <= tol) & + (np.abs(arr[..., 1].astype(np.int16) - g_bg) <= tol) & + (np.abs(arr[..., 2].astype(np.int16) - b_bg) <= tol) + ) + arr[bg_mask, 3] = 0 + img_to_process = Image.fromarray(arr, "RGBA") + + width, height = img_to_process.size max_w, max_h = PREVIEW_MAX_SIZE scale = min(max_w / width, max_h / height) if scale <= 0: scale = 1.0 size = (max(1, int(width * scale)), max(1, int(height * scale))) - self.preview_img = self.orig_img.resize(size, Image.LANCZOS) + self.preview_img = img_to_process.resize(size, Image.LANCZOS) def _rebuild_overlay(self) -> None: """Build color-match overlay using vectorized NumPy operations.""" @@ -203,7 +222,18 @@ class QtImageProcessor: arr = np.asarray(base, dtype=np.float32) # (H, W, 4) rgb = arr[..., :3] / 255.0 - alpha_ch = arr[..., 3] # alpha channel of the image + alpha_ch = arr[..., 3].copy() # alpha channel of the image + + if self.exclude_bg: + # Exclude specific background color + r_bg, g_bg, b_bg = self.exclude_bg_rgb + tol = self.exclude_bg_tolerance + bg_mask = ( + (np.abs(arr[..., 0] - r_bg) <= tol) & + (np.abs(arr[..., 1] - g_bg) <= tol) & + (np.abs(arr[..., 2] - b_bg) <= tol) + ) + alpha_ch[bg_mask] = 0 hsv = _rgb_to_hsv_numpy(rgb) # (H, W, 3): H°, S%, V% @@ -224,7 +254,7 @@ class QtImageProcessor: & (sat <= float(self.sat_max)) & (val >= float(self.val_min)) & (val <= float(self.val_max)) - & (alpha_ch > 0) + & (alpha_ch >= 128) ) # Exclusion mask (same pixel space as preview) @@ -232,8 +262,7 @@ class QtImageProcessor: keep_match = match_mask & ~excl_mask excl_match = match_mask & excl_mask - visible = alpha_ch > 0 - + visible = alpha_ch >= 128 matches_all = int(match_mask[visible].sum()) total_all = int(visible.sum()) matches_keep = int(keep_match[visible].sum()) @@ -270,7 +299,18 @@ class QtImageProcessor: arr = np.asarray(base, dtype=np.float32) rgb = arr[..., :3] / 255.0 - alpha_ch = arr[..., 3] + alpha_ch = arr[..., 3].copy() + + if self.exclude_bg: + # Exclude background color with tolerance + r_bg, g_bg, b_bg = self.exclude_bg_rgb + tol = self.exclude_bg_tolerance + bg_mask = ( + (np.abs(arr[..., 0] - r_bg) <= tol) & + (np.abs(arr[..., 1] - g_bg) <= tol) & + (np.abs(arr[..., 2] - b_bg) <= tol) + ) + alpha_ch[bg_mask] = 0 hsv = _rgb_to_hsv_numpy(rgb) @@ -291,15 +331,15 @@ class QtImageProcessor: & (sat <= float(self.sat_max)) & (val >= float(self.val_min)) & (val <= float(self.val_max)) - & (alpha_ch > 0) + & (alpha_ch >= 128) ) excl_mask = self._build_exclusion_mask_numpy(base.size) keep_match = match_mask & ~excl_mask excl_match = match_mask & excl_mask - visible = alpha_ch > 0 + visible = alpha_ch >= 128 matches_keep_count = int(keep_match[visible].sum()) keep_visible = visible & ~excl_mask brightness = float(val[keep_visible].mean()) if keep_visible.any() else 0.0 @@ -443,3 +483,21 @@ class QtImageProcessor: self._cached_mask = mask self._cached_mask_size = size return mask + + def set_exclude_bg_color(self, hex_code: str, tolerance: int = 5) -> None: + """Set the RGB channels for background exclusion from a hex string.""" + self.exclude_bg_tolerance = tolerance + if not hex_code.startswith("#") or len(hex_code) not in (7, 9): + return + try: + r = int(hex_code[1:3], 16) + g = int(hex_code[3:5], 16) + b = int(hex_code[5:7], 16) + self.exclude_bg_rgb = (r, g, b) + except ValueError: + pass + + @property + def exclude_bg_color_hex(self) -> str: + r, g, b = self.exclude_bg_rgb + return f"#{r:02x}{g:02x}{b:02x}" diff --git a/app/qt/main_window.py b/app/qt/main_window.py index eda792a..62923b2 100644 --- a/app/qt/main_window.py +++ b/app/qt/main_window.py @@ -516,7 +516,7 @@ class TitleBar(QtWidgets.QWidget): class MainWindow(QtWidgets.QMainWindow, I18nMixin): """Main application window containing all controls.""" - def __init__(self, language: str, defaults: dict, reset_exclusions: bool, overlay_color: str | None = None) -> None: + def __init__(self, language: str, defaults: dict, reset_exclusions: bool, overlay_color: str | None = None, exclude_bg_color: str | None = None, exclude_bg_tolerance: int = 5) -> None: super().__init__() self.init_i18n(language) self.setWindowTitle(self._t("app.title")) @@ -540,6 +540,8 @@ class MainWindow(QtWidgets.QMainWindow, I18nMixin): self.processor.reset_exclusions_on_switch = reset_exclusions # Always use red for the overlay regardless of the target color self.processor.set_overlay_color(DEFAULT_OVERLAY_HEX) + if exclude_bg_color: + self.processor.set_exclude_bg_color(exclude_bg_color, tolerance=exclude_bg_tolerance) self.content_layout = QtWidgets.QVBoxLayout(self.content) self.content_layout.setContentsMargins(24, 0, 24, 24) @@ -655,6 +657,13 @@ class MainWindow(QtWidgets.QMainWindow, I18nMixin): self.prefer_dark_action.setChecked(False) self.prefer_dark_action.triggered.connect(lambda: self._invoke_action("toggle_prefer_dark")) view_menu.addAction(self.prefer_dark_action) + + self.exclude_bg_action = QtGui.QAction("🖼 " + self._t("toolbar.exclude_bg", color=self.processor.exclude_bg_color_hex), self) + self.exclude_bg_action.setCheckable(True) + self.exclude_bg_action.setChecked(True) + self.exclude_bg_action.triggered.connect(lambda: self._invoke_action("toggle_exclude_bg")) + view_menu.addAction(self.exclude_bg_action) + view_menu.addSeparator() view_menu.addAction("📁 " + self._t("toolbar.open_app_folder"), lambda: self._invoke_action("open_app_folder")) @@ -806,6 +815,7 @@ class MainWindow(QtWidgets.QMainWindow, I18nMixin): "reset_sliders": self._reset_sliders, "toggle_theme": self.toggle_theme, "toggle_prefer_dark": self.toggle_prefer_dark, + "toggle_exclude_bg": self.toggle_exclude_bg, "open_app_folder": self.open_app_folder, "show_previous_image": self.show_previous_image, "show_next_image": self.show_next_image, @@ -856,7 +866,7 @@ class MainWindow(QtWidgets.QMainWindow, I18nMixin): def open_image(self) -> None: filters = "Images (*.png *.jpg *.jpeg *.bmp *.webp)" - default_dir = str(Path("images").absolute()) if Path("images").exists() else "" + default_dir = str(Path("analyses").absolute()) if Path("analyses").exists() else "" path_str, _ = QtWidgets.QFileDialog.getOpenFileName(self, self._t("dialog.open_image_title"), default_dir, filters) if not path_str: return @@ -873,7 +883,7 @@ class MainWindow(QtWidgets.QMainWindow, I18nMixin): self._refresh_views() def open_folder(self) -> None: - default_dir = str(Path("images").absolute()) if Path("images").exists() else "" + default_dir = str(Path("analyses").absolute()) if Path("analyses").exists() else "" directory = QtWidgets.QFileDialog.getExistingDirectory(self, self._t("dialog.open_folder_title"), default_dir) if not directory: return @@ -895,20 +905,28 @@ class MainWindow(QtWidgets.QMainWindow, I18nMixin): def export_settings(self) -> None: item_name = "" + default_dir = Path("analyses").absolute() if Path("analyses").exists() else Path() + if self._current_image_path: - # Try to get folder name first, otherwise file name - if self._current_image_path.parent.name and self._current_image_path.parent.name != "images": - item_name = self._current_image_path.parent.name + if self._current_image_path.parent.name == "images": + root_dir = self._current_image_path.parent.parent + item_name = root_dir.name + default_dir = root_dir / "settings" + elif self._current_image_path.parent.name and self._current_image_path.parent.name != "images": + root_dir = self._current_image_path.parent + item_name = root_dir.name + default_dir = root_dir / "settings" else: item_name = self._current_image_path.stem - + default_dir = self._current_image_path.parent / "settings" + + default_dir.mkdir(parents=True, exist_ok=True) default_filename = f"icra_settings_{item_name}.json" if item_name else "icra_settings.json" - default_dir = str(Path("images").absolute()) if Path("images").exists() else "" path_str, _ = QtWidgets.QFileDialog.getSaveFileName( self, self._t("dialog.export_settings_title"), - str(Path(default_dir) / default_filename), + str(default_dir / default_filename), self._t("dialog.json_filter") ) if not path_str: @@ -935,11 +953,25 @@ class MainWindow(QtWidgets.QMainWindow, I18nMixin): QtWidgets.QMessageBox.warning(self, self._t("dialog.error_title"), str(e)) def import_settings(self) -> None: - default_dir = str(Path("images").absolute()) if Path("images").exists() else "" + default_dir = Path("analyses").absolute() if Path("analyses").exists() else Path() + + if self._current_image_path: + if self._current_image_path.parent.name == "images": + root_dir = self._current_image_path.parent.parent + default_dir = root_dir / "settings" + elif self._current_image_path.parent.name and self._current_image_path.parent.name != "images": + root_dir = self._current_image_path.parent + default_dir = root_dir / "settings" + else: + default_dir = self._current_image_path.parent / "settings" + + if not default_dir.exists(): + default_dir = Path("analyses").absolute() if Path("analyses").exists() else Path() + path_str, _ = QtWidgets.QFileDialog.getOpenFileName( self, self._t("dialog.import_settings_title"), - default_dir, + str(default_dir), self._t("dialog.json_filter") ) if not path_str: @@ -989,12 +1021,22 @@ class MainWindow(QtWidgets.QMainWindow, I18nMixin): return folder_path = self.processor.preview_paths[0].parent - default_filename = f"icra_stats_{folder_path.name}.csv" + + if folder_path.name == "images": + root_dir = folder_path.parent + item_name = root_dir.name + else: + root_dir = folder_path + item_name = root_dir.name + + target_dir = root_dir / "results" + target_dir.mkdir(parents=True, exist_ok=True) + default_filename = f"icra_results_{item_name}.csv" if item_name else "icra_results.csv" csv_path, _ = QtWidgets.QFileDialog.getSaveFileName( self, self._t("dialog.export_stats_title"), - str(folder_path / default_filename), + str(target_dir / default_filename), self._t("dialog.csv_filter") ) if not csv_path: @@ -1299,6 +1341,15 @@ class MainWindow(QtWidgets.QMainWindow, I18nMixin): self.processor._rebuild_overlay() self._refresh_overlay_only() + def toggle_exclude_bg(self) -> None: + self.processor.exclude_bg = not self.processor.exclude_bg + self.exclude_bg_action.setChecked(self.processor.exclude_bg) + if self._current_image_path: + # Re-scale/re-apply transparency + self.processor._build_preview() + self.processor._rebuild_overlay() + self._refresh_views() + def clear_exclusions(self) -> None: self.image_view.clear_shapes() self.processor.set_exclusions([]) diff --git a/app/qt/pattern_puller.py b/app/qt/pattern_puller.py index b0b28df..38e18e4 100644 --- a/app/qt/pattern_puller.py +++ b/app/qt/pattern_puller.py @@ -30,17 +30,32 @@ class PatternDownloadWorker(QtCore.QThread): filename = self.save_dir / f"{seed}.png" if filename.exists(): - return True, None + try: + from PIL import Image + with Image.open(filename) as img: + img.verify() + return True, None + except Exception: + filename.unlink(missing_ok=True) try: req = urllib.request.Request(url, headers={'User-Agent': 'Mozilla/5.0 ICRA/1.0'}) with urllib.request.urlopen(req, timeout=10) as response: with open(filename, 'wb') as f: f.write(response.read()) - return True, None + + try: + from PIL import Image + with Image.open(filename) as img: + img.verify() + return True, None + except Exception: + filename.unlink(missing_ok=True) + return False, "Downloaded file is invalid/corrupt" except urllib.error.HTTPError as e: return False, f"HTTP {e.code}" except Exception as e: + filename.unlink(missing_ok=True) return False, f"Network error: {e}" def run(self) -> None: @@ -140,7 +155,7 @@ class PatternPullerDialog(QtWidgets.QDialog, I18nMixin): QtWidgets.QMessageBox.warning(self, "Error", self._t("dialog.puller_invalid_url", default="Invalid URL format.")) return - save_dir = Path("images") / slug + save_dir = Path("analyses") / slug / "images" self.start_btn.setEnabled(False) self.url_input.setEnabled(False) diff --git a/config.toml b/config.toml index b91c6f3..3c93ba6 100644 --- a/config.toml +++ b/config.toml @@ -7,6 +7,10 @@ language = "en" reset_exclusions_on_image_change = false # Hex color code for the match overlay (e.g. "#ff0000" for Red, "#00ff00" for Green) overlay_color = "#ff0000" +# Hex color code for the background to be excluded (default #1f2937) +exclude_bg_color = "#1f2937" +# Tolerance for background color matching (0-255, default 5) +exclude_bg_tolerance = 5 [defaults] # Override any of the following keys to tweak the initial slider values: diff --git a/tests/test_file_paths.py b/tests/test_file_paths.py new file mode 100644 index 0000000..0eceb39 --- /dev/null +++ b/tests/test_file_paths.py @@ -0,0 +1,72 @@ +import pytest +from pathlib import Path +from unittest.mock import MagicMock, patch + +from PySide6 import QtWidgets +from app.qt.main_window import MainWindow +from app.qt.pattern_puller import PatternDownloadWorker + + +@pytest.fixture +def qt_app(): + from PySide6.QtWidgets import QApplication + import sys + app = QApplication.instance() + if app is None: + app = QApplication(sys.argv) + yield app + + +def test_export_settings_path_generation(qt_app, tmp_path): + mock_widget = QtWidgets.QWidget() + mock_widget.title_label = MagicMock() + mock_widget.apply_theme = MagicMock() + with patch('app.qt.main_window.TitleBar', return_value=mock_widget): + window = MainWindow(language="en", defaults={}, reset_exclusions=False) + + with patch("PySide6.QtWidgets.QFileDialog.getSaveFileName", return_value=("", "")) as mock_get_save: + + # Test case 1: New subfolder structure (analyses/m4a1-s/images/1.png) + root = tmp_path / "analyses" / "m4a1_s" + img_dir = root / "images" + img_dir.mkdir(parents=True) + img_path = img_dir / "1.png" + + window._current_image_path = img_path + window.export_settings() + + # Verify settings directory was created + assert (root / "settings").exists() + + # Verify default path given to QFileDialog + args, kwargs = mock_get_save.call_args + # args[2] is the default path string + expected_path = str(root / "settings" / "icra_settings_m4a1_s.json") + assert args[2] == expected_path + + +def test_export_folder_path_generation(qt_app, tmp_path): + mock_widget = QtWidgets.QWidget() + mock_widget.title_label = MagicMock() + mock_widget.apply_theme = MagicMock() + with patch('app.qt.main_window.TitleBar', return_value=mock_widget): + window = MainWindow(language="en", defaults={}, reset_exclusions=False) + + with patch("PySide6.QtWidgets.QFileDialog.getSaveFileName", return_value=("", "")) as mock_get_save: + # Mock processor paths + root = tmp_path / "analyses" / "m4a1_s" + img_dir = root / "images" + img_dir.mkdir(parents=True) + window.processor.preview_paths = [img_dir / "1.png"] + + window.export_folder() + + assert (root / "results").exists() + args, kwargs = mock_get_save.call_args + expected_path = str(root / "results" / "icra_results_m4a1_s.csv") + assert args[2] == expected_path + + +def test_pattern_download_worker_dir(tmp_path): + worker = PatternDownloadWorker(slug="test-slug", save_dir=tmp_path / "analyses" / "test-slug" / "images") + assert worker.save_dir == tmp_path / "analyses" / "test-slug" / "images"