spatial-computing.py

from __future__ import annotations

"""Camera-based eye cursor and hand gesture demo for macOS."""

import argparse
from collections import deque
import json
import math
from pathlib import Path
import sys
import time

import cv2
import mediapipe as mp
import pyautogui

from gaze_features import AdaptiveCursorFilter, GazeModelPredictor, extract_gaze_features
from spatial_core import (
    AppConfig,
    CalibrationSession,
    GestureEvent,
    GestureEventType,
    GestureRecognizer,
    apply_poly,
    clamp_to_screen,
    desktop_direction_for_action,
    load_calibration,
    scroll_amount_for_action,
)

try:
    from pynput import keyboard

    HAS_PYNPUT = True
except ImportError:
    HAS_PYNPUT = False


CALIBRATION_PATH = Path(__file__).with_name("gaze_calib_poly.json")
WINDOW_NAME = "Spatial Computing Demo"


def require_legacy_mediapipe_api() -> None:
    if not hasattr(mp, "solutions"):
        raise RuntimeError(
            "This project requires MediaPipe 0.10.21 and Python 3.12. "
            "Create the documented virtual environment and install requirements.txt."
        )


class SystemActions:
    def __init__(self, config: AppConfig) -> None:
        self.config = config
        pyautogui.PAUSE = 0
        pyautogui.FAILSAFE = True

    def move(self, x: float, y: float) -> bool:
        try:
            pyautogui.moveTo(x, y, _pause=False)
            return True
        except pyautogui.FailSafeException:
            return False

    def perform(self, event: GestureEvent) -> bool:
        try:
            if event.kind == GestureEventType.CLICK:
                pyautogui.click()
            elif event.kind == GestureEventType.SCROLL:
                pyautogui.scroll(scroll_amount_for_action(event.amount, self.config))
            elif event.kind in (GestureEventType.SWIPE_LEFT, GestureEventType.SWIPE_RIGHT):
                if sys.platform == "darwin":
                    direction = desktop_direction_for_action(event.kind, self.config)
                    pyautogui.hotkey("ctrl", direction)
            return True
        except pyautogui.FailSafeException:
            return False


class DemoActions(SystemActions):
    def __init__(self, config: AppConfig) -> None:
        self.config = config
        self.cursor = (0.0, 0.0)

    def move(self, x: float, y: float) -> bool:
        self.cursor = (x, y)
        return True

    def perform(self, event: GestureEvent) -> bool:
        return True


