result = model(frame, verbose=False)[0]
detections = sv.Detections.from_ultralytics(result)

detections.xyxy
detections.class_id
detections.confidence

class_name = model.names[int(class_id)]

label = f"{class_name} {confidence:.2f}"

tracker = sv.ByteTrack()

detections = tracker.update_with_detections(detections)

tracker_id = detections.tracker_id[i]
display_id = f"学生{tracker_id}"

label = f"{display_id} | {class_name} {confidence:.2f}"

def crop_xyxy(frame, xyxy):
    h, w = frame.shape[:2]
    x1, y1, x2, y2 = xyxy.astype(int).tolist()

    x1 = max(0, min(x1, w - 1))
    y1 = max(0, min(y1, h - 1))
    x2 = max(0, min(x2, w))
    y2 = max(0, min(y2, h))

    if x2 <= x1 or y2 <= y1:
        return None

    return frame[y1:y2, x1:x2].copy()

extra_result = extra_model(crop, verbose=False)[0]

top1 = int(extra_result.probs.top1)
conf = float(extra_result.probs.top1conf)
label = extra_model.names[top1]

boxes = extra_result.boxes
best_index = int(np.argmax(boxes.conf.cpu().numpy()))
class_id = int(boxes.cls[best_index].cpu().item())
conf = float(boxes.conf[best_index].cpu().item())
label = extra_model.names[class_id]

store.insert_event(
    frame_index=frame_index,
    tracker_id=tracker_id,
    display_id=display_id,
    yolo_class=yolo_class,
    yolo_conf=yolo_conf,
    xyxy=xyxy,
    extra_label=extra_label,
    extra_conf=extra_conf,
    face_id=face_id,
    face_score=face_score,
)

score = np.dot(current_embedding, saved_embedding)

embedding = embedding / np.linalg.norm(embedding)

cap = cv2.VideoCapture(source)

source = 0

source = "rtsp://username:password@192.168.1.100:554/stream1"

source = "source.mp4"

while True:
    ret, frame = cap.read()
    if not ret:
        break

    annotated, detections, labels = process_frame(frame)

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
edge_yolo_label_extra_face_demo.py

用途：
    这是一个“YOLO 标记 + supervision 追踪 + 额外识别模型 + 可选人脸/人员比对 + SQLite 保存结果”的完整示例。

你可以把它用于学习下面几个问题：

    1. YOLO 训练出来的标记怎么加到画面上？
       例如：person、cup、helmet、phone、你的自定义商品类别等。

    2. YOLO 检测框出来后，怎么再接一个额外识别模型？
       例如：
           - YOLO 先检测 person
           - 再裁剪 person 区域
           - 把裁剪图交给另一个分类模型判断：学生 / 老师 / 是否穿校服 / 是否戴安全帽

    3. 怎么把 ByteTrack 的临时追踪 ID 显示成“学生1、学生2、学生3”？
       注意：这是本次程序运行内的临时编号，不是长期身份识别。

    4. 如果确实需要跨多节课识别“是不是同一个人”，人脸比对应该怎么接？
       这个文件提供了可选 face 模式。
       但要注意：人脸 embedding 也属于生物识别信息。
       即使不保存照片，只保存向量，也不是“完全不采集人脸信息”。
       实际项目里需要明确告知、授权、最小化采集、本地加密、设置删除期限，并提供非人脸替代方式。

默认模式：
    默认使用 identity-mode=track：
        - 不做人脸识别
        - 不保存人脸图像
        - 只用 ByteTrack 给当前视频流里的目标分配临时 ID
        - 显示：学生1、学生2、学生3

可选 face 模式：
    使用 identity-mode=face：
        - 会启用 InsightFace 提取人脸 embedding
        - 会把 embedding 存入 SQLite，用于下次比对
        - 不会默认保存人脸图片
        - 这是生物识别处理，需要你自己确保合规和授权

安装：

    uv init
    uv add ultralytics supervision opencv-python numpy

如果你要启用人脸比对：

    uv add insightface onnxruntime

