emotion.py

import time
import threading
from dataclasses import dataclass, field
from typing import Optional, List, Tuple, Dict

import gradio as gr
import numpy as np
import cv2

import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision

from deepface import DeepFace


# =========================
# 參數（可微調）
# =========================
ANALYZE_EVERY_SEC = 1.5          # ⭐ 建議 1.5~2.0：HF CPU 比較順
EMA_ALPHA = 0.45
CONF_TH = 40.0                   # DeepFace emotion confidence 門檻
SWITCH_CONFIRM_SEC = 2.5         # 連續穩定多久才算完成
ZOOM_FACTOR = 1.5                # ✅ 保留：只用在「顯示」
UI_TARGET_W = 720                # ✅ 回傳畫面縮小寬度，提升順暢度

# MediaPipe face detector 門檻（太高很容易漏）
MP_MIN_DET_CONF = 0.3

# 模型路徑（確認 repo 中模型真的在這裡）
MP_MODEL_PATH = "blaze_face_short_range.tflite"


# =========================
# 情緒中文 + 分數映射（7 -> 5）
# =========================
EMO_ZH = {
    "angry": "生氣",
    "disgust": "厭惡",
    "fear": "害怕",
    "happy": "開心",
    "sad": "難過",
    "surprise": "驚訝",
    "neutral": "平靜",
    "unknown": "未知",
}

EMO_TO_SCORE = {
    "angry": 1,
    "disgust": 1,
    "fear": 1,
    "sad": 2,
    "neutral": 3,
    "surprise": 4,
    "happy": 5,
    "unknown": 3,
}

SCORE_LABELS = {
    1: "1（心情差）",
    2: "2（不太好）",
    3: "3（普通）",
    4: "4（不錯）",
    5: "5（超棒）",
}
SCORE_CHOICES = [SCORE_LABELS[i] for i in [1, 2, 3, 4, 5]]


# =========================
# 狀態
# =========================
@dataclass
class AppState:
    mode: str = "攝影機辨識"

    running: bool = True
    finished: bool = False
    last_analyze_t: float = 0.0

    # smoothing
    ema_scores: Optional[np.ndarray] = None
    emo_keys: List[str] = field(default_factory=list)

    # stable/candidate
    stable_emo: str = "unknown"
    candidate_emo: Optional[str] = None
    candidate_since: Optional[float] = None

    # final
    final_emo: Optional[str] = None
    final_conf: Optional[float] = None
    final_score: int = 3


# =========================
# UI HTML
# =========================
def _hint_html(msg: str) -> str:
    return f"""
    <div style="border-radius:14px;padding:14px;border:1px dashed #DADADA;background:#FAFAFA;color:#000000;">
      {msg}
    </div>
    """

def _camera_hint_before_start() -> str:
    return _hint_html(
        "請先按「Click to access webcam / 允許」開啟攝影機；"
        "接著，請按「錄製（Record）」開始。"
    )

def _camera_hint_need_face() -> str:
    return _hint_html(
        "⚠️ 目前未偵測到人臉，請將臉部移至畫面中央、靠近鏡頭，並避免瀏海遮住眉眼，才可進行情緒辨識。"
    )

def _camera_hint_done() -> str:
    return _hint_html("已完成辨識，可按「重新攝影機辨識」再次進行辨識。")

def _camera_hint_stopped() -> str:
    return _hint_html("已停止辨識。可按「重新攝影機辨識」再試一次。")

def _score_radio_value(score: int) -> str:
    score = int(score)
    score = min(5, max(1, score))
    return SCORE_LABELS[score]