class SpatialComputingApp:
    def __init__(
        self, demo_mode: bool = False, camera: int = 0, gaze_model: Path | None = None
    ) -> None:
        require_legacy_mediapipe_api()
        self.config = AppConfig()
        self.screen_size = pyautogui.size()
        self.actions = DemoActions(self.config) if demo_mode else SystemActions(self.config)
        self.demo_mode = demo_mode
        self.gaze_predictor = GazeModelPredictor(gaze_model) if gaze_model else None
        self.cursor_filter = AdaptiveCursorFilter()
        self.control_enabled = True
        self.theta = load_calibration(CALIBRATION_PATH)
        self.calibration = CalibrationSession(self.screen_size, self.config)
        self.gaze_history: deque[tuple[float, float]] = deque(
            maxlen=self.config.calibration_samples
        )
        self.gesture = GestureRecognizer(self.config)
        self.filtered_cursor: tuple[float, float] | None = None
        self.latest_gaze: tuple[float, float] | None = None
        self.calibration_score: float | None = None
        self.toast = ""
        self.toast_until = 0.0
        self.intro = True
        self.calibration_pinch = False
        self._pressed_keys: set[object] = set()

        self.face = mp.solutions.face_mesh.FaceMesh(max_num_faces=1, refine_landmarks=True)
        self.hands = mp.solutions.hands.Hands(
            static_image_mode=False,
            max_num_hands=1,
            min_detection_confidence=0.5,
            min_tracking_confidence=0.5,
        )
        self.cap = cv2.VideoCapture(camera)
        if not self.cap.isOpened():
            self.face.close()
            self.hands.close()
            raise RuntimeError(f"Cannot open camera {camera}")

        self.listener = None
        if HAS_PYNPUT:
            self.listener = keyboard.Listener(on_press=self._on_press, on_release=self._on_release)
            self.listener.start()

    def run(self) -> None:
        previous_frame_at = time.monotonic()
        fps = 0.0
        try:
            while True:
                ok, frame = self.cap.read()
                if not ok:
                    raise RuntimeError("Camera stopped returning frames")
                frame = cv2.flip(frame, 1)
                now = time.monotonic()
                elapsed = max(now - previous_frame_at, 0.0001)
                previous_frame_at = now
                fps = fps * 0.9 + (1 / elapsed) * 0.1
                key = self._process_frame(frame, now, fps)
                if key == ord("q"):
                    break
                if key in (13, ord(" ")) and self.intro:
                    self.intro = False
                    if not self._tracker_ready():
                        self._start_recalibration()
                elif key == ord("r"):
                    self._start_recalibration()
                elif key == ord(" ") and self.theta is None:
                    self._capture_calibration()
        finally:
            self.cleanup()

    def _process_frame(self, frame, now: float, fps: float) -> int:
        height, width = frame.shape[:2]
        rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        face_result = self.face.process(rgb)
        hand_result = self.hands.process(rgb)

        if face_result.multi_face_landmarks:
            self._process_face(frame, face_result.multi_face_landmarks[0], width, height)

        pinch_distance, pinch_point = self._process_hand(frame, hand_result, width, height)
        if not self.intro:
            if not self._tracker_ready():
                pinching = pinch_distance is not None and pinch_distance < self.config.pinch_threshold_px
                if pinching and not self.calibration_pinch:
                    self._capture_calibration()
                self.calibration_pinch = pinching
            else:
                for event in self.gesture.update(now, pinch_distance, pinch_point):
                    self._handle_event(event, now)

        self._draw_overlay(frame, now, fps, pinch_point)
        cv2.imshow(WINDOW_NAME, frame)
        return cv2.waitKey(1) & 0xFF

    def _process_face(self, frame, face_landmarks, width: int, height: int) -> None:
        landmarks = face_landmarks.landmark
        if self.config.show_face_dots:
            for point in landmarks[:468]:
                cv2.circle(frame, (int(point.x * width), int(point.y * height)), 1, (255, 255, 255), -1)

        gaze = ((landmarks[473].x + landmarks[468].x) / 2, (landmarks[473].y + landmarks[468].y) / 2)
        self.latest_gaze = gaze
        self.gaze_history.append(gaze)
        cv2.circle(frame, (int(gaze[0] * width), int(gaze[1] * height)), 5, (80, 255, 120), -1)

        model_prediction = None
        if self.gaze_predictor is not None:
            features = extract_gaze_features(frame, landmarks)
            if features is not None:
                normalized, confidence = self.gaze_predictor.predict(features)
                if confidence >= self.config.gaze_model_confidence:
                    model_prediction = (
                        normalized[0] * self.screen_size[0],
                        normalized[1] * self.screen_size[1],
                    )

        ready = model_prediction is not None or (
            self.gaze_predictor is None and self.theta is not None
        )
        if ready and self.control_enabled and not self.intro:
            raw_cursor = model_prediction or apply_poly(self.theta, *gaze)
            cursor = clamp_to_screen(
                *raw_cursor,
                *self.screen_size,
                margin=self.config.cursor_safe_margin_px,
            )
            if self.gaze_predictor is not None:
                self.filtered_cursor = self.cursor_filter.update(cursor)
            elif self.filtered_cursor is None:
                self.filtered_cursor = cursor
            else:
                alpha = self.config.smooth_alpha
                self.filtered_cursor = (
                    alpha * cursor[0] + (1 - alpha) * self.filtered_cursor[0],
                    alpha * cursor[1] + (1 - alpha) * self.filtered_cursor[1],
                )
            if not self.actions.move(*self.filtered_cursor):
                self.control_enabled = False
                self.filtered_cursor = None
                self._show_toast("FAIL-SAFE: CONTROL PAUSED", time.monotonic(), 3)

    def _process_hand(self, frame, result, width: int, height: int):
        if not result.multi_hand_landmarks:
            return None, None
        hand = result.multi_hand_landmarks[0]
        mp.solutions.drawing_utils.draw_landmarks(frame, hand, mp.solutions.hands.HAND_CONNECTIONS)
        thumb, index = hand.landmark[4], hand.landmark[8]
        point = (index.x * width, index.y * height)
        distance = math.hypot((thumb.x - index.x) * width, (thumb.y - index.y) * height)
        return distance, point

    def _handle_event(self, event: GestureEvent, now: float) -> None:
        labels = {
            GestureEventType.CLICK: "CLICK",
            GestureEventType.SCROLL: "SCROLL INVERTED",
            GestureEventType.TOGGLE_CONTROL: "CONTROL TOGGLED",
        }
        if event.kind in (GestureEventType.SWIPE_LEFT, GestureEventType.SWIPE_RIGHT):
            direction = desktop_direction_for_action(event.kind, self.config)
            labels[event.kind] = f"DESKTOP {direction.upper()}"
        if event.kind == GestureEventType.TOGGLE_CONTROL:
            self.control_enabled = not self.control_enabled
            self.filtered_cursor = None
            self.cursor_filter.reset()
            self._show_toast("CONTROL RESUMED" if self.control_enabled else "CONTROL PAUSED", now)
        elif self.control_enabled:
            if self.actions.perform(event):
                self._show_toast(labels[event.kind], now)
            else:
                self.control_enabled = False
                self.filtered_cursor = None
                self._show_toast("FAIL-SAFE: CONTROL PAUSED", now, 3)

    def _capture_calibration(self) -> None:
        if self.calibration.capture(self.gaze_history):
            self.gaze_history.clear()
            if self.calibration.complete:
                self.theta, self.calibration_score = self.calibration.fit()
                try:
                    CALIBRATION_PATH.write_text(json.dumps(self.theta))
                except OSError:
                    self._show_toast("CALIBRATED - COULD NOT SAVE", time.monotonic(), 3)
                    return
                self.filtered_cursor = None
                self._show_toast(f"CALIBRATED - ERROR {self.calibration_score:.0f}px", time.monotonic(), 3)

    def _tracker_ready(self) -> bool:
        return self.gaze_predictor is not None or self.theta is not None

    def _start_recalibration(self) -> None:
        if self.gaze_predictor is not None:
            self.filtered_cursor = None
            self.cursor_filter.reset()
            self._show_toast("ML TRACKER FILTER RESET", time.monotonic())
            return
        self.intro = False
        self.theta = None
        self.calibration = CalibrationSession(self.screen_size, self.config)
        self.gaze_history.clear()
        self.gesture.reset()
        self.filtered_cursor = None
        self.cursor_filter.reset()
        self.calibration_pinch = False
        self._show_toast("RECALIBRATION STARTED", time.monotonic())

    def _draw_overlay(self, frame, now: float, fps: float, pinch_point) -> None:
        height, width = frame.shape[:2]
        overlay = frame.copy()
        cv2.rectangle(overlay, (0, 0), (width, 92), (8, 12, 22), -1)
        cv2.addWeighted(overlay, 0.78, frame, 0.22, 0, frame)
        mode = "DEMO MODE" if self.demo_mode else "LIVE CONTROL"
        if self.gaze_predictor is not None:
            mode += " - ML GAZE"
        status = "PAUSED" if not self.control_enabled else mode
        self._text(frame, status, (24, 34), 0.75, (80, 210, 255))
        self._text(frame, f"{fps:.0f} FPS", (width - 110, 32), 0.55, (220, 220, 220))
        self._text(frame, "Q quit   R recalibrate   double-pinch pause/resume", (24, 70), 0.48, (190, 195, 205))

        if pinch_point and self.gesture.pinching:
            cv2.circle(frame, (int(pinch_point[0]), int(pinch_point[1])), 24, (255, 190, 40), 3)

        if self.intro:
            self._draw_intro(frame)
        elif not self._tracker_ready():
            self._draw_calibration(frame)

        if self.toast and now < self.toast_until:
            size = cv2.getTextSize(self.toast, cv2.FONT_HERSHEY_SIMPLEX, 0.8, 2)[0]
            x = (width - size[0]) // 2
            cv2.rectangle(frame, (x - 18, height - 78), (x + size[0] + 18, height - 32), (20, 25, 35), -1)
            self._text(frame, self.toast, (x, height - 46), 0.8, (80, 240, 180), 2)

    def _draw_intro(self, frame) -> None:
        height, width = frame.shape[:2]
        overlay = frame.copy()
        cv2.rectangle(overlay, (width // 8, height // 5), (width * 7 // 8, height * 4 // 5), (8, 12, 22), -1)
        cv2.addWeighted(overlay, 0.88, frame, 0.12, 0, frame)
        x, y = width // 8 + 35, height // 5 + 55
        self._text(frame, "SPATIAL COMPUTING", (x, y), 1.05, (80, 240, 180), 2)
        lines = [
            "Look to move the cursor",
            "Pinch: click   Hold + vertical: scroll",
            "Hold + horizontal: switch desktop",
            "Double-pinch: pause or resume",
            "macOS live control requires Camera + Accessibility permissions",
            "Press SPACE or ENTER to start",
        ]
        for offset, line in enumerate(lines, start=1):
            color = (255, 210, 90) if offset == len(lines) else (225, 228, 235)
            self._text(frame, line, (x, y + 48 * offset), 0.56, color)

    def _draw_calibration(self, frame) -> None:
        height, width = frame.shape[:2]
        target = self.calibration.current_target
        if target is None:
            return
        x = int(target[0] / self.screen_size[0] * width)
        y = int(target[1] / self.screen_size[1] * height)
        cv2.circle(frame, (x, y), 18, (60, 60, 255), 3)
        cv2.circle(frame, (x, y), 4, (255, 255, 255), -1)
        progress = f"CALIBRATION {self.calibration.index + 1}/{len(self.calibration.targets)}"
        self._text(frame, progress, (24, 118), 0.7, (80, 240, 180), 2)
        self._text(frame, "Look at the target, then pinch or press SPACE", (24, 150), 0.48, (230, 230, 230))

    @staticmethod
    def _text(frame, text: str, origin, scale: float, color, thickness: int = 1) -> None:
        cv2.putText(frame, text, origin, cv2.FONT_HERSHEY_SIMPLEX, scale, color, thickness, cv2.LINE_AA)

    def _show_toast(self, message: str, now: float, duration: float = 1.0) -> None:
        self.toast = message
        self.toast_until = now + duration

    def _on_press(self, key) -> None:
        self._pressed_keys.add(key)
        cmd = getattr(keyboard.Key, "cmd", None)
        shift = getattr(keyboard.Key, "shift", None)
        if cmd in self._pressed_keys and shift in self._pressed_keys:
            char = getattr(key, "char", "")
            if char and char.lower() == "r":
                self._start_recalibration()

    def _on_release(self, key) -> None:
        self._pressed_keys.discard(key)

    def cleanup(self) -> None:
        self.cap.release()
        self.face.close()
        self.hands.close()
        if self.listener:
            self.listener.stop()
        cv2.destroyAllWindows()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--demo", action="store_true", help="Show gestures without controlling macOS")
    parser.add_argument("--camera", type=int, default=0, help="Camera device index")
    parser.add_argument("--gaze-model", type=Path, help="TorchScript personal gaze model")
    return parser.parse_args()


if __name__ == "__main__":
    arguments = parse_args()
    try:
        SpatialComputingApp(
            demo_mode=arguments.demo,
            camera=arguments.camera,
            gaze_model=arguments.gaze_model,
        ).run()
    except KeyboardInterrupt:
        print("\nInterrupted.")
    except Exception as error:
        print(f"Error: {error}", file=sys.stderr)
        raise SystemExit(1) from None