emotion.py

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

import gradio as gr
import numpy as np
from PIL import Image, ImageDraw

from deepface import DeepFace


# ===== 避免 css 未定義造成 NameError =====
css = ""


# =========================
# 參數（優先順：快）
# =========================
ANALYZE_EVERY_SEC = 1.2     # ✅ 降頻率（原 0.8）
DOWNSAMPLE_W = 480          # ✅ 降解析（原 720）
EMA_ALPHA = 0.45
CONF_TH = 40.0
SWITCH_CONFIRM_SEC = 2.5
IOU_TH = 0.25
ZOOM_FACTOR = 1.25          # ✅ 稍微降 zoom（減少失真與負擔）

DETECTOR_BACKEND = "opencv" # ✅ 改快（原 mtcnn）
ALIGN_FACE = True

# 防誤判（保留）
MIN_FACE_AREA_RATIO = 0.05
MIN_DET_CONF = 0.90


# =========================
# Warm-up（優先順：按 Record 不要等）
# =========================
_WARMED = False

def _warmup_models():
    global _WARMED
    if _WARMED:
        return
    _WARMED = True
    try:
        dummy = np.zeros((224, 224, 3), dtype=np.uint8)
        # detector warmup
        try:
            DeepFace.extract_faces(
                img_path=dummy,
                detector_backend=DETECTOR_BACKEND,
                enforce_detection=False,
                align=ALIGN_FACE,
            )
        except Exception:
            pass
        # emotion model warmup
        try:
            DeepFace.analyze(
                img_path=dummy,
                actions=["emotion"],
                enforce_detection=False,
            )
        except Exception:
            pass
    except Exception:
        pass

threading.Thread(target=_warmup_models, daemon=True).start()


# =========================
# 情緒中文 + 分數映射
# =========================
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:
    running: bool = True
    finished: bool = False
    last_analyze_t: float = 0.0

    track_bbox: Optional[Tuple[int, int, int, int]] = None

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

    stable_emo: str = "unknown"
    candidate_emo: Optional[str] = None
    candidate_since: Optional[float] = None

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

    last_face_t: float = 0.0
    face_in_box: bool = False


# =========================
# 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_done() -> str:
    return _hint_html("已完成辨識，可按「重新攝影機辨識」再次進行辨識。")

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

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 _score_radio_value(score: int) -> str:
    score = int(score)
    score = min(5, max(1, score))
    return SCORE_LABELS[score]


# =========================
# 影像工具
# =========================
def _downsample_rgb(frame: np.ndarray, target_w: int) -> np.ndarray:
    if frame is None:
        return None
    h, w = frame.shape[:2]
    if w <= target_w:
        return frame
    scale = target_w / w
    new_h = max(1, int(h * scale))
    img = Image.fromarray(frame)
    img = img.resize((target_w, new_h))
    return np.array(img)

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
    return frame[start_y:start_y + new_h, start_x:start_x + new_w]


# =========================
# 引導框（永遠畫）
# =========================
GUIDE_W_RATIO = 0.55
GUIDE_H_RATIO = 0.70
GUIDE_Y_OFFSET = 0.00

