nchdlhbctm/TraceDetect-AI

1import cv2
2import numpy as np
3import os
4from PIL import Image
5from image_module import analyze_image
6
7
8def analyze_video(video_path, num_samples=10):
9    """
10    使用 OpenCV 对视频进行均匀抽帧，并复用图像引擎进行鉴别
11    :param video_path: 视频文件的路径
12    :param num_samples: 准备抽取的代表性帧数（默认 10 帧）
13    """
14    # 1. 打开视频文件
15    cap = cv2.VideoCapture(video_path)
16    if not cap.isOpened():
17        return {"error": "无法打开视频文件"}
18
19    # 2. 获取视频的基础信息
20    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
21    fps = cap.get(cv2.CAP_PROP_FPS)
22
23    # 3. 计算均匀抽帧的索引 (比如从 300 帧里均匀选 10 个时间点)
24    if total_frames < num_samples:
25        num_samples = total_frames  # 如果视频太短，有几帧抽几帧
26    intervals = np.linspace(0, total_frames - 1, num_samples, dtype=int)
27
28    frame_scores = []
29
30    # 4. 开始逐帧提取
31    for frame_idx in intervals:
32        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)  # 跳转到指定帧
33        ret, frame = cap.read()
34
35        if ret:
36            # OpenCV 默认读取的是 BGR 格式，我们需要转成正常的 RGB 格式
37            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
38            pil_img = Image.fromarray(frame_rgb)
39
40            # 临时保存为图片文件，喂给咱们之前写好的图像模块
41            temp_path = f"temp_frame_{frame_idx}.jpg"
42            pil_img.save(temp_path)
43
44            try:
45                # 🌟 核心：直接调用咱们炼好的图像鉴别引擎！
46                result = analyze_image(temp_path)
47                frame_scores.append(result['final_probability'])
48            finally:
49                # 阅后即焚，清理临时文件
50                if os.path.exists(temp_path):
51                    os.remove(temp_path)
52
53    cap.release()
54
55    if not frame_scores:
56        return {"error": "未能成功提取任何视频帧"}
57
58    # 5. 综合计算这 10 张图的得分
59    avg_score = np.mean(frame_scores)
60    max_score = np.max(frame_scores)  # 记录最可疑的一帧
61
62    return {
63        "avg_probability": avg_score,
64        "max_probability": max_score,
65        "sampled_frames": num_samples,
66        "total_frames": total_frames,
67        "fps": fps
68    }