zjsh_video_collection/video_new.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
# @Time    : 2025/9/11 16:48
# @Author  : reenrr
# @File    : video_new.py
'''
import cv2
import time
import os
from PIL import Image
import shutil  # 用于检查磁盘空间
from cls_inference.cls_inference import yolov11_cls_inference
import numpy as np

# ================== 配置参数 ==================
url = "rtsp://admin:XJ123456@192.168.1.50:554/streaming/channels/101"
check_interval = 100  # 每隔 N 帧检查一次“盖板对齐”状态
output_video_dir = os.path.join("camera01_videos")  # 视频保存目录

# 多帧确认参数
required_consecutive_frames = 3  # 需要连续 N 帧检测为"盖板对齐"才开始录制

# 摄像头重连参数
max_retry_seconds = 10  # 最大重试时间为10秒
retry_interval_seconds = 1  # 每隔1秒尝试重新连接一次

# 分类模型参数
cls_model_path = "/userdata/data_collection/cls_inference/yolov11_cls.rknn"  # 分类模型路径
target_size = (640, 640)

# 视频录制参数
video_fps = 25  # 视频帧率
video_codec = cv2.VideoWriter_fourcc(*'mp4v')  # MP4 编码
video_duration = 60  # 每次检测到符合条件后录制的秒数
frame_per_video = video_fps * video_duration  # 每个视频的总帧数


def rotate_frame_180(pil_image):
    """
    统一处理：将PIL图像旋转180度，并转为OpenCV录制所需的BGR格式
    input: pil_image (PIL.Image对象，RGB格式)
    output: rotated_bgr (numpy数组，BGR格式，旋转180度后)
    """
    # 1. 旋转180度（expand=True避免图像被裁剪）
    rotated_pil = pil_image.rotate(180, expand=True)
    # 2. 转为numpy数组（此时是RGB格式）
    rotated_rgb = np.array(rotated_pil)
    # 3. 转为BGR格式（OpenCV VideoWriter要求BGR输入）
    rotated_bgr = cv2.cvtColor(rotated_rgb, cv2.COLOR_RGB2BGR)
    return rotated_bgr


if __name__ == '__main__':
    # 视频录制状态变量
    is_recording = False  # 是否正在录制视频
    video_writer = None  # 视频写入对象
    recorded_frames = 0  # 当前视频已录制帧数

    # 创建视频目录
    os.makedirs(output_video_dir, exist_ok=True)
    print(f"✅ 已创建/确认视频目录: {output_video_dir}")

    while True:  # 外层循环：处理摄像头断连重连
        cap = cv2.VideoCapture(url)

        # 设置RTSP流缓存大小
        cap.set(cv2.CAP_PROP_BUFFERSIZE, 5)  # 5MB缓存
        # 强制指定解码方式
        cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'H264'))

        start_time = time.time()

        # 摄像头连接重试逻辑
        while not cap.isOpened():
            if time.time() - start_time >= max_retry_seconds:
                print(f"已尝试重新连接 {max_retry_seconds} 秒，仍无法获取视频流，程序退出。")
                # 若退出时正在录制，先释放视频写入对象
                if video_writer is not None:
                    video_writer.release()
                    video_writer = None
                cap.release()
                exit()

            print("无法打开摄像头，正在尝试重新连接...")
            time.sleep(retry_interval_seconds)
            cap.release()
            cap = cv2.VideoCapture(url)

        # 获取摄像头实际参数
        frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        actual_fps = cap.get(cv2.CAP_PROP_FPS)
        # 修正帧率
        if actual_fps > 0:
            video_fps = int(actual_fps)
            frame_per_video = video_fps * video_duration
        print(f"✅ 开始读取视频流（分辨率：{frame_width}x{frame_height}，帧率：{video_fps}）...")

        frame_count = 0  # 总帧计数器
        consecutive_ready_count = 0  # 连续"盖板对齐"计数
        confirmed_frames = []  # 存储确认的"盖板对齐"帧（用于录制开始前的帧）

        try:
            while True:
                ret, frame = cap.read()
                if not ret:
                    print("读取帧失败，可能是流中断或摄像头断开")
                    # 若断连时正在录制，先释放视频写入对象
                    if video_writer is not None:
                        video_writer.release()
                        video_writer = None
                        is_recording = False
                        recorded_frames = 0
                        print("⚠️  流中断时正在录制，已保存当前视频。")
                    # 重置检测状态
                    consecutive_ready_count = 0
                    confirmed_frames = []
                    cap.release()
                    break  # 跳出内层循环，重新连接摄像头

                frame_count += 1

                # 转换为RGB并创建PIL图像（用于后续处理）
                rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                pil_image = Image.fromarray(rgb_frame)
                # 预先旋转180度（为后续可能的录制做准备）
                rotated_bgr = rotate_frame_180(pil_image)

                # 按间隔检测“盖板对齐”状态（仅未录制时执行）
                if not is_recording and frame_count % check_interval == 0:
                    try:
                        # 调用模型判断状态（使用旋转后的图像）
                        class_name = yolov11_cls_inference(cls_model_path, rotated_bgr, target_size)

                        # 校验类别有效性
                        if not isinstance(class_name, str) or class_name not in ["cover_ready", "cover_noready"]:
                            print(f"跳过检测：模型返回无效类别（{class_name}） (总帧：{frame_count})")
                            consecutive_ready_count = 0
                            confirmed_frames = []
                            continue

                        # 检测到"盖板对齐"
                        if class_name == "cover_ready":
                            consecutive_ready_count += 1
                            # 保存当前确认的帧
                            confirmed_frames.append(rotated_bgr)
                            print(
                                f"检测到'盖板对齐'，连续计数: {consecutive_ready_count}/{required_consecutive_frames} (总帧：{frame_count})")

                            # 达到所需连续帧数，开始录制
                            if consecutive_ready_count >= required_consecutive_frames:
                                # 检查磁盘空间
                                total, used, free = shutil.disk_usage(output_video_dir)
                                if free < 1024 * 1024 * 1024 * 5:
                                    print(f"❌ 磁盘空间严重不足（仅剩 {free / (1024 ** 3):.2f} GB），停止录制。")
                                    consecutive_ready_count = 0
                                    confirmed_frames = []
                                    raise SystemExit(1)

                                # 生成视频文件名
                                timestamp = time.strftime("%Y%m%d_%H%M%S")
                                video_filename = f"video_{timestamp}.mp4"
                                video_filepath = os.path.join(output_video_dir, video_filename)

                                # 初始化视频写入器
                                video_writer = cv2.VideoWriter(
                                    video_filepath, video_codec, video_fps, (frame_width, frame_height)
                                )
                                if not video_writer.isOpened():
                                    print(f"⚠️  视频写入器初始化失败，无法录制视频（路径：{video_filepath}）")
                                    consecutive_ready_count = 0
                                    confirmed_frames = []
                                    continue

                                # 写入之前确认的帧
                                for confirmed_frame in confirmed_frames:
                                    video_writer.write(confirmed_frame)

                                is_recording = True
                                # 已录制帧数为确认帧数量
                                recorded_frames = len(confirmed_frames)
                                print(
                                    f"📹 开始录制视频（已连续{required_consecutive_frames}帧'盖板对齐'，保存路径：{video_filepath}）")
                                print(f"已写入 {recorded_frames} 帧确认帧")

                                # 重置计数器，为下一次检测做准备
                                consecutive_ready_count = 0
                                confirmed_frames = []

                        # 检测到"盖板不对齐"，重置计数器
                        else:
                            print(f"盖板状态：{class_name} (总帧：{frame_count})")
                            if consecutive_ready_count > 0:
                                print(
                                    f"检测到'盖板不对齐'，重置连续计数 (当前计数: {consecutive_ready_count}/{required_consecutive_frames})")
                            consecutive_ready_count = 0
                            confirmed_frames = []

                    except Exception as e:
                        print(f"分类模型调用异常: {e}（总帧：{frame_count}）")
                        consecutive_ready_count = 0
                        confirmed_frames = []
                        continue

                # 若正在录制，持续写入旋转后的帧
                if is_recording:
                    video_writer.write(rotated_bgr)
                    recorded_frames += 1

                    # 检查是否达到录制时长
                    if recorded_frames >= frame_per_video:
                        video_writer.release()
                        video_writer = None
                        is_recording = False
                        recorded_frames = 0
                        print(f"✅ 视频录制完成（已达 {video_duration} 秒，共录制 {frame_per_video} 帧）")

        except KeyboardInterrupt:
            print("\n用户中断程序")
            # 中断时若正在录制，先保存视频
            if video_writer is not None:
                video_writer.release()
                print("⚠️  用户中断时正在录制，已保存当前视频。")
            break

        finally:
            # 释放资源
            cap.release()
            cv2.destroyAllWindows()
            if video_writer is not None:
                video_writer.release()
            is_recording = False
            recorded_frames = 0
            consecutive_ready_count = 0
            confirmed_frames = []
            print(f"视频流已关闭，共处理总帧：{frame_count}")

    print("程序结束")