calibration_overlay.py

import threading
import time
import tkinter as tk
from typing import Callable, Dict, List, Optional, Tuple


PointDef = Dict[str, object]
Sample = Dict[str, object]


class CalibrationOverlay:
    """Fullscreen transparent overlay to collect gaze samples at fixed targets."""

    DEFAULT_POINTS: List[PointDef] = [
        {"name": "center", "pos": (0.5, 0.5), "radius": 18},
        {"name": "left", "pos": (0.2, 0.5), "radius": 16},
        {"name": "right", "pos": (0.8, 0.5), "radius": 16},
        {"name": "top", "pos": (0.5, 0.25), "radius": 16},
        {"name": "bottom", "pos": (0.5, 0.75), "radius": 16},
        {"name": "topleft", "pos": (0.15, 0.25), "radius": 14},
        {"name": "topright", "pos": (0.85, 0.25), "radius": 14},
        {"name": "bottomleft", "pos": (0.15, 0.75), "radius": 14},
        {"name": "bottomright", "pos": (0.85, 0.75), "radius": 14},
    ]

    def __init__(
        self,
        screen_size: Tuple[int, int],
        sample_getter: Callable[[], Optional[Sample]],
        on_complete: Callable[[Optional[List[Dict[str, object]]]], None],
        points: Optional[List[PointDef]] = None,
    ) -> None:
        self.screen_w, self.screen_h = screen_size
        self.sample_getter = sample_getter
        self.on_complete = on_complete
        self.points = points or self.DEFAULT_POINTS

        self._thread: Optional[threading.Thread] = None
        self._running = False
        self._root: Optional[tk.Tk] = None
        self._canvas: Optional[tk.Canvas] = None
        self._current_idx = 0
        self._records: List[Dict[str, object]] = []
        self._bg_color = "#010203"  # a near-black color to serve as transparent key

    # -------- Public API --------
    def start(self) -> None:
        if self._running:
            return
        self._running = True
        self._thread = threading.Thread(target=self._run_loop, name="CalibrationOverlay", daemon=True)
        self._thread.start()

    # -------- Internal helpers --------
    def _run_loop(self) -> None:
        try:
            self._root = tk.Tk()
        except tk.TclError as exc:
            print(f"[Calibration Overlay] Failed to create Tk window: {exc}")
            self._running = False
            self.on_complete(None)
            return

        root = self._root
        root.withdraw()
        root.overrideredirect(True)
        root.attributes("-topmost", True)
        try:
            root.attributes("-transparentcolor", self._bg_color)
        except tk.TclError:
            # Fallback to alpha transparency if transparentcolor unsupported
            root.attributes("-alpha", 0.3)
            self._bg_color = "#000000"
        root.geometry(f"{self.screen_w}x{self.screen_h}+0+0")

        self._canvas = tk.Canvas(root, width=self.screen_w, height=self.screen_h,
                                 bg=self._bg_color, highlightthickness=0)
        self._canvas.pack(fill="both", expand=True)

        root.bind("<space>", self._capture_point)
        root.bind("<Return>", self._capture_point)
        root.bind("<Escape>", self._cancel)
        root.bind("<BackSpace>", self._undo)

        root.after(10, self._prepare_canvas)
        root.deiconify()
        root.mainloop()
        self._running = False

    def _prepare_canvas(self) -> None:
        if not self._canvas:
            return
        self._canvas.delete("all")
        for idx, point in enumerate(self.points):
            x_norm, y_norm = point["pos"]
            radius = point.get("radius", 16)
            x = x_norm * self.screen_w
            y = y_norm * self.screen_h
            fill = "#ffff00" if idx == self._current_idx else "#ff8800"
            outline = "#ffffff" if idx == self._current_idx else "#444444"
            self._canvas.create_oval(
                x - radius, y - radius, x + radius, y + radius,
                fill=fill, outline=outline, width=3,
                tags=(f"point_{idx}",)
            )
        self._draw_instruction()

    def _draw_instruction(self) -> None:
        if not self._canvas:
            return
        message = (
            "Nhìn vào chấm vàng, nhấn SPACE để ghi mẫu."
            "\nESC huỷ, BACKSPACE quay lại."
        )
        self._canvas.create_text(
            self.screen_w / 2,
            self.screen_h * 0.08,
            text=message,
            fill="#ffffff",
            font=("Segoe UI", 24, "bold"),
            tags=("instruction",)
        )

    def _advance_point(self) -> None:
        if not self._canvas:
            return
        self._current_idx += 1
        if self._current_idx >= len(self.points):
            self._finish()
            return
        self._prepare_canvas()

    def _capture_point(self, _event=None) -> None:
        if self._current_idx >= len(self.points):
            return
        sample = self.sample_getter()
        if not sample:
            self._flash_message("Không có dữ liệu mắt ở thời điểm này!", duration_ms=900)
            return
        point = self.points[self._current_idx]
        record = {
            "point": point["name"],
            "target_norm": point["pos"],
            "captured_at": time.time(),
            "sample": sample,
        }
        self._records.append(record)
        self._advance_point()

    def _undo(self, _event=None) -> None:
        if not self._records or self._current_idx == 0:
            return
        self._records.pop()
        self._current_idx = max(0, self._current_idx - 1)
        self._prepare_canvas()

    def _cancel(self, _event=None) -> None:
        self._teardown()
        self.on_complete(None)

    def _finish(self) -> None:
        records = list(self._records)
        self._teardown()
        self.on_complete(records)

    def _flash_message(self, text: str, duration_ms: int = 600) -> None:
        if not self._canvas:
            return
        tag = "flash"
        self._canvas.delete(tag)
        self._canvas.create_text(
            self.screen_w / 2,
            self.screen_h * 0.92,
            text=text,
            fill="#ff6666",
            font=("Segoe UI", 24, "bold"),
            tags=(tag,)
        )
        self._canvas.after(duration_ms, lambda: self._canvas and self._canvas.delete(tag))

    def _teardown(self) -> None:
        if self._root:
            try:
                self._root.destroy()
            except tk.TclError:
                pass
            self._root = None
        self._canvas = None
        self._running = False


__all__ = ["CalibrationOverlay"]