如果是 NVIDIA GPU 并且你知道环境支持，也可以自己改为 onnxruntime-gpu。

运行示例：

    1. 摄像头 + YOLO + 临时学生编号

        uv run python edge_yolo_label_extra_face_demo.py \
            --source 0 \
            --model yolov8n.pt \
            --identity-mode track

    2. 摄像头 + 自己训练的 YOLO

        uv run python edge_yolo_label_extra_face_demo.py \
            --source 0 \
            --model best.pt \
            --identity-mode track

    3. RTSP 摄像头

        uv run python edge_yolo_label_extra_face_demo.py \
            --source rtsp://username:password@192.168.1.100:554/stream1 \
            --model best.pt

    4. 本地视频

        uv run python edge_yolo_label_extra_face_demo.py \
            --source source.mp4 \
            --model best.pt \
            --output result.mp4

    5. YOLO 检测 + 额外分类模型

        uv run python edge_yolo_label_extra_face_demo.py \
            --source 0 \
            --model person_detector.pt \
            --extra-model uniform_classifier.pt

    6. 人脸比对模式，不保存人脸图片，只保存 embedding

        uv run python edge_yolo_label_extra_face_demo.py \
            --source 0 \
            --model yolov8n.pt \
            --identity-mode face \
            --face-db face_identities.sqlite3

整体流程：

    视频帧 frame
        ↓
    YOLO 主模型检测
        ↓
    sv.Detections.from_ultralytics(result)
        ↓
    过滤置信度 / 类别
        ↓
    ByteTrack 追踪，得到 tracker_id
        ↓
    对每个检测框裁剪 crop
        ↓
    可选：额外模型识别 crop
        ↓
    可选：人脸模型提取 embedding 并比对
        ↓
    拼接显示标签
        ↓
    画框、显示、保存事件到 SQLite

重要概念：

    YOLO 的 class_id / class_name 是“类别标记”：
        person、cup、helmet、student、teacher、phone、product_A

    ByteTrack 的 tracker_id 是“视频内临时编号”：
        当前这次运行里，某个目标可能是 #1、#2、#3

    Face embedding 的 face_id 是“跨运行身份编号”：
        如果启用人脸比对，系统可能会把某个人记为 face_001

    这三者不是同一个东西：

        class_name 代表“这是什么”
        tracker_id 代表“这个目标在当前视频里是谁”
        face_id    代表“这个人是否和数据库里的某个 face_id 匹配”

推荐学习顺序：

    第一步：只跑 --identity-mode track
    第二步：接你自己的 best.pt
    第三步：接 --extra-model 做二次识别
    第四步：确认确实需要并合规后，再考虑 --identity-mode face
