signspeak/live_debug.py

from __future__ import annotations

import os
from typing import Any

import numpy as np

from .asl.asl_detector import ASLDetector
from .asl.mediapipe_utils import draw_holistic_landmarks, extract_keypoints_from_holistic


class LiveASLSession:
    def __init__(self) -> None:
        self.asl = ASLDetector()
        self.interpreter: Any | None = None
        self.frame_keypoints: list[np.ndarray] = []
        self.latest_prediction = ""
        self.latest_top_candidate = ""
        self.latest_top_predictions: list[dict[str, Any]] = []
        self.latest_emotion = "unknown"
        self.min_prediction_frames = max(1, int(os.getenv("LIVE_ASL_MIN_FRAMES", "4")))
        self.max_prediction_frames = max(self.min_prediction_frames, int(os.getenv("LIVE_ASL_MAX_FRAMES", "12")))
        self.predict_every = max(1, int(os.getenv("LIVE_ASL_PREDICT_EVERY", "1")))
        self.emotion_every = max(1, int(os.getenv("LIVE_EMOTION_EVERY", "45")))
        self.latest_status = f"Waiting for live frames: 0/{self.min_prediction_frames}."
        self.frames_seen = 0
        self.mp: Any | None = None
        self.holistic: Any | None = None
        self._load_mediapipe()

    def process_frame(self, frame_rgb: np.ndarray) -> tuple[np.ndarray, str]:
        output = frame_rgb.copy()
        self.frames_seen += 1
        if self.mp is None or self.holistic is None:
            self.latest_status = "MediaPipe unavailable."
            return self._draw(output), self.latest_status

        try:
            results = self.holistic.process(frame_rgb)
            draw_holistic_landmarks(self.mp, output, results)
            keypoints = extract_keypoints_from_holistic(results, missing_value=np.nan)
            self.frame_keypoints.append(keypoints)
            self.frame_keypoints = self.frame_keypoints[-self.max_prediction_frames :]
            if self.frames_seen % self.emotion_every == 0:
                self._update_emotion(frame_rgb)
            if len(self.frame_keypoints) < self.min_prediction_frames:
                self.latest_status = f"Waiting for live frames: {len(self.frame_keypoints)}/{self.min_prediction_frames}."
            elif self.frames_seen % self.predict_every == 0:
                self._predict_latest()
        except Exception as exc:
            self.latest_status = f"Live ASL error: {type(exc).__name__}: {exc}"

        return self._draw(output), self._status_text()

    def _load_mediapipe(self) -> None:
        try:
            import mediapipe as mp

            self.mp = mp
            self.holistic = mp.solutions.holistic.Holistic(
                model_complexity=0,
                min_detection_confidence=0.5,
                min_tracking_confidence=0.5,
            )
        except Exception as exc:
            self.latest_status = f"MediaPipe unavailable: {type(exc).__name__}: {exc}"

    def _predict_latest(self) -> None:
        if not self.asl.model_path.exists():
            self.latest_prediction = ""
            self.latest_status = "ASL model missing."
            return

        if self.interpreter is None:
            self.interpreter = self.asl._load_interpreter()

        keypoints = np.asarray(self.frame_keypoints, dtype=np.float32)
        output = self.asl._predict(self.interpreter, keypoints)
        probs = self.asl._softmax_if_needed(np.asarray(output).reshape(-1))
        top_idx = int(np.argmax(probs))
        label = self.asl._label_for_index(top_idx)
        confidence = float(probs[top_idx])
        top_predictions = self.asl._top_predictions(probs, limit=3)
        self.latest_top_predictions = top_predictions
        self.latest_top_candidate = f"{label} ({confidence:.0%})"

        if confidence >= self.asl.confidence_threshold:
            self.latest_prediction = f"{label} ({confidence:.0%})"
            self.latest_status = (
                f"Accepted: {label} at {confidence:.2f} "
                f"with {len(self.frame_keypoints)} live frames."
            )
        else:
            self.latest_prediction = ""
            self.latest_status = (
                f"Top candidate: {label} at {confidence:.2f}; "
                f"waiting for {self.asl.confidence_threshold:.2f}."
            )

    def _update_emotion(self, frame_rgb: np.ndarray) -> None:
        try:
            os.environ.setdefault("TF_USE_LEGACY_KERAS", "1")
            import cv2
            from deepface import DeepFace

            frame_bgr = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)
            analysis = DeepFace.analyze(
                img_path=frame_bgr,
                actions=["emotion"],
                enforce_detection=False,
                silent=True,
            )
            if isinstance(analysis, list):
                analysis = analysis[0] if analysis else {}
            dominant = analysis.get("dominant_emotion") if isinstance(analysis, dict) else None
            if dominant:
                self.latest_emotion = str(dominant)
        except Exception:
            self.latest_emotion = "unavailable"

    def _draw(self, frame_rgb: np.ndarray) -> np.ndarray:
        import cv2

        height, width = frame_rgb.shape[:2]
        cv2.rectangle(frame_rgb, (0, 0), (width, 146), (8, 11, 16), -1)
        cv2.putText(
            frame_rgb,
            f"Sign: {self.latest_prediction or '-'}",
            (14, 36),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.86,
            (45, 212, 191),
            2,
            cv2.LINE_AA,
        )
        cv2.putText(
            frame_rgb,
            f"Top: {self.latest_top_candidate or '-'}",
            (14, 68),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.58,
            (129, 140, 248),
            2,
            cv2.LINE_AA,
        )
        cv2.putText(
            frame_rgb,
            f"Emotion: {self.latest_emotion}",
            (14, 98),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.58,
            (245, 158, 11),
            2,
            cv2.LINE_AA,
        )
        cv2.putText(
            frame_rgb,
            self.latest_status[:96],
            (14, 128),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.48,
            (248, 250, 252),
            1,
            cv2.LINE_AA,
        )
        return frame_rgb

    def _status_text(self) -> str:
        top_lines = [
            f"- {item.get('label')}: {float(item.get('confidence', 0.0) or 0.0):.2f}"
            for item in self.latest_top_predictions
        ]
        top_block = "\n".join(top_lines) if top_lines else "- None yet"
        return (
            f"{self.latest_status}\n"
            f"Accepted sign: {self.latest_prediction or 'None'}\n"
            f"Top candidates:\n{top_block}\n"
            f"Frames in rolling buffer: {len(self.frame_keypoints)}/{self.max_prediction_frames}\n"
            f"Acceptance threshold: {self.asl.confidence_threshold:.2f}\n"
            f"Emotion: {self.latest_emotion}"
        )


_live_session: LiveASLSession | None = None


def process_live_debug_frame(frame: np.ndarray | None) -> tuple[np.ndarray | None, str]:
    global _live_session

    if frame is None:
        return None, "Waiting for camera frame."

    if _live_session is None:
        _live_session = LiveASLSession()

    return _live_session.process_frame(frame)