def _result_card_html(emo_key: str, conf: Optional[float]) -> str:
    zh = EMO_ZH.get(emo_key, emo_key)
    score = EMO_TO_SCORE.get(emo_key, 3)
    conf_txt = "" if conf is None else f"{conf:.1f}%"
    score_desc = SCORE_LABELS[score].split("（")[1].rstrip("）")

    return f"""
    <div style="
        border-radius:16px;
        padding:20px;
        border:1px solid #E6E6E6;
        background:#FFFFFF;
        color:#000000;
    ">
      <div style="font-size:14px;color:#000000;margin-bottom:8px;">
        辨識結果
      </div>

      <div style="
        font-size:42px;
        font-weight:800;
        line-height:1.1;
        margin-bottom:14px;
        color:#000000;
      ">
        {zh}
      </div>

      <div style="font-size:16px;color:#000000;margin-bottom:6px;">
        心情分數：{score}（{score_desc}）
      </div>

      <div style="font-size:14px;color:#000000;">
        信心值：{conf_txt}
      </div>
    </div>
    """


# =========================
# 影像處理（保命 + 顯示優化）
# =========================
def _to_rgb_uint8(frame: np.ndarray) -> np.ndarray:
    """統一成 MediaPipe Tasks 最穩格式：RGB、uint8、0~255、contiguous。"""
    if frame is None:
        return frame

    # RGBA -> RGB
    if frame.ndim == 3 and frame.shape[2] == 4:
        frame = frame[:, :, :3]

    # 灰階 -> RGB
    if frame.ndim == 2:
        frame = np.stack([frame, frame, frame], axis=-1)

    # dtype / range 修正
    if frame.dtype != np.uint8:
        mx = float(np.max(frame)) if frame.size else 0.0
        if mx <= 1.5:  # 多半是 0~1
            frame = (frame * 255.0).clip(0, 255).astype(np.uint8)
        else:          # 多半是 0~255 但 float
            frame = np.clip(frame, 0, 255).astype(np.uint8)

    return np.ascontiguousarray(frame)

def _zoom_center(frame: np.ndarray, zoom: float) -> np.ndarray:
    if frame is None or zoom <= 1.0:
        return frame
    h, w = frame.shape[:2]
    new_w = int(w / zoom)
    new_h = int(h / zoom)
    start_x = (w - new_w) // 2
    start_y = (h - new_h) // 2
    cropped = frame[start_y : start_y + new_h, start_x : start_x + new_w]
    return np.ascontiguousarray(cropped)

def _resize_for_ui(frame_rgb: np.ndarray, target_w: int = UI_TARGET_W) -> np.ndarray:
    """回傳給 Gradio 的畫面縮到固定寬度，減少傳輸/渲染負擔。"""
    if frame_rgb is None:
        return frame_rgb
    h, w = frame_rgb.shape[:2]
    if w <= target_w:
        return frame_rgb
    scale = target_w / w
    new_h = max(1, int(h * scale))
    return cv2.resize(frame_rgb, (target_w, new_h), interpolation=cv2.INTER_AREA)


# =========================
# MediaPipe Tasks Face Detector
# =========================
BaseOptions = python.BaseOptions
FaceDetector = vision.FaceDetector
FaceDetectorOptions = vision.FaceDetectorOptions
VisionRunningMode = vision.RunningMode

_face_detector = FaceDetector.create_from_options(
    FaceDetectorOptions(
        base_options=BaseOptions(model_asset_path=MP_MODEL_PATH),
        running_mode=VisionRunningMode.IMAGE,
        min_detection_confidence=MP_MIN_DET_CONF,
    )
)

def mp_tasks_detect_and_draw(frame_rgb: np.ndarray):
    """
    - 最大臉策略（bbox 面積最大）
    - 有臉：畫白框 + 紅點 keypoints
    - 回 bbox_px (x, y, w, h) 供 ROI 裁切
    """
    frame_rgb = _to_rgb_uint8(frame_rgb)
    h, w = frame_rgb.shape[:2]

    mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=frame_rgb)
    result = _face_detector.detect(mp_image)

    if not result.detections:
        return frame_rgb, None

    best_det = None
    best_area = -1.0
    best_bbox = None

    for det in result.detections:
        bbox = det.bounding_box
        area = float(bbox.width) * float(bbox.height)
        if area > best_area:
            best_area = area
            best_det = det
            best_bbox = bbox

    if best_det is None or best_bbox is None:
        return frame_rgb, None

    lx = int(best_bbox.origin_x)
    ly = int(best_bbox.origin_y)
    bw = int(best_bbox.width)
    bh = int(best_bbox.height)

    lx = max(0, min(w - 1, lx))
    ly = max(0, min(h - 1, ly))
    bw = max(1, min(w - lx, bw))
    bh = max(1, min(h - ly, bh))

    bgr = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)
    cv2.rectangle(bgr, (lx, ly), (lx + bw, ly + bh), (255, 255, 255), 3)

    for kp in best_det.keypoints:
        cx = int(kp.x * w)
        cy = int(kp.y * h)
        cv2.circle(bgr, (cx, cy), 6, (0, 0, 255), -1)

    drawn = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
    return drawn, (lx, ly, bw, bh)


