Files
zjsh_video_collection/test_01.py

423 lines
20 KiB
Python
Raw Normal View History

2025-09-26 20:41:44 +08:00
#!/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()