Shape_Model_2D/pm2d/gui.py

"""GUI standalone OpenCV per Pattern Matching 2D.

Flusso:
1. Apri immagine modello (file dialog tk)
2. Selezione ROI con cv2.selectROI
3. Apri immagine scena
4. Esegui matching
5. Visualizza risultati (baricentro, angolo, score, bbox)

Tutta la logica algoritmica vive in pm2d.matcher.EdgeShapeMatcher.
"""

from __future__ import annotations

import sys
from pathlib import Path
from tkinter import Tk, filedialog

import cv2
import numpy as np

from pm2d.matcher import EdgeShapeMatcher
from pm2d.line_matcher import LineShapeMatcher, Match


WINDOW_MODEL = "Modello (selezionare ROI - INVIO conferma, c annulla)"
WINDOW_RESULT = "Risultato matching"


def pick_file(title: str, initialdir: str | None = None) -> str | None:
    """Tk file picker (root nascosto)."""
    root = Tk()
    root.withdraw()
    path = filedialog.askopenfilename(
        title=title,
        initialdir=initialdir,
        filetypes=[
            ("Immagini", "*.png *.jpg *.jpeg *.bmp *.tif *.tiff"),
            ("Tutti i file", "*.*"),
        ],
    )
    root.destroy()
    return path or None


def load_image(path: str) -> np.ndarray:
    img = cv2.imread(path, cv2.IMREAD_COLOR)
    if img is None:
        raise FileNotFoundError(f"Impossibile leggere immagine: {path}")
    return img


def select_roi(image: np.ndarray) -> np.ndarray | None:
    """Apre finestra di selezione ROI. Ritorna ROI BGR o None se annullato."""
    disp = _fit_for_display(image, max_side=1200)
    scale = disp.shape[1] / image.shape[1]
    r = cv2.selectROI(WINDOW_MODEL, disp, showCrosshair=True, fromCenter=False)
    cv2.destroyWindow(WINDOW_MODEL)
    x, y, w, h = r
    if w == 0 or h == 0:
        return None
    # Riporta a coordinate immagine originale
    x0 = int(round(x / scale))
    y0 = int(round(y / scale))
    w0 = int(round(w / scale))
    h0 = int(round(h / scale))
    x0 = max(0, x0); y0 = max(0, y0)
    w0 = max(1, min(w0, image.shape[1] - x0))
    h0 = max(1, min(h0, image.shape[0] - y0))
    return image[y0:y0 + h0, x0:x0 + w0].copy()


def _fit_for_display(image: np.ndarray, max_side: int = 1200) -> np.ndarray:
    h, w = image.shape[:2]
    m = max(h, w)
    if m <= max_side:
        return image
    s = max_side / m
    return cv2.resize(image, (int(w * s), int(h * s)), interpolation=cv2.INTER_AREA)


def draw_matches(scene: np.ndarray, matches: list[Match]) -> np.ndarray:
    """Disegna baricentro, asse orientamento, bbox ruotato per ogni match."""
    out = scene.copy()
    for i, m in enumerate(matches):
        color = _color_for(i)
        # Bbox ruotato: il template ruotato di angle_deg ha bbox assi-allineato
        # nel sistema variante; per disegnarlo esatto, ricaviamo il rettangolo
        # ruotato del template originale attorno al baricentro.
        x, y, w, h = m.bbox
        # box assi-allineato della variante
        cv2.rectangle(out, (x, y), (x + w, y + h), color, 1, cv2.LINE_AA)
        # Baricentro
        cx, cy = int(round(m.cx)), int(round(m.cy))
        cv2.drawMarker(out, (cx, cy), color, cv2.MARKER_CROSS, 22, 2, cv2.LINE_AA)
        cv2.circle(out, (cx, cy), 4, color, -1, cv2.LINE_AA)
        # Asse orientamento (lunghezza ~ metà altezza bbox)
        L = max(h, w) // 2
        ang_rad = np.deg2rad(m.angle_deg)
        ex = int(round(cx + L * np.cos(ang_rad)))
        ey = int(round(cy - L * np.sin(ang_rad)))  # y invertita immagine
        cv2.arrowedLine(out, (cx, cy), (ex, ey), color, 2, cv2.LINE_AA, tipLength=0.2)
        # Etichetta
        label = f"#{i+1} {m.angle_deg:.0f}d s={m.scale:.2f} {m.score:.2f}"
        cv2.putText(out, label, (cx + 8, cy - 8),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1, cv2.LINE_AA)
    return out