"""

import argparse
import json
import sqlite3
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Optional, Tuple, List, Dict, Any

import cv2
import numpy as np
import supervision as sv
from ultralytics import YOLO


# =========================
# 1. 工具函数
# =========================

def parse_source(source: str):
    """
    OpenCV 的 VideoCapture 支持：
        0                 本机摄像头
        source.mp4         本地视频
        rtsp://...         RTSP 流

    argparse 读进来的都是字符串。
    如果用户传入 "0"，这里转成 int 0。
    """
    if source.isdigit():
        return int(source)
    return source


def safe_float(value) -> float:
    """
    Ultralytics / PyTorch / NumPy 里有些数值是 tensor 或 ndarray。
    这个函数统一转成 Python float。
    """
    try:
        if hasattr(value, "detach"):
            value = value.detach()
        if hasattr(value, "cpu"):
            value = value.cpu()
        if hasattr(value, "item"):
            return float(value.item())
        return float(value)
    except Exception:
        return 0.0


def normalize_vector(vec: np.ndarray) -> np.ndarray:
    """
    把向量归一化，便于余弦相似度计算。
    """
    vec = np.asarray(vec, dtype=np.float32)
    norm = np.linalg.norm(vec)
    if norm <= 1e-12:
        return vec
    return vec / norm


def crop_xyxy(frame: np.ndarray, xyxy: np.ndarray) -> Optional[np.ndarray]:
    """
    根据 [x1, y1, x2, y2] 从原图裁剪目标区域。
    """
    h, w = frame.shape[:2]

    x1, y1, x2, y2 = xyxy.astype(int).tolist()

    x1 = max(0, min(x1, w - 1))
    y1 = max(0, min(y1, h - 1))
    x2 = max(0, min(x2, w))
    y2 = max(0, min(y2, h))

    if x2 <= x1 or y2 <= y1:
        return None

    return frame[y1:y2, x1:x2].copy()


def get_model_name(model: YOLO, class_id: int) -> str:
    """
    从 Ultralytics 模型里根据 class_id 取 class_name。
    """
    names = model.names

    if isinstance(names, dict):
        return str(names.get(class_id, f"class_{class_id}"))

    if isinstance(names, list) and 0 <= class_id < len(names):
        return str(names[class_id])

    return f"class_{class_id}"


def make_box_annotator():
    """
    supervision 版本之间命名有过变化：
        较新版本推荐 BoxAnnotator
        一些版本里也有 BoundingBoxAnnotator

    这里做兼容。
    """
    if hasattr(sv, "BoxAnnotator"):
        return sv.BoxAnnotator()
    return sv.BoundingBoxAnnotator()


# =========================
# 2. SQLite 保存
# =========================

class EventStore:
    """
    保存检测事件。

    默认保存字段：
        - 时间
        - 帧号
        - tracker_id
        - YOLO 类别和置信度
        - 检测框
        - 额外模型识别结果
        - face_id / face_score，如果启用了人脸模式

    注意：
        为了避免每一帧都写数据库，主循环里会做保存间隔控制。
    """

    def __init__(self, db_path: str):
        self.db_path = db_path
        self.conn = sqlite3.connect(db_path)
        self.init_tables()

    def init_tables(self):
        self.conn.execute(
            """
            CREATE TABLE IF NOT EXISTS events (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                ts REAL NOT NULL,
                frame_index INTEGER NOT NULL,
                tracker_id INTEGER,
                display_id TEXT,
                yolo_class TEXT,
                yolo_conf REAL,
                xyxy_json TEXT,
                extra_label TEXT,
                extra_conf REAL,
                face_id TEXT,
                face_score REAL
            )
            """
        )
        self.conn.commit()

    def insert_event(
        self,
        frame_index: int,
        tracker_id: Optional[int],
        display_id: str,
        yolo_class: str,
        yolo_conf: float,
        xyxy: np.ndarray,
        extra_label: Optional[str],
        extra_conf: Optional[float],
        face_id: Optional[str],
        face_score: Optional[float],
    ):
        self.conn.execute(
            """
            INSERT INTO events (
                ts, frame_index, tracker_id, display_id,
                yolo_class, yolo_conf, xyxy_json,
                extra_label, extra_conf, face_id, face_score
            )
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                time.time(),
                frame_index,
                tracker_id,
                display_id,
                yolo_class,
                yolo_conf,
                json.dumps([float(x) for x in xyxy.tolist()], ensure_ascii=False),
                extra_label,
                extra_conf,
                face_id,
                face_score,
            ),
        )
        self.conn.commit()

    def close(self):
        self.conn.close()


# =========================
# 3. 人脸 embedding 存储和比对
# =========================

