用python调用显卡GPU截取视频中的人像图片

最近接到一个项目，需要从几万个零散的视频中提取有人走动的的内容保存为图片，方便查找线索。
这对于python的cv2加上YoLo8模型，简直不要太好用，先测试一个文件，然后批量调用就行。

下载示例文件：video.zip (10MB)，把这个压缩包里的文件解压到C盘根目录，执行video/video.bat或者手动运行下面的命令
命令：python c:/video/video.py c:/video/20210305.mp4 c:/video

链接:https://caiyun.139.com/w/i/2pU95BMDyRc86
提取码:037q

需要用到的组件，python版本3.10，其它版本可能不兼容，需要排错
最好先更新一下pip（建议）
python.exe -m pip install –upgrade pip

安装OpenCV (cv2，必须)
用于视频处理和图像操作。
安装命令：pip install opencv-python==4.8.0.76
Ultralytics YOLO（必须）
提供 YOLOv8 模型支持。
安装命令：pip install ultralytics       (大概需下载好几百兆的文件，没有国外加速会下很久很久。。。）
PyTorch
YOLOv8 的深度学习框架依赖。

如果需要显卡CUDA加速，需要安装（文件很大。。。。，我不装了）

python.exe -m pip install torch torchvision torchaudio –index-url https://download.pytorch.org/whl/cu118

这几个组件版本分别如下，版本不兼容可能会导致CUDA无法运行：

numpy 1.26.4
opencv-python 4.8.0.76
torch 2.2.0+cu118
torchvision 0.17.0+cu118
ultralytics 8.2.89 

import sys
import os
import cv2
from ultralytics import YOLO
import torch

def analyze_video_yolo(video_path, output_dir, save_interval_sec=1, min_confidence=0.5):
    if not os.path.exists(video_path):
        print(f"视频文件 {video_path} 不存在")
        return
    try:
        os.makedirs(output_dir, exist_ok=True)
        test_file = os.path.join(output_dir, "test.txt")
        with open(test_file, 'w') as f:
            f.write("test")
        os.remove(test_file)
    except Exception as e:
        print(f"输出目录 {output_dir} 不可写: {e}")
        return
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"无法打开视频：{video_path}")
        return
    fps = cap.get(cv2.CAP_PROP_FPS)
    if fps == 0:
        print("无法获取视频FPS，默认30")
        fps = 30
    save_interval_frames = int(fps * save_interval_sec)
    if torch.cuda.is_available():
        device = 'cuda'
        print("使用 GPU（CUDA）进行人体检测")
    else:
        device = 'cpu'
        print("未检测到可用的 GPU，将使用 CPU 进行人体检测")
    try:
        model = YOLO("yolov8n.pt")
    except Exception as e:
        print(f"模型加载失败: {e}")
        return

    frame_count = 0
    saved_count = 0

    print(f"开始分析视频: {video_path}，FPS: {fps:.2f}")

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

        # 每隔一定帧检测一次并保存（如果检测到人）
        if frame_count % save_interval_frames == 0:
            results = model(frame, device=device)[0]

            person_boxes = [
                box for box in results.boxes
                if int(box.cls) == 0 and box.conf.item() > min_confidence
            ]

            if person_boxes:
                output_path = os.path.join(output_dir, f"frame_{frame_count}.jpg")
                try:
                    cv2.imwrite(output_path, frame)
                    saved_count += 1
                    print(f"[帧 {frame_count}] 检测到 {len(person_boxes)} 人体，保存整帧截图")
                except Exception as e:
                    print(f"保存图像失败: {e}")

        # 绘制检测框用于显示
        results = model(frame, device=device)[0]
        for box in results.boxes:
            if int(box.cls) == 0 and box.conf.item() > min_confidence:
                x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                conf = box.conf.item()
                cv2.putText(frame, f"Person {conf:.2f}", (x1, y1 - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)

        cv2.imshow("YOLOv8 Detection", frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            print("检测中止，手动退出。")
            break

        frame_count += 1

    cap.release()
    cv2.destroyAllWindows()
    print(f"完成！总帧数: {frame_count}，保存截图: {saved_count}")

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("用法: python video_yolo.py [视频路径] [输出目录] [可选:截图间隔秒数（默认1秒）]")
        sys.exit(1)

    video_path = sys.argv[1]
    output_dir = sys.argv[2]
    interval_sec = float(sys.argv[3]) if len(sys.argv) > 3 else 1

    analyze_video_yolo(video_path, output_dir, interval_sec)