ICRA/app/qt/image_processor.py

"""Minimal image processing pipeline adapted for the Qt frontend."""

from __future__ import annotations

import colorsys
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, Tuple

from PIL import Image, ImageDraw
from PySide6 import QtGui

from app.logic import PREVIEW_MAX_SIZE


@dataclass
class Stats:
    matches_all: int = 0
    total_all: int = 0
    matches_keep: int = 0
    total_keep: int = 0

    def summary(self) -> str:
        if self.total_all == 0:
            return "Matches with exclusions: —"
        with_pct = (self.matches_keep / self.total_keep * 100) if self.total_keep else 0.0
        without_pct = (self.matches_all / self.total_all * 100) if self.total_all else 0.0
        return (
            f"Matches with exclusions: {with_pct:.1f}%   ·   "
            f"Matches overall: {without_pct:.1f}%"
        )


class QtImageProcessor:
    """Process images and build overlays for the Qt UI."""

    def __init__(self) -> None:
        self.orig_img: Image.Image | None = None
        self.preview_img: Image.Image | None = None
        self.overlay_img: Image.Image | None = None
        self.preview_paths: list[Path] = []
        self.current_index: int = -1
        self.stats = Stats()

        # HSV thresholds and overlay alpha
        self.hue_min = 0
        self.hue_max = 360
        self.sat_min = 25
        self.val_min = 15
        self.val_max = 100
        self.alpha = 120

        self.exclude_shapes: list[dict[str, object]] = []

    # thresholds -------------------------------------------------------------

    def set_threshold(self, key: str, value: int) -> None:
        setattr(self, key, value)
        if self.preview_img is not None:
            self._rebuild_overlay()

    # image handling --------------------------------------------------------

    def load_single_image(self, path: Path) -> None:
        image = Image.open(path).convert("RGBA")
        self.orig_img = image
        self.preview_paths = [path]
        self.current_index = 0
        self._build_preview()
        self._rebuild_overlay()

    def load_folder(self, paths: Iterable[Path], start_index: int = 0) -> None:
        self.preview_paths = list(paths)
        if not self.preview_paths:
            raise ValueError("No images in folder.")
        self.current_index = max(0, min(start_index, len(self.preview_paths) - 1))
        self._load_current()

    def next_image(self) -> None:
        if not self.preview_paths:
            return
        self.current_index = (self.current_index + 1) % len(self.preview_paths)
        self._load_current()

    def previous_image(self) -> None:
        if not self.preview_paths:
            return
        self.current_index = (self.current_index - 1) % len(self.preview_paths)
        self._load_current()

    def _load_current(self) -> None:
        path = self.preview_paths[self.current_index]
        self.load_single_image(path)

    # preview/overlay -------------------------------------------------------

    def _build_preview(self) -> None:
        if self.orig_img is None:
            self.preview_img = None
            return
        width, height = self.orig_img.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)

    def _rebuild_overlay(self) -> None:
        if self.preview_img is None:
            self.overlay_img = None
            return
        base = self.preview_img.convert("RGBA")
        overlay = Image.new("RGBA", base.size, (0, 0, 0, 0))
        draw = ImageDraw.Draw(overlay)
        pixels = base.load()
        width, height = base.size
        highlight = (255, 0, 0, int(self.alpha))
        matches_all = total_all = 0

        for y in range(height):
            for x in range(width):
                r, g, b, a = pixels[x, y]
                if a == 0:
                    continue
                total_all += 1
                if self._matches(r, g, b):
                    draw.point((x, y), fill=highlight)
                    matches_all += 1

        self.overlay_img = overlay
        self.stats = Stats(matches_all=matches_all, total_all=total_all, matches_keep=matches_all, total_keep=total_all)

    # helpers ----------------------------------------------------------------

    def _matches(self, r: int, g: int, b: int) -> bool:
        h, s, v = colorsys.rgb_to_hsv(r / 255.0, g / 255.0, b / 255.0)
        hue = (h * 360.0) % 360.0
        if self.hue_min <= self.hue_max:
            hue_ok = self.hue_min <= hue <= self.hue_max
        else:
            hue_ok = hue >= self.hue_min or hue <= self.hue_max
        sat_ok = s * 100.0 >= self.sat_min
        val_ok = self.val_min <= v * 100.0 <= self.val_max
        return hue_ok and sat_ok and val_ok

    # exported data ----------------------------------------------------------

    def preview_pixmap(self) -> QtGui.QPixmap:
        return self._to_pixmap(self.preview_img)

    def overlay_pixmap(self) -> QtGui.QPixmap:
        if self.preview_img is None or self.overlay_img is None:
            return QtGui.QPixmap()
        merged = Image.alpha_composite(self.preview_img.convert("RGBA"), self.overlay_img)
        return self._to_pixmap(merged)

    @staticmethod
    def _to_pixmap(image: Image.Image | None) -> QtGui.QPixmap:
        if image is None:
            return QtGui.QPixmap()
        buffer = image.tobytes("raw", "RGBA")
        qt_image = QtGui.QImage(buffer, image.width, image.height, QtGui.QImage.Format_RGBA8888)
        return QtGui.QPixmap.fromImage(qt_image)