cfc/cfc-langgraph/src/graphs/emotion.py

"""
情绪识别 LangGraph — DeepFace 人脸情绪分析

流程：
  START → detect_faces → analyze_emotion → END

输出：7 类基础情绪概率 + 主导情绪 + 中文标签
"""
import base64
import io
import os
from typing import TypedDict, Optional

import requests
from PIL import Image
from langgraph.graph import StateGraph, START, END

# DeepFace 延迟导入（未安装时 graceful fallback）
try:
    from deepface import DeepFace
    DEEPFACE_AVAILABLE = True
except ImportError:
    DEEPFACE_AVAILABLE = False


# ── 中文情绪映射 ──────────────────────────────────────────────
EMOTION_ZH = {
    "angry": "愤怒",
    "disgust": "厌恶",
    "fear": "恐惧",
    "happy": "开心",
    "sad": "悲伤",
    "surprise": "惊讶",
    "neutral": "平静",
}


# ── Graph State ───────────────────────────────────────────────
class GraphState(TypedDict):
    request: dict
    image_bytes: Optional[bytes]       # 原始图片字节
    image_path: Optional[str]          # 临时文件路径（供 DeepFace 使用）
    emotion_result: Optional[dict]     # DeepFace analyze() 结果
    emotions: list                     # 结构化输出
    dominant_emotion: str
    dominant_label_zh: str
    all_emotions: dict                 # 7 类完整概率
    error: Optional[str]


# ── Nodes ─────────────────────────────────────────────────────

def load_image(state: GraphState) -> GraphState:
    """将 image_url 下载到内存 / 解码 base64"""
    req = state["request"]
    img_bytes = None

    if req.get("image_url"):
        try:
            resp = requests.get(req["image_url"], timeout=10)
            resp.raise_for_status()
            img_bytes = resp.content
        except Exception as e:
            return {**state, "error": f"图片下载失败: {e}"}

    elif req.get("image_base64"):
        try:
            img_bytes = base64.b64decode(req["image_base64"])
        except Exception as e:
            return {**state, "error": f"Base64 解码失败: {e}"}

    else:
        return {**state, "error": "image_url 或 image_base64 至少提供一项"}

    if img_bytes:
        state["image_bytes"] = img_bytes
        # 写入临时文件供 DeepFace 使用
        tmp_path = f"/tmp/emotion_{os.getpid()}.jpg"
        with open(tmp_path, "wb") as f:
            f.write(img_bytes)
        state["image_path"] = tmp_path

    return state


def analyze_emotion(state: GraphState) -> GraphState:
    """调用 DeepFace 进行情绪识别"""
    if not DEEPFACE_AVAILABLE or not state.get("image_path"):
        return {**state, "error": "DeepFace 未安装或图片未加载"}

    try:
        result = DeepFace.analyze(
            img_path=state["image_path"],
            actions=["emotion"],
            detector_backend="opencv",   # 内置 Haar Cascade，无需额外模型
            enforce_detection=False,
            silent=True,
        )
        # DeepFace.analyze 返回 list（可能多张人脸）
        if isinstance(result, list) and len(result) > 0:
            state["emotion_result"] = result[0]
        else:
            return {**state, "error": "未检测到人脸"}
    except Exception as e:
        return {**state, "error": f"DeepFace 分析失败: {e}"}

    return state


def format_result(state: GraphState) -> GraphState:
    """将 DeepFace 结果格式化为结构化输出"""
    if not state.get("emotion_result"):
        return state

    emo = state["emotion_result"]
    # DeepFace 返回 0-100 百分比（如 {'happy': 85.0, ...}），需归一化为 0.0-1.0
    raw_emotions = emo.get("emotion", {})

    # 归一化：值 >1 视为百分比，除以 100
    norm = {}
    for k, v in raw_emotions.items():
        norm[k] = round(v / 100.0, 4) if v > 1.0 else round(v, 4)

    # 找出主导情绪
    dominant = max(norm, key=norm.get) if norm else "neutral"
    dominant_zh = EMOTION_ZH.get(dominant, dominant)

    # 结构化列表
    emotions = [
        {"emotion": k, "confidence": v}
        for k, v in sorted(norm.items(), key=lambda x: -x[1])
    ]

    state["dominant_emotion"] = dominant
    state["dominant_label_zh"] = dominant_zh
    state["all_emotions"] = norm
    state["emotions"] = emotions
    return state


def cleanup(state: GraphState) -> GraphState:
    """清理临时文件"""
    import os as _os
    path = state.get("image_path")
    if path and _os.path.exists(path):
        try:
            _os.remove(path)
        except Exception:
            pass
    return state


# ── Graph 构建 ────────────────────────────────────────────────

def build_emotion_graph():
    graph = StateGraph(GraphState)
    graph.add_node("load_image", load_image)
    graph.add_node("analyze_emotion", analyze_emotion)
    graph.add_node("format_result", format_result)
    graph.add_node("cleanup", cleanup)

    graph.add_edge(START, "load_image")
    graph.add_edge("load_image", "analyze_emotion")
    graph.add_edge("analyze_emotion", "format_result")
    graph.add_edge("format_result", "cleanup")
    graph.add_edge("cleanup", END)
    return graph.compile()


_emotion_graph = None


def get_emotion_graph():
    global _emotion_graph
    if _emotion_graph is None:
        _emotion_graph = build_emotion_graph()
    return _emotion_graph