class FaceIdentityStore:
    """
    保存 face embedding，用于跨运行比对。

    重要提醒：
        这里保存的是人脸 embedding，不是普通匿名信息。
        它仍然可以用于识别同一个人。
        实际项目里需要合规处理。

    数据表：
        face_id     例如 face_001
        embedding   float32 向量的二进制
        dim         向量维度
        created_at  创建时间
    """

    def __init__(self, db_path: str):
        self.db_path = db_path
        self.conn = sqlite3.connect(db_path)
        self.init_tables()

    def init_tables(self):
        self.conn.execute(
            """
            CREATE TABLE IF NOT EXISTS face_identities (
                face_id TEXT PRIMARY KEY,
                embedding BLOB NOT NULL,
                dim INTEGER NOT NULL,
                created_at REAL NOT NULL
            )
            """
        )
        self.conn.commit()

    def _load_all(self) -> List[Tuple[str, np.ndarray]]:
        rows = self.conn.execute(
            "SELECT face_id, embedding, dim FROM face_identities"
        ).fetchall()

        result = []
        for face_id, blob, dim in rows:
            emb = np.frombuffer(blob, dtype=np.float32, count=dim)
            emb = normalize_vector(emb)
            result.append((face_id, emb))

        return result

    def _next_face_id(self) -> str:
        row = self.conn.execute(
            "SELECT COUNT(*) FROM face_identities"
        ).fetchone()
        count = int(row[0]) if row else 0
        return f"face_{count + 1:03d}"

    def create_identity(self, embedding: np.ndarray) -> str:
        embedding = normalize_vector(embedding).astype(np.float32)
        face_id = self._next_face_id()

        self.conn.execute(
            """
            INSERT INTO face_identities (face_id, embedding, dim, created_at)
            VALUES (?, ?, ?, ?)
            """,
            (
                face_id,
                embedding.tobytes(),
                int(embedding.shape[0]),
                time.time(),
            ),
        )
        self.conn.commit()

        return face_id

    def find_or_create(
        self,
        embedding: np.ndarray,
        threshold: float,
        auto_create: bool = True,
    ) -> Tuple[Optional[str], float, bool]:
        """
        用余弦相似度比对人脸 embedding。

        返回：
            face_id: 匹配到的 ID，或者新建的 ID
            best_score: 相似度
            created: 是否新建身份

        threshold：
            阈值需要你自己在真实场景里测试。
            太低：容易把不同人认成同一个
            太高：同一个人也可能匹配不上

        常见做法：
            先用测试集统计不同阈值下的误识别率和漏识别率，再决定。
        """
        embedding = normalize_vector(embedding)

        identities = self._load_all()

        if not identities:
            if auto_create:
                return self.create_identity(embedding), 1.0, True
            return None, 0.0, False

        best_id = None
        best_score = -1.0

        for face_id, saved_embedding in identities:
            score = float(np.dot(embedding, saved_embedding))
            if score > best_score:
                best_score = score
                best_id = face_id

        if best_score >= threshold:
            return best_id, best_score, False

        if auto_create:
            return self.create_identity(embedding), best_score, True

        return None, best_score, False

    def close(self):
        self.conn.close()


# =========================
# 4. 可选：InsightFace 人脸识别
# =========================

class FaceRecognizer:
    """
    InsightFace 封装。

    依赖安装：
        uv add insightface onnxruntime

    第一次使用时，InsightFace 可能会下载模型。

    providers:
        CPU:
            ["CPUExecutionProvider"]

        GPU:
            需要自己安装 onnxruntime-gpu，并确认 CUDA 环境可用。
    """

    def __init__(self, use_gpu: bool = False):
        try:
            from insightface.app import FaceAnalysis
        except Exception as e:
            raise RuntimeError(
                "未安装 insightface。请先执行：uv add insightface onnxruntime"
            ) from e

        providers = ["CPUExecutionProvider"]

        self.app = FaceAnalysis(
            name="buffalo_l",
            providers=providers,
        )

        # ctx_id=-1 表示 CPU
        # ctx_id=0  通常表示第 0 张 GPU，但前提是环境支持
        ctx_id = 0 if use_gpu else -1

        self.app.prepare(
            ctx_id=ctx_id,
            det_size=(640, 640),
        )

    def extract_best_face_embedding(
        self,
        bgr_image: np.ndarray,
    ) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]:
        """
        输入 BGR 图片，返回最大人脸的 embedding 和 face_bbox。

        返回：
            embedding: np.ndarray 或 None
            face_bbox: [x1, y1, x2, y2]，坐标是相对于 bgr_image 的
        """
        faces = self.app.get(bgr_image)

        if not faces:
            return None, None

        def face_area(face):
            x1, y1, x2, y2 = face.bbox
            return max(0, x2 - x1) * max(0, y2 - y1)

        best_face = max(faces, key=face_area)

        if hasattr(best_face, "normed_embedding") and best_face.normed_embedding is not None:
            embedding = np.asarray(best_face.normed_embedding, dtype=np.float32)
        else:
            embedding = normalize_vector(np.asarray(best_face.embedding, dtype=np.float32))

        bbox = np.asarray(best_face.bbox, dtype=np.float32)

        return embedding, bbox


