sentinel-home-ai/fam-core/tools/backfill_frames.py

"""为历史事件补抽关键帧图（视频仍存于 NAS 时）

Edge 注入关键帧 base64 上线前落库的事件没有帧图，本脚本从原始视频
按 frame_timestamp - event_start_time 偏移重新抽帧，补齐到 UI 静态目录。

用法:
    venv/bin/python tools/backfill_frames.py [--dry-run] [--event-id N]

特性:
- 幂等: 单帧文件已存在即跳过，帧数齐全的事件整条跳过
- NAS ffmpeg41 无 image2 muxer，必须用 -f singlejpeg 输出 jpg
- 抽帧尺寸与 NAS 预压缩一致（480p 等比缩放），-q:v 5 约 60-100KB/张
"""
import os
import re
import sys
import argparse
import subprocess
from datetime import datetime

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'src'))

import pymysql
from fam_core.config_loader import load_config

FFMPEG = '/var/packages/CodecPack/target/bin/ffmpeg41'


def video_duration(path: str) -> float:
    """解析 ffmpeg header 里的 Duration（无 ffprobe 环境）"""
    try:
        r = subprocess.run([FFMPEG, '-i', path], capture_output=True, text=True, timeout=60)
        m = re.search(r'Duration:\s*(\d+):(\d+):(\d+)', r.stderr)
        if m:
            h, mi, s = (int(x) for x in m.groups())
            return h * 3600 + mi * 60 + s
    except Exception:
        pass
    return 0.0


def extract_frame(video: str, offset: float, out_path: str) -> bool:
    cmd = [
        FFMPEG, '-y',
        '-ss', f'{offset:.1f}',
        '-i', video,
        '-frames:v', '1',
        '-vf', 'scale=854:480:force_original_aspect_ratio=decrease,scale=trunc(iw/2)*2:trunc(ih/2)*2',
        '-q:v', '5',
        '-f', 'singlejpeg',
        out_path,
    ]
    try:
        r = subprocess.run(cmd, capture_output=True, timeout=120)
        return r.returncode == 0 and os.path.getsize(out_path) > 1024
    except Exception:
        return False


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--dry-run', action='store_true', help='只统计不落盘')
    ap.add_argument('--event-id', type=int, default=None, help='只处理指定事件')
    args = ap.parse_args()

    cfg = load_config()
    frame_dir = cfg.get('storage', {}).get(
        'frame_image_dir', '/volume1/web/sentinel-home-ai/fam-ui/static/frames')
    db = cfg['database']
    conn = pymysql.connect(
        host=db.get('host', '127.0.0.1'),
        port=db.get('port', 3306),
        user=db.get('user', 'root'),
        password=db.get('password', ''),
        database=db.get('database', 'sentinel_home_ai'),
        unix_socket=db.get('unix_socket'),
        charset='utf8mb4',
        cursorclass=pymysql.cursors.DictCursor,
    )

    try:
        with conn.cursor() as cur:
            sql = """
                SELECT me.event_id, me.task_id, me.event_start_time, pt.video_path,
                       (SELECT COUNT(*) FROM event_details ed
                         WHERE ed.event_id = me.event_id) AS detail_count
                FROM monitor_events me
                LEFT JOIN process_tasks pt ON pt.task_id = me.task_id
            """
            params = ()
            if args.event_id:
                sql += ' WHERE me.event_id = %s'
                params = (args.event_id,)
            sql += ' ORDER BY me.event_id'
            cur.execute(sql, params)
            events = cur.fetchall()

        stat = {'skip_complete': 0, 'skip_no_video': 0, 'extracted': 0, 'failed': 0, 'events': 0}

        for ev in events:
            eid = ev['event_id']
            edir = os.path.join(frame_dir, f'event_{eid}')
            existing = {f for f in os.listdir(edir) if f.endswith('.jpg')} if os.path.isdir(edir) else set()

            with conn.cursor() as cur:
                cur.execute(
                    """SELECT frame_index, frame_timestamp FROM event_details
                       WHERE event_id = %s ORDER BY frame_index""",
                    (eid,))
                frames = cur.fetchall()

            missing = [f for f in frames if f'frame_{f["frame_index"]}.jpg' not in existing]
            if not missing:
                stat['skip_complete'] += 1
                continue

            video = ev.get('video_path') or ''
            source = '原始视频'
            if not video or not os.path.isfile(video):
                # 原始视频被监控保留策略清理时，回退到 dispatcher 压缩缓存（480p 副本）
                cached = os.path.join(
                    f'/tmp/fam_compressed/task_{ev["task_id"]}', os.path.basename(video)) if video else ''
                if cached and os.path.isfile(cached):
                    video = cached
                    source = '压缩缓存'
            if not video or not os.path.isfile(video):
                print(f'[event {eid}] 原始视频与压缩缓存均不存在，跳过 {len(missing)} 帧: {video}')
                stat['skip_no_video'] += 1
                continue

            dur = video_duration(video)
            start = ev['event_start_time']
            if isinstance(start, str):
                start = datetime.strptime(start[:19], '%Y-%m-%d %H:%M:%S')

            os.makedirs(edir, exist_ok=True)
            stat['events'] += 1
            print(f'[event {eid}] 补 {len(missing)}/{len(frames)} 帧 · {source} (视频 {os.path.basename(video)}, {dur:.0f}s)')

            for f in missing:
                out_path = os.path.join(edir, f'frame_{f["frame_index"]}.jpg')
                ts = f['frame_timestamp']
                if isinstance(ts, str):
                    ts = datetime.strptime(ts[:19], '%Y-%m-%d %H:%M:%S')
                offset = (ts - start).total_seconds()
                offset = max(1.0, min(offset, max(1.0, dur - 2)))
                if args.dry_run:
                    print(f'  dry-run frame_{f["frame_index"]} @ {offset:.0f}s')
                    continue
                ok = extract_frame(video, offset, out_path)
                stat['extracted' if ok else 'failed'] += 1
                if not ok:
                    print(f'  失败 frame_{f["frame_index"]} @ {offset:.0f}s')
                    if os.path.exists(out_path):
                        os.remove(out_path)

        print(f"\n完成: 事件 {stat['events']} 个已补 | 抽帧成功 {stat['extracted']} 失败 {stat['failed']} "
              f"| 齐全跳过 {stat['skip_complete']} | 视频缺失跳过 {stat['skip_no_video']}")
    finally:
        conn.close()


if __name__ == '__main__':
    main()