# =========================
# 背景推論（避免卡頓） + ✅ 沒臉時清空
# =========================
_ANALYZE_LOCK = threading.Lock()
_ANALYZE_THREAD: Optional[threading.Thread] = None
_PENDING_FACE: Optional[np.ndarray] = None
_LAST_ANALYZE: Dict[str, Optional[float]] = {"ts": 0.0, "conf": None, "emo": "unknown"}

def _clear_last_analyze():
    """✅ 沒有人臉時，清空上一筆情緒，避免殘留顯示（你影片遇到的 bug）"""
    global _LAST_ANALYZE
    _LAST_ANALYZE = {"ts": 0.0, "conf": None, "emo": "unknown"}

def _analyze_worker():
    global _PENDING_FACE, _ANALYZE_THREAD, _LAST_ANALYZE

    while True:
        with _ANALYZE_LOCK:
            face = _PENDING_FACE
            _PENDING_FACE = None

        if face is None:
            break

        try:
            r = DeepFace.analyze(
                img_path=face,
                actions=["emotion"],
                enforce_detection=False,
                detector_backend="skip",  # ⭐ 你已裁 ROI，別再偵測一次
            )
            if isinstance(r, list):
                r = r[0]
            emo_dict = r.get("emotion", None)

            if isinstance(emo_dict, dict) and len(emo_dict) > 0:
                top_key = max(emo_dict, key=lambda k: float(emo_dict[k]))
                top_conf = float(emo_dict[top_key])
                _LAST_ANALYZE = {"ts": time.time(), "emo": top_key, "conf": top_conf}
        except Exception:
            pass

    _ANALYZE_THREAD = None

def _enqueue_analyze(face_roi: np.ndarray):
    global _PENDING_FACE, _ANALYZE_THREAD

    with _ANALYZE_LOCK:
        _PENDING_FACE = face_roi

    if _ANALYZE_THREAD is None:
        _ANALYZE_THREAD = threading.Thread(target=_analyze_worker, daemon=True)
        _ANALYZE_THREAD.start()


# =========================
# state helper
# =========================
def _reset_state_for_camera(st: AppState) -> AppState:
    st.running = True
    st.finished = False
    st.last_analyze_t = 0.0

    st.ema_scores = None
    st.emo_keys = []

    st.stable_emo = "unknown"
    st.candidate_emo = None
    st.candidate_since = None

    st.final_emo = None
    st.final_conf = None
    st.final_score = 3

    _clear_last_analyze()
    return st