# =========================
# 5. 额外识别模型
# =========================

@dataclass
class ExtraRecognitionResult:
    label: Optional[str]
    confidence: Optional[float]


class ExtraRecognizer:
    """
    额外识别模型封装。

    支持两种常见情况：

    情况 A：额外模型是分类模型
        例如：
            uniform_classifier.pt
            phone_classifier.pt
            student_teacher_classifier.pt

        输入：主 YOLO 检测框裁剪出来的 crop
        输出：top1 类别

    情况 B：额外模型还是检测模型
        例如：
            主模型检测 person
            额外模型在 person crop 里检测 face / badge / helmet / phone

        输入：crop
        输出：最高置信度的检测类别

    说明：
        这个类只是一个通用示例。
        实际项目里你可以按自己的业务改成更精细的逻辑。
    """

    def __init__(self, model_path: str):
        self.model = YOLO(model_path)

    def recognize(self, crop: np.ndarray) -> ExtraRecognitionResult:
        if crop is None or crop.size == 0:
            return ExtraRecognitionResult(label=None, confidence=None)

        result = self.model(crop, verbose=False)[0]

        # 分类模型：result.probs 不为空
        if getattr(result, "probs", None) is not None:
            top1 = int(result.probs.top1)
            conf = safe_float(result.probs.top1conf)
            label = get_model_name(self.model, top1)
            return ExtraRecognitionResult(label=label, confidence=conf)

        # 检测模型：result.boxes 不为空
        boxes = getattr(result, "boxes", None)
        if boxes is not None and len(boxes) > 0:
            confs = boxes.conf
            best_index = int(np.argmax(confs.cpu().numpy()))
            class_id = int(boxes.cls[best_index].cpu().item())
            conf = safe_float(boxes.conf[best_index])
            label = get_model_name(self.model, class_id)
            return ExtraRecognitionResult(label=label, confidence=conf)

        return ExtraRecognitionResult(label=None, confidence=None)


# =========================
# 6. 主应用
# =========================