def _color_for(i: int) -> tuple[int, int, int]:
    palette = [
        (0, 255, 0), (0, 200, 255), (255, 100, 100),
        (255, 200, 0), (200, 0, 255), (100, 255, 200),
        (255, 0, 0), (0, 255, 255),
    ]
    return palette[i % len(palette)]


def show_results(scene: np.ndarray, matches: list[Match]) -> None:
    print(f"\n=== {len(matches)} match trovati ===")
    for i, m in enumerate(matches):
        print(f"  #{i+1}: cx={m.cx:.1f} cy={m.cy:.1f} "
              f"angle={m.angle_deg:.1f}d scale={m.scale:.2f} score={m.score:.3f}")
    overlay = draw_matches(scene, matches)
    disp = _fit_for_display(overlay, max_side=1400)
    cv2.imshow(WINDOW_RESULT, disp)
    print("\nPremere un tasto sulla finestra per chiudere.")
    cv2.waitKey(0)
    cv2.destroyAllWindows()


def run(
    initial_dir: str | None = None,
    angle_step_deg: float = 5.0,
    angle_range_deg: tuple[float, float] = (0.0, 360.0),
    scale_range: tuple[float, float] = (1.0, 1.0),
    scale_step: float = 0.1,
    num_features: int = 96,
    weak_grad: float = 30.0,
    strong_grad: float = 60.0,
    spread_radius: int = 5,
    pyramid_levels: int = 3,
    min_score: float = 0.55,
    max_matches: int = 25,
    backend: str = "line",
) -> None:
    """Entry-point GUI completo."""
    print("[1/4] Selezionare immagine MODELLO...")
    model_path = pick_file("Immagine MODELLO", initialdir=initial_dir)
    if not model_path:
        print("Annullato."); return
    model_img = load_image(model_path)
    print(f"      caricato: {model_path}  shape={model_img.shape}")

    print("[2/4] Selezionare ROI sul modello (trascinare, INVIO conferma).")
    roi = select_roi(model_img)
    if roi is None:
        print("ROI vuota, annullato."); return
    print(f"      ROI: {roi.shape[1]}x{roi.shape[0]} px")

    print("[3/4] Selezionare immagine SCENA...")
    scene_path = pick_file("Immagine SCENA",
                           initialdir=str(Path(model_path).parent))
    if not scene_path:
        print("Annullato."); return
    scene = load_image(scene_path)
    print(f"      caricato: {scene_path}  shape={scene.shape}")

    print(f"[4/4] Train + match (backend={backend})...")
    if backend == "edge":
        matcher: EdgeShapeMatcher | LineShapeMatcher = EdgeShapeMatcher(
            angle_step_deg=angle_step_deg, angle_range_deg=angle_range_deg,
            scale_range=scale_range, scale_step=scale_step,
        )
    else:
        matcher = LineShapeMatcher(
            num_features=num_features,
            weak_grad=weak_grad, strong_grad=strong_grad,
            angle_step_deg=angle_step_deg, angle_range_deg=angle_range_deg,
            scale_range=scale_range, scale_step=scale_step,
            spread_radius=spread_radius, pyramid_levels=pyramid_levels,
        )
    import time
    t0 = time.time()
    n = matcher.train(roi)
    print(f"      train: {n} varianti in {time.time()-t0:.2f}s")
    t0 = time.time()
    matches = matcher.find(scene, min_score=min_score, max_matches=max_matches)
    print(f"      find:  {len(matches)} match in {time.time()-t0:.2f}s")
    show_results(scene, matches)


if __name__ == "__main__":
    test_dir = "/home/adriano/Documenti/Git_XYZ/VisionSuite/Shape_model_2d/Test"
    run(initial_dir=test_dir if Path(test_dir).is_dir() else None)