# =========================
# streaming callback
# =========================
def on_stream(frame_rgb: np.ndarray, st: AppState):
    # stopped / finished
    if st.finished or (not st.running):
        cam_visible = False

        if st.final_emo:
            result_html = _result_card_html(st.final_emo, st.final_conf)
            mood_val = _score_radio_value(EMO_TO_SCORE.get(st.final_emo, 3))
            hint_html = _camera_hint_done()
        else:
            result_html = _hint_html("尚未完成辨識。")
            mood_val = _score_radio_value(3)
            hint_html = _camera_hint_stopped()

        return (
            gr.update(value=None, visible=cam_visible),
            gr.update(value=result_html),
            gr.update(value=mood_val),
            gr.update(value=hint_html),
            st,
            gr.update(value="重新攝影機辨識"),
        )

    # no frame yet
    if frame_rgb is None:
        return (
            gr.update(value=None, visible=True),
            gr.update(value=_hint_html("等待影像…（尚未開始串流）")),
            gr.update(value=_score_radio_value(3)),
            gr.update(value=_camera_hint_before_start()),
            st,
            gr.update(value="重新攝影機辨識"),
        )

    frame_rgb = _to_rgb_uint8(frame_rgb)

    # ✅ 偵測用原圖
    frame_drawn, bbox_px = mp_tasks_detect_and_draw(frame_rgb)

    # ✅ 顯示用 zoom
    frame_show = _zoom_center(frame_drawn if bbox_px is not None else frame_rgb, ZOOM_FACTOR)
    frame_show = _resize_for_ui(frame_show, UI_TARGET_W)

    # ✅ 沒有臉：清空殘留情緒 + 重置候選狀態
    if bbox_px is None:
        _clear_last_analyze()

        st.stable_emo = "unknown"
        st.candidate_emo = None
        st.candidate_since = None
        st.ema_scores = None
        st.emo_keys = []

        return (
            gr.update(value=frame_show, visible=True),
            gr.update(value=_hint_html("尚未偵測到人臉")),
            gr.update(value=_score_radio_value(3)),
            gr.update(value=_camera_hint_need_face()),
            st,
            gr.update(value="重新攝影機辨識"),
        )

    now = time.time()

    # 到分析時間：只排隊給背景 thread
    if now - st.last_analyze_t >= ANALYZE_EVERY_SEC:
        st.last_analyze_t = now

        x, y, bw, bh = bbox_px
        pad_x = int(bw * 0.10)
        pad_y = int(bh * 0.10)
        H, W = frame_rgb.shape[:2]
        x1 = max(0, x - pad_x)
        y1 = max(0, y - pad_y)
        x2 = min(W, x + bw + pad_x)
        y2 = min(H, y + bh + pad_y)
        face_roi = frame_rgb[y1:y2, x1:x2].astype(np.uint8)

        _enqueue_analyze(face_roi)

    # ✅ Gate：只有「有臉」才讀取 _LAST_ANALYZE（避免殘留）
    top_key = _LAST_ANALYZE.get("emo", "unknown") or "unknown"
    top_conf = _LAST_ANALYZE.get("conf", None)

    # 尚未推論到結果：先顯示 unknown
    if top_conf is None:
        show = st.stable_emo or "unknown"
        return (
            gr.update(value=frame_show, visible=True),
            gr.update(value=_result_card_html(show, None)),
            gr.update(value=_score_radio_value(EMO_TO_SCORE.get(show, 3))),
            gr.update(value=_camera_hint_before_start()),
            st,
            gr.update(value="重新攝影機辨識"),
        )

    # 低信心就不切換
    if top_conf < CONF_TH:
        show = st.stable_emo or "unknown"
        return (
            gr.update(value=frame_show, visible=True),
            gr.update(value=_result_card_html(show, None)),
            gr.update(value=_score_radio_value(EMO_TO_SCORE.get(show, 3))),
            gr.update(value=_camera_hint_before_start()),
            st,
            gr.update(value="重新攝影機辨識"),
        )

    # === EMA smoothing（用 7 類別向量）===
    all_keys = ["angry", "disgust", "fear", "happy", "sad", "surprise", "neutral"]
    vec = np.zeros((len(all_keys),), dtype=np.float32)
    if top_key in all_keys:
        vec[all_keys.index(top_key)] = float(top_conf)

    if st.ema_scores is None or st.emo_keys != all_keys:
        st.ema_scores = vec
        st.emo_keys = all_keys
    else:
        st.ema_scores = EMA_ALPHA * vec + (1.0 - EMA_ALPHA) * st.ema_scores

    stable_idx = int(np.argmax(st.ema_scores))
    stable_key = st.emo_keys[stable_idx] if st.emo_keys else "unknown"
    st.stable_emo = stable_key

    # 延遲確認，避免跳動
    if st.candidate_emo != stable_key:
        st.candidate_emo = stable_key
        st.candidate_since = now
    else:
        if st.candidate_since is None:
            st.candidate_since = now

    confirmed = (
        st.candidate_since is not None
        and (now - st.candidate_since) >= SWITCH_CONFIRM_SEC
    )

    # ✅ 完成一次辨識後停止
    if confirmed and st.final_emo is None:
        st.final_emo = stable_key
        st.final_conf = float(top_conf) if top_conf is not None else None
        st.final_score = EMO_TO_SCORE.get(stable_key, 3)

        st.running = False
        st.finished = True

        return (
            gr.update(value=None, visible=False),
            gr.update(value=_result_card_html(stable_key, st.final_conf)),
            gr.update(value=_score_radio_value(EMO_TO_SCORE.get(stable_key, 3))),
            gr.update(value=_camera_hint_done()),
            st,
            gr.update(value="重新攝影機辨識"),
        )

    return (
        gr.update(value=frame_show, visible=True),
        gr.update(value=_result_card_html(stable_key, float(top_conf))),
        gr.update(value=_score_radio_value(EMO_TO_SCORE.get(stable_key, 3))),
        gr.update(value=_camera_hint_before_start()),
        st,
        gr.update(value="重新攝影機辨識"),
    )