class App:
    def __init__(self, args):
        self.args = args

        self.source = parse_source(args.source)
        self.main_model = YOLO(args.model)

        self.tracker = sv.ByteTrack()

        if hasattr(sv, "DetectionsSmoother"):
            self.smoother = sv.DetectionsSmoother()
        else:
            self.smoother = None

        self.box_annotator = make_box_annotator()
        self.label_annotator = sv.LabelAnnotator()

        self.trace_annotator = sv.TraceAnnotator() if hasattr(sv, "TraceAnnotator") else None

        self.event_store = EventStore(args.event_db)

        self.extra_recognizer = None
        if args.extra_model:
            self.extra_recognizer = ExtraRecognizer(args.extra_model)

        self.face_recognizer = None
        self.face_store = None

        if args.identity_mode == "face":
            self.face_recognizer = FaceRecognizer(use_gpu=args.face_gpu)
            self.face_store = FaceIdentityStore(args.face_db)

        self.last_saved_at: Dict[str, float] = {}

    def close(self):
        self.event_store.close()
        if self.face_store is not None:
            self.face_store.close()

    def filter_detections(self, detections: sv.Detections) -> sv.Detections:
        """
        过滤检测结果：
            - 置信度过滤
            - 类别白名单过滤
        """
        if len(detections) == 0:
            return detections

        detections = detections[detections.confidence >= self.args.conf]

        if self.args.class_whitelist:
            classes = [int(x.strip()) for x in self.args.class_whitelist.split(",") if x.strip()]
            if len(classes) > 0 and len(detections) > 0:
                detections = detections[np.isin(detections.class_id, classes)]

        return detections

    def get_display_id(
        self,
        tracker_id: Optional[int],
        face_id: Optional[str],
    ) -> str:
        """
        决定画面上显示什么 ID。

        identity-mode=track:
            显示：学生1、学生2、学生3
            注意：这是临时编号。

        identity-mode=face:
            如果识别到 face_id，显示 face_001
            如果没识别到，退回显示 tracker_id

        identity-mode=none:
            不显示学生编号，只显示类别。
        """
        if self.args.identity_mode == "none":
            return ""

        if self.args.identity_mode == "face" and face_id:
            return face_id

        if tracker_id is not None:
            return f"{self.args.label_prefix}{tracker_id}"

        return f"{self.args.label_prefix}?"

    def should_save_event(self, key: str) -> bool:
        """
        控制保存频率。

        不建议每一帧都写数据库。
        默认每个目标每 save_interval 秒保存一次。
        """
        now = time.time()
        last = self.last_saved_at.get(key, 0)

        if now - last >= self.args.save_interval:
            self.last_saved_at[key] = now
            return True

        return False

    def recognize_face_if_enabled(
        self,
        crop: Optional[np.ndarray],
    ) -> Tuple[Optional[str], Optional[float]]:
        """
        如果启用了 face 模式，则对目标裁剪图做：
            人脸检测
            embedding 提取
            和数据库比对
            匹配不到则创建新 face_id

        默认不保存人脸图片。
        """
        if self.args.identity_mode != "face":
            return None, None

        if crop is None or crop.size == 0:
            return None, None

        if self.face_recognizer is None or self.face_store is None:
            return None, None

        embedding, face_bbox = self.face_recognizer.extract_best_face_embedding(crop)

        if embedding is None:
            return None, None

        face_id, score, created = self.face_store.find_or_create(
            embedding=embedding,
            threshold=self.args.face_threshold,
            auto_create=True,
        )

        # 默认不保存人脸图像。
        # 只在你明确打开 --save-face-debug 时，为调试保存裁剪图。
        if self.args.save_face_debug and face_bbox is not None and face_id:
            self.save_debug_face_crop(crop, face_bbox, face_id)

        return face_id, score

    def save_debug_face_crop(
        self,
        person_crop: np.ndarray,
        face_bbox: np.ndarray,
        face_id: str,
    ):
        """
        调试用：保存人脸裁剪图。

        注意：
            真实系统不建议默认保存人脸图片。
            这个开关只适合本地开发调试，并且要有明确授权。
        """
        out_dir = Path("debug_faces")
        out_dir.mkdir(parents=True, exist_ok=True)

        x1, y1, x2, y2 = face_bbox.astype(int).tolist()

        h, w = person_crop.shape[:2]
        x1 = max(0, min(x1, w - 1))
        y1 = max(0, min(y1, h - 1))
        x2 = max(0, min(x2, w))
        y2 = max(0, min(y2, h))

        if x2 <= x1 or y2 <= y1:
            return

        face_crop = person_crop[y1:y2, x1:x2].copy()
        filename = out_dir / f"{face_id}_{int(time.time() * 1000)}.jpg"
        cv2.imwrite(str(filename), face_crop)

    def process_frame(
        self,
        frame: np.ndarray,
        frame_index: int,
    ) -> Tuple[np.ndarray, sv.Detections, List[str]]:
        """
        处理一帧：
            1. 主 YOLO 检测
            2. 转 supervision.Detections
            3. 过滤
            4. ByteTrack 追踪
            5. 可选平滑
            6. 每个框做额外模型识别 / 人脸比对
            7. 拼标签
            8. 画框
        """
        result = self.main_model(frame, verbose=False)[0]

        detections = sv.Detections.from_ultralytics(result)
        detections = self.filter_detections(detections)

        detections = self.tracker.update_with_detections(detections)

        if self.smoother is not None and len(detections) > 0:
            detections = self.smoother.update_with_detections(detections)

        labels = []

        for i in range(len(detections)):
            xyxy = detections.xyxy[i]
            class_id = int(detections.class_id[i])
            yolo_conf = float(detections.confidence[i])
            yolo_class = get_model_name(self.main_model, class_id)

            tracker_id = None
            if detections.tracker_id is not None:
                raw_tracker_id = detections.tracker_id[i]
                if raw_tracker_id is not None:
                    tracker_id = int(raw_tracker_id)

            crop = crop_xyxy(frame, xyxy)

            # 额外识别模型
            extra_label = None
            extra_conf = None

            if self.extra_recognizer is not None and crop is not None:
                extra_result = self.extra_recognizer.recognize(crop)
                extra_label = extra_result.label
                extra_conf = extra_result.confidence

            # 人脸/人员比对
            # 建议只在主 YOLO 类别是 person 时启用。
            face_id = None
            face_score = None

            if self.args.identity_mode == "face" and yolo_class == self.args.face_target_class:
                face_id, face_score = self.recognize_face_if_enabled(crop)

            display_id = self.get_display_id(
                tracker_id=tracker_id,
                face_id=face_id,
            )

            # 拼接画面标签
            label_parts = []

            if display_id:
                label_parts.append(display_id)

            label_parts.append(f"{yolo_class} {yolo_conf:.2f}")

            if extra_label:
                if extra_conf is not None:
                    label_parts.append(f"{extra_label} {extra_conf:.2f}")
                else:
                    label_parts.append(extra_label)

            if face_id and face_score is not None:
                label_parts.append(f"face_score {face_score:.2f}")

            label = " | ".join(label_parts)
            labels.append(label)

            # 保存事件
            # 使用 face_id 优先，其次 tracker_id，最后使用检测框索引
            save_key = face_id or (f"track_{tracker_id}" if tracker_id is not None else f"det_{i}")

            if self.should_save_event(save_key):
                self.event_store.insert_event(
                    frame_index=frame_index,
                    tracker_id=tracker_id,
                    display_id=display_id,
                    yolo_class=yolo_class,
                    yolo_conf=yolo_conf,
                    xyxy=xyxy,
                    extra_label=extra_label,
                    extra_conf=extra_conf,
                    face_id=face_id,
                    face_score=face_score,
                )

        annotated = frame.copy()

        if self.trace_annotator is not None and len(detections) > 0:
            annotated = self.trace_annotator.annotate(
                scene=annotated,
                detections=detections,
            )

        annotated = self.box_annotator.annotate(
            scene=annotated,
            detections=detections,
        )

        annotated = self.label_annotator.annotate(
            scene=annotated,
            detections=detections,
            labels=labels,
        )

        cv2.putText(
            annotated,
            f"frame: {frame_index}",
            (20, 40),
            cv2.FONT_HERSHEY_SIMPLEX,
            1,
            (0, 255, 0),
            2,
        )

        return annotated, detections, labels

    def run(self):
        cap = cv2.VideoCapture(self.source)

        if not cap.isOpened():
            raise RuntimeError(f"无法打开视频源：{self.args.source}")

        writer = None
        frame_index = 0

        try:
            while True:
                ret, frame = cap.read()

                if not ret:
                    print("视频流结束或读取失败")
                    break

                annotated, detections, labels = self.process_frame(
                    frame=frame,
                    frame_index=frame_index,
                )

                if writer is None and self.args.output:
                    h, w = annotated.shape[:2]
                    fps = cap.get(cv2.CAP_PROP_FPS)
                    if fps is None or fps <= 1:
                        fps = self.args.output_fps

                    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
                    writer = cv2.VideoWriter(
                        self.args.output,
                        fourcc,
                        float(fps),
                        (w, h),
                    )

                if writer is not None:
                    writer.write(annotated)

                if not self.args.no_window:
                    cv2.imshow("YOLO + supervision + extra recognition", annotated)

                    if cv2.waitKey(1) & 0xFF == ord("q"):
                        break

                frame_index += 1

        finally:
            cap.release()

            if writer is not None:
                writer.release()

            if not self.args.no_window:
                cv2.destroyAllWindows()

            self.close()


