423 lines
20 KiB
Python
423 lines
20 KiB
Python
#!/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() |