#!/usr/bin/env python # -*- coding: utf-8 -*- ''' # @Time : 2025/9/11 16:48 # @Author : reenrr # @File : test_01.py # @Description: 结构化重构 - 录制6个10秒子视频（累计60秒） ''' 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 class VideoRecorder: """视频录制器类：封装配置、状态和核心逻辑""" def __init__(self): # -------------------------- 1. 配置参数（集中管理，便于修改） -------------------------- self.config = { # 检测参数 "detection_interval": 10, # 检测间隔（秒） "detection_frame_count": 3, # 每次检测抽帧数量 "required_all_noready": True, # 需所有帧为“盖板不对齐” # 摄像头参数 "rtsp_url": "rtsp://admin:XJ123456@192.168.1.50:554/streaming/channels/101", "output_dir": "camera01_videos", # 视频保存目录 "max_retry_sec": 10, # 摄像头重连超时（秒） "retry_interval_sec": 1, # 重连间隔（秒） # 模型参数 "cls_model_path": "/userdata/data_collection/cls_inference/yolov11_cls.rknn", "model_target_size": (640, 640), # 录制参数 "video_fps": 25, # 初始帧率 "video_codec": cv2.VideoWriter_fourcc(*'mp4v'), "single_duration": 10, # 单段视频时长（秒） "total_duration": 60, # 总目标时长（秒） "min_valid_duration": 10 # 最小有效视频时长（秒） } # 计算衍生配置（避免重复计算） self.config["single_frames"] = self.config["video_fps"] * self.config["single_duration"] self.config["total_frames"] = self.config["video_fps"] * self.config["total_duration"] # -------------------------- 2. 全局状态（断连后需保留） -------------------------- self.state = { "total_recorded_frames": 0, # 累计录制帧数 "current_segment": 0, # 当前分段编号 "is_recording": False, # 是否正在录制 "current_video_path": None, # 当前分段视频路径 "cached_frames": [], # 断连时缓存的帧 "last_resolution": (0, 0), # 上次摄像头分辨率（宽,高） "recording_start_time": 0 # 当前分段录制开始时间 } # -------------------------- 3. 临时状态（单次连接内有效，断连后重置） -------------------------- self.temp = { "cap": None, # 摄像头对象 "video_writer": None, # 视频写入对象 "recorded_frames": 0, # 当前分段已录帧数 "frame_count": 0, # 摄像头总读帧数 "confirmed_frames": [], # 检测通过的起始帧 "last_detection_time": time.time(),# 上次检测时间 "detection_window": [] # 检测用帧缓存窗口 } def init_environment(self): """初始化环境：创建目录、打印配置信息""" # 创建视频保存目录 os.makedirs(self.config["output_dir"], exist_ok=True) # 打印目标信息 print("=" * 60) print(f"✅ 环境初始化完成") print(f"📁 视频保存目录：{os.path.abspath(self.config['output_dir'])}") print(f"🎯 录制目标：{self.config['total_duration']}秒（{self.config['total_duration']//self.config['single_duration']}段×{self.config['single_duration']}秒）") print(f"🔍 检测条件：每{self.config['detection_interval']}秒检测，需{self.config['detection_frame_count']}帧全为'盖板不对齐'") print("=" * 60) def rotate_frame(self, pil_image): """工具函数：将PIL图像旋转180度并转为OpenCV的BGR格式""" rotated_pil = pil_image.rotate(180, expand=True) rotated_rgb = np.array(rotated_pil) return cv2.cvtColor(rotated_rgb, cv2.COLOR_RGB2BGR) def connect_camera(self): """摄像头连接：含重连逻辑，返回是否连接成功""" # 初始化摄像头 self.temp["cap"] = cv2.VideoCapture(self.config["rtsp_url"]) self.temp["cap"].set(cv2.CAP_PROP_BUFFERSIZE, 5) # 设置RTSP缓存 self.temp["cap"].set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'H264')) # 强制H264解码 # 重连逻辑 start_retry_time = time.time() while not self.temp["cap"].isOpened(): # 超时退出 if time.time() - start_retry_time >= self.config["max_retry_sec"]: print(f"\n❌ 摄像头重连超时（{self.config['max_retry_sec']}秒），程序退出") self.release_resources() # 释放资源 print(f"📊 退出时累计录制：{self.state['total_recorded_frames']/self.config['video_fps']:.1f}秒") return False # 重试提示 retry_sec = int(time.time() - start_retry_time) print(f"🔄 正在重连摄像头（已重试{retry_sec}秒）...") time.sleep(self.config["retry_interval_sec"]) # 释放旧连接，重新初始化 self.temp["cap"].release() self.temp["cap"] = cv2.VideoCapture(self.config["rtsp_url"]) # 连接成功：获取实际摄像头参数 frame_width = int(self.temp["cap"].get(cv2.CAP_PROP_FRAME_WIDTH)) frame_height = int(self.temp["cap"].get(cv2.CAP_PROP_FRAME_HEIGHT)) actual_fps = self.temp["cap"].get(cv2.CAP_PROP_FPS) # 修正帧率（以摄像头实际帧率为准） if actual_fps > 0: self.config["video_fps"] = int(actual_fps) self.config["single_frames"] = self.config["video_fps"] * self.config["single_duration"] self.config["total_frames"] = self.config["video_fps"] * self.config["total_duration"] # 检查分辨率变化 resolution_changed = False if self.state["last_resolution"] != (0, 0): if (frame_width, frame_height) != self.state["last_resolution"]: print(f"⚠️ 摄像头分辨率变化：{self.state['last_resolution']} → ({frame_width},{frame_height})") resolution_changed = True # 分辨率变化：重置录制状态 self.state["is_recording"] = False self.state["cached_frames"] = [] self.temp["detection_window"] = [] # 更新分辨率记录 self.state["last_resolution"] = (frame_width, frame_height) # 打印连接成功信息 print(f"\n✅ 摄像头连接成功") print(f"📊 分辨率：{frame_width}×{frame_height} | 实际帧率：{self.config['video_fps']}fps") print(f"📈 当前累计进度：{self.state['total_recorded_frames']/self.config['video_fps']:.1f}/{self.config['total_duration']}秒") return True, resolution_changed def restore_recording(self, resolution_changed): """重连后恢复录制状态：重新初始化写入器并写入缓存帧""" if not (self.state["is_recording"] and self.state["current_video_path"] and not resolution_changed): return # 无需恢复 print(f"🔄 恢复录制：{self.state['current_video_path']}") # 重新初始化视频写入器 frame_width, frame_height = self.state["last_resolution"] self.temp["video_writer"] = cv2.VideoWriter( self.state["current_video_path"], self.config["video_codec"], self.config["video_fps"], (frame_width, frame_height) ) # 恢复失败：重置状态 if not self.temp["video_writer"].isOpened(): print(f"⚠️ 视频写入器恢复失败，放弃当前分段") self.state["is_recording"] = False self.temp["recorded_frames"] = 0 self.state["cached_frames"] = [] return # 写入缓存帧 if self.state["cached_frames"]: print(f"🔄 恢复缓存的{len(self.state['cached_frames'])}帧") for frame in self.state["cached_frames"]: self.temp["video_writer"].write(frame) self.temp["recorded_frames"] = len(self.state["cached_frames"]) self.state["cached_frames"] = [] # 清空缓存 def run_detection(self): """执行检测逻辑：从缓存窗口抽帧，调用模型判断是否开始录制""" # 检测触发条件：时间间隔达标 + 缓存帧足够 + 未录满总目标 if (time.time() - self.temp["last_detection_time"] < self.config["detection_interval"]) or \ (len(self.temp["detection_window"]) < self.config["detection_frame_count"]) or \ (self.state["total_recorded_frames"] >= self.config["total_frames"]): return print(f"\n==== 开始检测（总读帧：{self.temp['frame_count']} | 累计进度：{self.state['total_recorded_frames']/self.config['video_fps']:.1f}秒） ====") # 均匀抽帧（避免连续帧重复） sample_step = max(1, len(self.temp["detection_window"]) // self.config["detection_frame_count"]) sample_frames = self.temp["detection_window"][::sample_step][:self.config["detection_frame_count"]] print(f"📋 检测窗口：{len(self.temp['detection_window'])}帧 → 抽帧：{len(sample_frames)}帧") # 模型分类：统计“盖板不对齐”帧数 noready_count = 0 valid_detection = True for idx, frame in enumerate(sample_frames): try: class_name = yolov11_cls_inference( self.config["cls_model_path"], frame, self.config["model_target_size"] ) except Exception as e: print(f"❌ 模型调用异常：{str(e)}") valid_detection = False break # 校验分类结果有效性 if not isinstance(class_name, str) or class_name not in ["cover_ready", "cover_noready"]: print(f"❌ 抽帧{idx+1}：分类结果无效（{class_name}）") valid_detection = False break # 统计结果 if class_name == "cover_noready": noready_count += 1 print(f"✅ 抽帧{idx+1}：{class_name}（符合条件）") else: print(f"❌ 抽帧{idx+1}：{class_name}（不符合）") # 检测未通过：重置状态 if not valid_detection or noready_count != self.config["detection_frame_count"]: if valid_detection: print(f"❌ 检测未通过（仅{noready_count}/{self.config['detection_frame_count']}帧符合条件）") else: print(f"❌ 检测未通过（存在无效结果）") self.temp["confirmed_frames"] = [] self.temp["last_detection_time"] = time.time() self.temp["detection_window"] = [] return # 检测通过：准备开始录制 print(f"✅ 检测通过！准备录制新分段") self.temp["confirmed_frames"] = sample_frames # 保存起始帧 self.start_recording() # 启动录制 # 重置检测状态 self.temp["last_detection_time"] = time.time() self.temp["detection_window"] = [] def start_recording(self): """启动新分段录制：初始化写入器、写入起始帧""" # 检查磁盘空间 total_disk, used_disk, free_disk = shutil.disk_usage(self.config["output_dir"]) if free_disk < 1024 * 1024 * 1024 * 5: # 剩余<5GB print(f"❌ 磁盘空间不足（仅剩{free_disk/(1024**3):.2f}GB），退出程序") self.release_resources() raise SystemExit(1) # 生成分段视频路径（含时间戳和分段号） self.state["current_segment"] = self.state["total_recorded_frames"] // self.config["single_frames"] + 1 timestamp = time.strftime("%Y%m%d_%H%M%S") self.state["current_video_path"] = os.path.join( self.config["output_dir"], f"video_{timestamp}_part{self.state['current_segment']}.mp4" ) # 初始化视频写入器 frame_width, frame_height = self.state["last_resolution"] self.temp["video_writer"] = cv2.VideoWriter( self.state["current_video_path"], self.config["video_codec"], self.config["video_fps"], (frame_width, frame_height) ) # 写入器初始化失败：跳过本次录制 if not self.temp["video_writer"].isOpened(): print(f"⚠️ 视频写入器初始化失败（路径：{self.state['current_video_path']}）") self.temp["confirmed_frames"] = [] return # 写入检测通过的起始帧 for frame in self.temp["confirmed_frames"]: self.temp["video_writer"].write(frame) self.temp["recorded_frames"] = len(self.temp["confirmed_frames"]) self.state["is_recording"] = True self.state["recording_start_time"] = time.time() # 打印录制启动信息 print(f"\n📹 开始录制第{self.state['current_segment']}段视频（目标{self.config['single_duration']}秒）") print(f"📁 视频路径：{self.state['current_video_path']}") print(f"🔢 已写入起始帧：{self.temp['recorded_frames']}帧") self.temp["confirmed_frames"] = [] # 清空起始帧缓存 def process_recording(self, rotated_frame): """处理录制逻辑：写入当前帧，判断分段是否完成""" if not (self.state["is_recording"] and self.temp["video_writer"]): return # 写入当前帧 self.temp["video_writer"].write(rotated_frame) self.temp["recorded_frames"] += 1 # 检查分段是否完成（帧数达标 或 时间达标，取其一） actual_duration = time.time() - self.state["recording_start_time"] if self.temp["recorded_frames"] >= self.config["single_frames"] or actual_duration >= self.config["single_duration"]: self.finish_segment() def finish_segment(self): """完成当前分段录制：释放写入器、更新累计进度、检查总目标""" # 释放当前分段写入器 self.temp["video_writer"].release() self.temp["video_writer"] = None self.state["is_recording"] = False # 计算实际录制信息 actual_duration = self.temp["recorded_frames"] / self.config["video_fps"] self.state["total_recorded_frames"] += self.temp["recorded_frames"] # 打印分段完成信息 print(f"\n✅ 第{self.state['current_segment']}段录制完成") print(f"🔍 实际录制：{self.temp['recorded_frames']}帧 ≈ {actual_duration:.1f}秒（目标{self.config['single_duration']}秒）") print(f"📊 累计进度：{self.state['total_recorded_frames']/self.config['video_fps']:.1f}/{self.config['total_duration']}秒") # 检查是否达到总目标 if self.state["total_recorded_frames"] >= self.config["total_frames"]: print(f"\n🎉 已完成{self.config['total_duration']}秒录制目标！") # 重置累计状态（如需重复录制，保留此逻辑；单次录制可改为退出） self.state["total_recorded_frames"] = 0 self.state["current_segment"] = 0 # 重置录制临时状态 self.temp["recorded_frames"] = 0 self.temp["last_detection_time"] = time.time() self.temp["detection_window"] = [] self.state["recording_start_time"] = 0 def handle_disconnect(self): """处理摄像头断连：缓存帧、释放资源""" print(f"\n⚠️ 摄像头断连（读取帧失败）") # 缓存当前录制的帧 if self.state["is_recording"]: # 收集已录帧（含当前分段所有帧） if self.temp["video_writer"]: self.temp["video_writer"].release() self.temp["video_writer"] = None # 缓存已录帧（重连后恢复） self.state["cached_frames"] = self.temp["detection_window"].copy() # 用检测窗口缓存当前帧 print(f"⏸️ 缓存{len(self.state['cached_frames'])}帧，重连后继续录制") # 重置单次连接临时状态 self.temp["frame_count"] = 0 self.temp["detection_window"] = [] if self.temp["cap"]: self.temp["cap"].release() self.temp["cap"] = None def release_resources(self): """释放所有资源：摄像头、视频写入器""" if self.temp["cap"] and self.temp["cap"].isOpened(): self.temp["cap"].release() if self.temp["video_writer"]: self.temp["video_writer"].release() print(f"\n🔌 所有资源已释放") def run(self): """主运行逻辑：初始化→循环（连接→录制→断连处理）""" # 初始化环境 self.init_environment() try: while True: # 1. 连接摄像头（含重连） connect_success, resolution_changed = self.connect_camera() if not connect_success: break # 连接失败，退出程序 # 2. 重连后恢复录制状态 self.restore_recording(resolution_changed) # 3. 主循环：读取帧→处理 last_frame_time = time.time() while True: # 读取帧 ret, frame = self.temp["cap"].read() if not ret: self.handle_disconnect() break # 断连，跳出内层循环重连 # 打印实时帧率（每10帧更新一次） current_time = time.time() fps = 1 / (current_time - last_frame_time) if (current_time - last_frame_time) > 0 else 0 last_frame_time = current_time if self.temp["frame_count"] % 10 == 0 and self.temp["frame_count"] > 0: print(f"📊 帧率：{fps:.1f} | 总读帧：{self.temp['frame_count']}", end='\r') self.temp["frame_count"] += 1 # 帧预处理（旋转180度） rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) rotated_frame = self.rotate_frame(Image.fromarray(rgb_frame)) # 加入检测窗口缓存（用于检测和断连缓存） self.temp["detection_window"].append(rotated_frame) # 限制缓存大小（避免内存溢出） if len(self.temp["detection_window"]) > self.config["detection_frame_count"] * 2: self.temp["detection_window"] = self.temp["detection_window"][-self.config["detection_frame_count"] * 2:] # 未录制：执行检测 if not self.state["is_recording"]: self.run_detection() # 正在录制：执行写入 else: self.process_recording(rotated_frame) # 捕获用户中断（Ctrl+C） except KeyboardInterrupt: print(f"\n\n👋 用户主动中断程序") # 处理未完成视频 if self.state["is_recording"] and self.state["current_video_path"] and os.path.exists(self.state["current_video_path"]): duration = self.temp["recorded_frames"] / self.config["video_fps"] if duration < self.config["min_valid_duration"]: os.remove(self.state["current_video_path"]) print(f"🗑️ 删除不足{self.config['min_valid_duration']}秒的视频：{self.state['current_video_path']}") else: print(f"⚠️ 保存未完成视频：{self.state['current_video_path']}（{duration:.1f}秒）") print(f"📊 中断时累计录制：{self.state['total_recorded_frames']/self.config['video_fps']:.1f}秒") # 捕获其他异常 except Exception as e: print(f"\n⚠️ 程序异常：{str(e)}") # 最终释放资源 finally: self.release_resources() print("\n📋 程序结束") if __name__ == '__main__': # 初始化并运行录制器 recorder = VideoRecorder() recorder.run()