def _guide_box(frame: np.ndarray) -> Tuple[int, int, int, int]:
    h, w = frame.shape[:2]
    gw = int(w * GUIDE_W_RATIO)
    gh = int(h * GUIDE_H_RATIO)
    cx = w // 2
    cy = int(h * (0.5 + GUIDE_Y_OFFSET))
    x1 = max(0, cx - gw // 2)
    y1 = max(0, cy - gh // 2)
    x2 = min(w - 1, cx + gw // 2)
    y2 = min(h - 1, cy + gh // 2)
    return x1, y1, x2, y2

def _bbox_center_in_box(bbox_xywh: Tuple[float, float, float, float], box_xyxy: Tuple[int, int, int, int]) -> bool:
    x, y, w, h = bbox_xywh
    cx = x + w / 2.0
    cy = y + h / 2.0
    x1, y1, x2, y2 = box_xyxy
    return (x1 <= cx <= x2) and (y1 <= cy <= y2)

def _draw_guide_overlay(frame_rgb: np.ndarray, detected_in_box: bool) -> np.ndarray:
    img = Image.fromarray(frame_rgb)
    draw = ImageDraw.Draw(img)

    x1, y1, x2, y2 = _guide_box(frame_rgb)
    color = (0, 200, 0) if detected_in_box else (220, 0, 0)

    h, w = frame_rgb.shape[:2]
    thickness = max(3, int(min(w, h) * 0.008))
    for t in range(thickness):
        draw.rectangle([x1 - t, y1 - t, x2 + t, y2 + t], outline=color)

    msg = "已偵測到人臉，開始辨識情緒中…" if detected_in_box else "未偵測到人臉，請把臉放進框內"
    tx = x1
    ty = max(0, y1 - int(thickness * 6))
    draw.rectangle([tx, ty, tx + 360, ty + 32], fill=(255, 255, 255))
    draw.text((tx + 6, ty + 6), msg, fill=(0, 0, 0))

    return np.array(img)


# =========================
# Tracking / IoU
# =========================
def _iou(a: Tuple[int, int, int, int], b: Tuple[int, int, int, int]) -> float:
    ax, ay, aw, ah = a
    bx, by, bw, bh = b
    ax2, ay2 = ax + aw, ay + ah
    bx2, by2 = bx + bw, by + bh

    ix1, iy1 = max(ax, bx), max(ay, by)
    ix2, iy2 = min(ax2, bx2), min(ay2, by2)
    iw, ih = max(0, ix2 - ix1), max(0, iy2 - iy1)
    inter = iw * ih
    if inter <= 0:
        return 0.0
    union = aw * ah + bw * bh - inter
    return inter / max(1e-6, union)

def _pick_face_by_tracking(faces: List[Dict], prev_bbox: Optional[Tuple[int, int, int, int]]):
    if not faces:
        return None, None
    cand = []
    for f in faces:
        fa = f.get("facial_area", {}) or {}
        x = int(fa.get("x", 0) or 0)
        y = int(fa.get("y", 0) or 0)
        w = int(fa.get("w", 0) or 0)
        h = int(fa.get("h", 0) or 0)
        if w > 0 and h > 0 and f.get("face") is not None:
            cand.append((f, (x, y, w, h), w * h))

    if not cand:
        return None, None

    if prev_bbox is not None:
        best = None
        best_i = 0.0
        for f, bb, area in cand:
            i = _iou(prev_bbox, bb)
            if i > best_i:
                best_i = i
                best = (f, bb)
        if best is not None and best_i >= IOU_TH:
            return best[0], best[1]

    cand.sort(key=lambda x: x[2], reverse=True)
    return cand[0][0], cand[0][1]


def _reset_state_for_camera(st: AppState) -> AppState:
    st.running = True
    st.finished = False
    st.last_analyze_t = 0.0
    st.track_bbox = None
    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
    st.last_face_t = 0.0
    st.face_in_box = False
    return st


def _is_real_face(pick: Dict, bbox_full_xywh: Tuple[float, float, float, float], frame_hw: Tuple[int, int]) -> bool:
    fh, fw = frame_hw
    x, y, w, h = bbox_full_xywh
    area_ratio = (w * h) / max(1.0, float(fw * fh))
    if area_ratio < MIN_FACE_AREA_RATIO:
        return False
    det_conf = pick.get("confidence", None)
    if det_conf is not None:
        try:
            if float(det_conf) < MIN_DET_CONF:
                return False
        except Exception:
            pass
    return True


# =========================
# Stream（優先順）
# =========================
def on_stream(frame_rgb: np.ndarray, st: AppState):
    if st.finished or (not st.running):
        cam_visible = False
        if st.final_emo:
            return (
                gr.update(visible=cam_visible),
                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="重新攝影機辨識"),
            )
        return (
            gr.update(visible=cam_visible),
            gr.update(value=_hint_html("尚未完成辨識。")),
            gr.update(value=_score_radio_value(3)),
            gr.update(value=_camera_hint_stopped()),
            st,
            gr.update(value="重新攝影機辨識"),
        )

    if frame_rgb is None:
        return (
            gr.update(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 = _zoom_center(frame_rgb, ZOOM_FACTOR)
    now = time.time()

    # ✅ 先畫框（永遠立即有框），不等 detector
    detected_recent = st.face_in_box and (now - st.last_face_t) < 1.0
    annotated_fast = _draw_guide_overlay(frame_rgb, detected_recent)

    # ✅ 如果還沒到分析時間，直接回傳畫框（超順）
    if now - st.last_analyze_t < ANALYZE_EVERY_SEC:
        return (
            gr.update(visible=True, value=annotated_fast),
            gr.update(value=_hint_html("已偵測到人臉，開始辨識情緒中…" if detected_recent else "未偵測到人臉，請把臉放進框內")),
            gr.update(value=_score_radio_value(3)),
            gr.update(value=_camera_hint_before_start()),
            st,
            gr.update(value="重新攝影機辨識"),
        )

    # 到時間才做一次重工作（低頻）
    st.last_analyze_t = now
    small = _downsample_rgb(frame_rgb, DOWNSAMPLE_W)

    try:
        faces = DeepFace.extract_faces(
            img_path=small,
            detector_backend=DETECTOR_BACKEND,
            enforce_detection=True,
            align=ALIGN_FACE,
        )
    except Exception:
        faces = []

    pick, bbox = _pick_face_by_tracking(faces, st.track_bbox)

    if pick is None or bbox is None or pick.get("face") is None:
        st.track_bbox = None
        st.face_in_box = False
        st.last_face_t = now
        st.stable_emo = "unknown"
        annotated = _draw_guide_overlay(frame_rgb, False)
        return (
            gr.update(visible=True, value=annotated),
            gr.update(value=_hint_html("未偵測到人臉，請把臉放進框內")),
            gr.update(value=_score_radio_value(3)),
            gr.update(value=_camera_hint_before_start()),
            st,
            gr.update(value="重新攝影機辨識"),
        )

    sh, sw = small.shape[:2]
    fh, fw = frame_rgb.shape[:2]
    sx = fw / max(1, sw)
    sy = fh / max(1, sh)
    x, y, w, h = bbox
    bbox_full = (x * sx, y * sy, w * sx, h * sy)

    if not _is_real_face(pick, bbox_full, (fh, fw)):
        st.track_bbox = None
        st.face_in_box = False
        st.last_face_t = now
        annotated = _draw_guide_overlay(frame_rgb, False)
        return (
            gr.update(visible=True, value=annotated),
            gr.update(value=_hint_html("未偵測到人臉，請把臉放進框內")),
            gr.update(value=_score_radio_value(3)),
            gr.update(value=_camera_hint_before_start()),
            st,
            gr.update(value="重新攝影機辨識"),
        )

    guide = _guide_box(frame_rgb)
    in_box = _bbox_center_in_box(bbox_full, guide)

    st.last_face_t = now
    st.face_in_box = in_box
    st.track_bbox = bbox

    if not in_box:
        annotated = _draw_guide_overlay(frame_rgb, False)
        return (
            gr.update(visible=True, value=annotated),
            gr.update(value=_hint_html("已偵測到人臉，但請把臉放進框內")),
            gr.update(value=_score_radio_value(3)),
            gr.update(value=_camera_hint_before_start()),
            st,
            gr.update(value="重新攝影機辨識"),
        )

    # 只有臉在框內才分析情緒
    face = pick["face"]
    if face.dtype != np.uint8:
        face = np.clip(face * 255.0, 0, 255).astype(np.uint8)

    try:
        result = DeepFace.analyze(
            img_path=face,
            actions=["emotion"],
            enforce_detection=True,
        )
    except Exception:
        annotated = _draw_guide_overlay(frame_rgb, False)
        st.face_in_box = False
        st.last_face_t = now
        return (
            gr.update(visible=True, value=annotated),
            gr.update(value=_hint_html("未偵測到人臉，請把臉放進框內")),
            gr.update(value=_score_radio_value(3)),
            gr.update(value=_camera_hint_before_start()),
            st,
            gr.update(value="重新攝影機辨識"),
        )

    if isinstance(result, list):
        result = result[0]
    emo_dict = result.get("emotion", None)
    if not isinstance(emo_dict, dict) or len(emo_dict) == 0:
        annotated = _draw_guide_overlay(frame_rgb, True)
        return (
            gr.update(visible=True, value=annotated),
            gr.update(value=_hint_html("已偵測到人臉，開始辨識情緒中…")),
            gr.update(value=_score_radio_value(3)),
            gr.update(value=_camera_hint_before_start()),
            st,
            gr.update(value="重新攝影機辨識"),
        )

    top_key = max(emo_dict, key=lambda k: float(emo_dict[k]))
    top_conf = float(emo_dict[top_key])

    if top_conf < CONF_TH:
        annotated = _draw_guide_overlay(frame_rgb, True)
        return (
            gr.update(visible=True, value=annotated),
            gr.update(value=_hint_html("已偵測到人臉，開始辨識情緒中…")),
            gr.update(value=_score_radio_value(3)),
            gr.update(value=_camera_hint_before_start()),
            st,
            gr.update(value="重新攝影機辨識"),
        )

    keys = list(emo_dict.keys())
    vec = np.array([float(emo_dict[k]) for k in keys], dtype=np.float32)

    if st.ema_scores is None or st.emo_keys != keys:
        st.ema_scores = vec
        st.emo_keys = 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]
    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 = top_conf
        st.final_score = EMO_TO_SCORE.get(stable_key, 3)
        st.running = False
        st.finished = True
        return (
            gr.update(visible=False),
            gr.update(value=_result_card_html(stable_key, top_conf)),
            gr.update(value=_score_radio_value(EMO_TO_SCORE.get(stable_key, 3))),
            gr.update(value=_camera_hint_done()),
            st,
            gr.update(value="重新攝影機辨識"),
        )

    annotated = _draw_guide_overlay(frame_rgb, True)
    return (
        gr.update(visible=True, value=annotated),
        gr.update(value=_result_card_html(stable_key, 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(visible=True, value=None),  # ✅ 清掉上一張畫面
        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(visible=False, value=None),
            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(visible=False, value=None),
        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
# =========================
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)