# =========================
# 7. 命令行参数
# =========================

def build_arg_parser():
    parser = argparse.ArgumentParser(
        description="YOLO + supervision + 额外识别 + 可选人脸比对 示例"
    )

    parser.add_argument(
        "--source",
        default="0",
        help="视频源：0 表示本机摄像头，也可以是 source.mp4 或 rtsp://...",
    )

    parser.add_argument(
        "--model",
        default="yolov8n.pt",
        help="主 YOLO 模型路径，例如 yolov8n.pt 或 best.pt",
    )

    parser.add_argument(
        "--extra-model",
        default=None,
        help="可选：额外识别模型路径，例如 uniform_classifier.pt；不传则不启用",
    )

    parser.add_argument(
        "--conf",
        type=float,
        default=0.35,
        help="主 YOLO 检测置信度阈值",
    )

    parser.add_argument(
        "--class-whitelist",
        default=None,
        help="只保留指定 class_id，例如 '0' 或 '0,2,3'；不传则不过滤类别",
    )

    parser.add_argument(
        "--identity-mode",
        choices=["none", "track", "face"],
        default="track",
        help=(
            "none: 不显示身份编号；"
            "track: 使用 ByteTrack 临时编号；"
            "face: 使用人脸 embedding 比对，需要 insightface"
        ),
    )

    parser.add_argument(
        "--label-prefix",
        default="学生",
        help="track 模式下的显示前缀，例如 学生、目标、person_",
    )

    parser.add_argument(
        "--event-db",
        default="events.sqlite3",
        help="检测事件保存数据库",
    )

    parser.add_argument(
        "--save-interval",
        type=float,
        default=5.0,
        help="同一个目标每隔多少秒保存一次事件，避免每帧写数据库",
    )

    parser.add_argument(
        "--output",
        default=None,
        help="可选：保存标注后的视频，例如 result.mp4",
    )

    parser.add_argument(
        "--output-fps",
        type=int,
        default=25,
        help="视频源读不到 FPS 时，输出视频使用的默认 FPS",
    )

    parser.add_argument(
        "--no-window",
        action="store_true",
        help="不弹出窗口，适合服务器 / 边缘设备无桌面环境",
    )

    # face 模式相关
    parser.add_argument(
        "--face-db",
        default="face_identities.sqlite3",
        help="face 模式下保存人脸 embedding 的数据库",
    )

    parser.add_argument(
        "--face-threshold",
        type=float,
        default=0.45,
        help="face embedding 余弦相似度阈值，需要用你的场景数据测试调整",
    )

    parser.add_argument(
        "--face-target-class",
        default="person",
        help="只有主 YOLO 类别等于这个名字时才做人脸比对，默认 person",
    )

    parser.add_argument(
        "--face-gpu",
        action="store_true",
        help="尝试使用 GPU 跑 InsightFace，前提是你安装并配置了 onnxruntime-gpu",
    )

    parser.add_argument(
        "--save-face-debug",
        action="store_true",
        help="调试用：保存人脸裁剪图到 debug_faces/。真实项目不建议默认开启。",
    )

    return parser


def main():
    parser = build_arg_parser()
    args = parser.parse_args()

    print("=" * 80)
    print("YOLO + supervision + 额外识别 + 可选人脸比对 示例")
    print("=" * 80)
    print(f"视频源: {args.source}")
    print(f"主模型: {args.model}")
    print(f"额外模型: {args.extra_model}")
    print(f"identity-mode: {args.identity_mode}")
    print(f"事件数据库: {args.event_db}")

    if args.identity_mode == "face":
        print()
        print("注意：你启用了 face 模式。")
        print("这会提取并保存人脸 embedding，用于跨运行比对。")
        print("即使不保存照片，embedding 也属于可用于识别个人的生物识别信息。")
        print("请确保已获得明确授权，并设置合适的数据保护和删除机制。")
        print()

    app = App(args)
    app.run()


if __name__ == "__main__":
    main()