# =========================
# Buttons
# =========================
def on_restart(st: AppState):
    st = _reset_state_for_camera(st)
    return (
        gr.update(value=None, visible=True),
        gr.update(value=_hint_html("等待影像…（尚未開始串流）")),
        gr.update(value=_score_radio_value(3)),
        gr.update(value=_camera_hint_before_start()),
        st,
        gr.update(value="重新攝影機辨識"),
    )

def on_stop(st: AppState):
    st.running = False
    st.finished = True

    if st.final_emo is None:
        return (
            gr.update(value=None, visible=False),
            gr.update(value=_hint_html("已停止辨識。")),
            gr.update(value=_score_radio_value(3)),
            gr.update(value=_camera_hint_stopped()),
            st,
            gr.update(value="重新攝影機辨識"),
        )

    return (
        gr.update(value=None, visible=False),
        gr.update(value=_result_card_html(st.final_emo, st.final_conf)),
        gr.update(value=_score_radio_value(EMO_TO_SCORE.get(st.final_emo, 3))),
        gr.update(value=_camera_hint_done()),
        st,
        gr.update(value="重新攝影機辨識"),
    )


# =========================
# Gradio UI
# =========================
css = ""

if __name__ == "__main__":
    with gr.Blocks(title="Emotion Detector", css=css) as demo:
        st = gr.State(AppState())

        gr.Markdown("## 情緒辨識")
        hint = gr.HTML(_camera_hint_before_start())

        cam = gr.Image(
            sources=["webcam"],
            streaming=True,
            type="numpy",
            label="攝影機畫面",
            visible=True,
        )

        with gr.Row():
            btn_restart = gr.Button("重新攝影機辨識", variant="primary")
            btn_stop = gr.Button("停止", variant="secondary")

        gr.Markdown("### 辨識結果")
        result = gr.HTML(_hint_html("等待影像…（尚未開始串流）"))

        gr.Markdown("### 心情分數（系統判定）")
        mood = gr.Radio(
            choices=SCORE_CHOICES,
            value=_score_radio_value(3),
            label="心情分數",
            interactive=False,
        )

        cam.stream(
            fn=on_stream,
            inputs=[cam, st],
            outputs=[cam, result, mood, hint, st, btn_restart],
            show_progress="minimal",
        )

        btn_restart.click(
            fn=on_restart,
            inputs=[st],
            outputs=[cam, result, mood, hint, st, btn_restart],
            show_progress="minimal",
        )

        btn_stop.click(
            fn=on_stop,
            inputs=[st],
            outputs=[cam, result, mood, hint, st, btn_restart],
            show_progress="minimal",
        )

    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)