yingshiyun/analysis/person_tracker.py

"""在 motion_filter 找出的活动窗口内跑 YOLO11n 检测 + ByteTrack 跟踪，
统计人物出现次数，写入 events 表（event_type='person_presence'）。

一个 track（同一个人从出现到消失的连续跟踪片段）对应一条 person_presence
事件。只在活动窗口内跑，空场景（motion_windows_json 为空）直接跳过，不跑模型。
"""
import json
from datetime import datetime, timedelta

import cv2

from analysis.config import (
    PERSON_MODEL_NAME,
    PERSON_DEVICE,
    PERSON_CONF_THRESHOLD,
    PERSON_SAMPLE_STRIDE,
    PERSON_MIN_TRACK_DURATION_SEC,
    PROJECT_ROOT,
)
from analysis.db import get_all_videos, insert_event, update_video_status

_model = None


def _get_model():
    """YOLO模型只在真正用到时才加载（避免只跑无活动视频时也触发权重下载/加载）。"""
    global _model
    if _model is None:
        from ultralytics import YOLO

        _model = YOLO(PERSON_MODEL_NAME)
    return _model


def _track_windows(file_path, windows, fps):
    """在给定的活动窗口内跑检测+跟踪，返回 [(track_id, start_sec, end_sec, max_conf), ...]。"""
    model = _get_model()
    cap = cv2.VideoCapture(str(file_path))
    if not cap.isOpened():
        return []

    spans = {}  # (window_index, track_id) -> [min_t, max_t, max_conf]

    for window_index, (start_sec, end_sec) in enumerate(windows):
        start_frame = int(start_sec * fps)
        end_frame = int(end_sec * fps)
        cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)

        frame_idx = start_frame
        first_call = True
        while frame_idx <= end_frame:
            ret, frame = cap.read()
            if not ret:
                break

            if (frame_idx - start_frame) % PERSON_SAMPLE_STRIDE == 0:
                # persist=False 强制在每个窗口开头重置跟踪器状态，跟踪器内部的track_id计数器
                # 也会跟着从1重新开始——不同窗口间同一个数字的track_id并不代表同一个人，
                # 所以下面的 spans key 必须带上 window_index，否则会把不同窗口的片段错误合并
                results = model.track(
                    frame, persist=not first_call, classes=[0], conf=PERSON_CONF_THRESHOLD,
                    tracker="bytetrack.yaml", device=PERSON_DEVICE, verbose=False,
                )
                first_call = False

                boxes = results[0].boxes
                if boxes is not None and boxes.id is not None:
                    t = frame_idx / fps
                    for track_id, conf in zip(boxes.id.tolist(), boxes.conf.tolist()):
                        key = (window_index, int(track_id))
                        if key not in spans:
                            spans[key] = [t, t, conf]
                        else:
                            spans[key][1] = t
                            spans[key][2] = max(spans[key][2], conf)

            frame_idx += 1

    cap.release()
    return [
        (track_id, start, end, conf)
        for (_window_index, track_id), (start, end, conf) in spans.items()
        if (end - start) >= PERSON_MIN_TRACK_DURATION_SEC
    ]


def process_video(video):
    """返回 (track_count, failed)。failed=True 表示该视频处理异常已标记 person_failed。"""
    try:
        windows = json.loads(video["motion_windows_json"] or "[]")
    except (json.JSONDecodeError, TypeError):
        windows = []
    if not windows:
        update_video_status(video["id"], "person_done")
        return 0, False

    file_path = PROJECT_ROOT / video["file_path"]
    if not file_path.exists():
        update_video_status(video["id"], "person_failed", error_message="文件不存在")
        return 0, True

    try:
        fps = video["fps"] or 15.0
        tracks = _track_windows(file_path, windows, fps)
    except Exception as e:
        update_video_status(video["id"], "person_failed", error_message=str(e))
        return 0, True

    start_time = None
    if video["start_time"]:
        try:
            start_time = datetime.strptime(video["start_time"], "%Y-%m-%d %H:%M:%S")
        except ValueError:
            pass  # 只精确到天的粗粒度值（"YYYY-MM-DD"），无法算出绝对时间戳，timestamp_abs留空

    for track_id, start_sec, end_sec, conf in tracks:
        timestamp_abs = (start_time + timedelta(seconds=start_sec)).isoformat() if start_time else None
        insert_event({
            "video_id": video["id"],
            "event_type": "person_presence",
            "label": None,
            "start_offset_sec": start_sec,
            "end_offset_sec": end_sec,
            "timestamp_abs": timestamp_abs,
            "confidence": conf,
            "track_id": track_id,
            "thumbnail_path": None,
            "extra_json": None,
            "created_at": datetime.now().isoformat(),
        })

    update_video_status(video["id"], "person_done")
    return len(tracks), False


def process_all(limit=None):
    videos = get_all_videos(status="motion_done")
    if limit:
        videos = videos[:limit]

    processed = 0
    total_tracks = 0
    failed = 0
    for video in videos:
        tracks, is_failed = process_video(video)
        total_tracks += tracks
        if is_failed:
            failed += 1
        processed += 1

    return {"processed": processed, "total_tracks": total_tracks, "failed": failed}


if __name__ == "__main__":
    print("===== 开始统计人物出现次数 =====")
    result = process_all()
    print(f"→ 处理 {result['processed']} 个视频，共识别出 {result['total_tracks']} 段人物出现")
    print("===== 统计完成 =====")