306 lines
16 KiB
Python
306 lines
16 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
'''
|
||
# @Time : 2025/9/11 16:48
|
||
# @Author : reenrr
|
||
# @File : test_01.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
|
||
|
||
# ================== 配置参数（严格按“定义在前，使用在后”排序）==================
|
||
# 1. 检测核心参数（先定义，后续打印和逻辑会用到）
|
||
detection_interval = 10 # 每隔10秒检查一次
|
||
detection_frame_count = 3 # 每次检测抽取3帧
|
||
required_all_noready = True # 要求所有帧都为“盖板不对齐”
|
||
|
||
# 2. 视频存储与摄像头基础参数
|
||
url = "rtsp://admin:XJ123456@192.168.1.50:554/streaming/channels/101"
|
||
output_video_dir = os.path.join("camera01_videos") # 视频保存目录
|
||
|
||
# 3. 摄像头重连参数
|
||
max_retry_seconds = 10
|
||
retry_interval_seconds = 1
|
||
|
||
# 4. 分类模型参数
|
||
cls_model_path = "/userdata/data_collection/cls_inference/yolov11_cls.rknn"
|
||
target_size = (640, 640)
|
||
|
||
# 5. 视频录制参数
|
||
video_fps = 25
|
||
video_codec = cv2.VideoWriter_fourcc(*'mp4v')
|
||
single_recording_duration = 10 # 每次录制10秒
|
||
total_target_duration = 60 # 累计目标60秒
|
||
single_recording_frames = video_fps * single_recording_duration
|
||
total_target_frames = video_fps * total_target_duration
|
||
|
||
|
||
def rotate_frame_180(pil_image):
|
||
"""将PIL图像旋转180度并转为OpenCV的BGR格式"""
|
||
rotated_pil = pil_image.rotate(180, expand=True)
|
||
rotated_rgb = np.array(rotated_pil)
|
||
rotated_bgr = cv2.cvtColor(rotated_rgb, cv2.COLOR_RGB2BGR)
|
||
return rotated_bgr
|
||
|
||
|
||
if __name__ == '__main__':
|
||
# 全局状态变量（断连重连后保留）
|
||
total_recorded_frames = 0 # 累计录制总帧数
|
||
current_segment = 0 # 当前视频分段编号
|
||
is_recording = False # 是否正在录制
|
||
current_video_filepath = None# 当前录制视频路径
|
||
|
||
# 单次连接内的临时状态变量
|
||
video_writer = None # 视频写入对象
|
||
recorded_frames = 0 # 当前分段已录制帧数
|
||
frame_count = 0 # 摄像头总读取帧数
|
||
confirmed_frames = [] # 检测通过的确认帧（用于录制起始）
|
||
last_detection_time = time.time() # 上次检测时间
|
||
detection_window_frames = [] # 检测窗口帧缓存
|
||
|
||
# 创建视频目录（确保目录存在）
|
||
os.makedirs(output_video_dir, exist_ok=True)
|
||
# 打印目标信息（此时detection_frame_count已提前定义）
|
||
print(f"✅ 已创建/确认视频目录: {output_video_dir}")
|
||
print(f"🎯 目标：累计录制{total_target_duration}秒，每次录制{single_recording_duration}秒，需{detection_frame_count}帧全为盖板不对齐")
|
||
|
||
# 外层循环：处理摄像头断连与重连
|
||
while True:
|
||
# 初始化摄像头连接
|
||
cap = cv2.VideoCapture(url)
|
||
cap.set(cv2.CAP_PROP_BUFFERSIZE, 5) # 设置RTSP缓存为5MB
|
||
cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'H264')) # 强制H264解码
|
||
|
||
# 摄像头连接重试逻辑
|
||
start_retry_time = time.time()
|
||
while not cap.isOpened():
|
||
# 超过最大重试时间则退出程序
|
||
if time.time() - start_retry_time >= max_retry_seconds:
|
||
print(f"\n❌ 已尝试重新连接 {max_retry_seconds} 秒，仍无法获取视频流，程序退出。")
|
||
# 退出前释放未关闭的视频写入器
|
||
if video_writer is not None:
|
||
video_writer.release()
|
||
print(f"📊 程序退出时累计录制：{total_recorded_frames/video_fps:.1f}秒")
|
||
exit()
|
||
|
||
# 每隔1秒重试一次
|
||
print(f"🔄 无法打开摄像头，正在尝试重新连接...（已重试{int(time.time()-start_retry_time)}秒）")
|
||
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)
|
||
single_recording_frames = video_fps * single_recording_duration
|
||
total_target_frames = video_fps * total_target_duration
|
||
# 打印重连成功信息
|
||
print(f"\n✅ 摄像头重连成功（分辨率：{frame_width}x{frame_height}，实际帧率：{video_fps}）")
|
||
print(f"📊 当前累计录制：{total_recorded_frames / video_fps:.1f}/{total_target_duration}秒")
|
||
|
||
# 重连后恢复录制状态（如果断连前正在录制）
|
||
if is_recording and current_video_filepath:
|
||
print(f"🔄 恢复录制状态，继续录制视频：{current_video_filepath}")
|
||
# 重新初始化视频写入器（确保参数与摄像头匹配）
|
||
video_writer = cv2.VideoWriter(
|
||
current_video_filepath, video_codec, video_fps, (frame_width, frame_height)
|
||
)
|
||
# 恢复失败则重置录制状态
|
||
if not video_writer.isOpened():
|
||
print(f"⚠️ 视频写入器重新初始化失败，无法继续录制")
|
||
is_recording = False
|
||
video_writer = None
|
||
recorded_frames = 0
|
||
|
||
# 内层循环：读取摄像头帧并处理（检测/录制）
|
||
try:
|
||
while True:
|
||
# 读取一帧图像
|
||
ret, frame = cap.read()
|
||
if not ret:
|
||
print(f"\n⚠️ 读取帧失败，可能是流中断或摄像头断开")
|
||
# 断连时保存当前录制进度（不释放全局状态）
|
||
if video_writer is not None:
|
||
video_writer.release()
|
||
video_writer = None
|
||
print(f"⏸️ 流中断，暂停录制（已保存当前进度）")
|
||
# 重置单次连接的临时状态（全局状态保留）
|
||
frame_count = 0
|
||
detection_window_frames = []
|
||
cap.release() # 释放当前摄像头连接
|
||
break # 跳出内层循环，进入重连流程
|
||
|
||
# 累计总帧数
|
||
frame_count += 1
|
||
|
||
# 预处理帧：转为RGB→PIL→旋转180度→转为BGR（适配录制和模型）
|
||
rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||
pil_image = Image.fromarray(rgb_frame)
|
||
rotated_bgr = rotate_frame_180(pil_image)
|
||
|
||
# -------------------------- 1. 未录制时：执行检测逻辑 --------------------------
|
||
if not is_recording:
|
||
# 将帧加入检测窗口缓存（用于10秒后的检测）
|
||
detection_window_frames.append(rotated_bgr)
|
||
# 限制缓存大小（避免内存占用过高，最多保留2倍检测帧数）
|
||
if len(detection_window_frames) > detection_frame_count * 2:
|
||
detection_window_frames = detection_window_frames[-detection_frame_count * 2:]
|
||
|
||
# 触发检测的条件：
|
||
# 1. 距离上次检测超过10秒；2. 检测窗口有足够帧；3. 累计录制未完成
|
||
if (time.time() - last_detection_time) >= detection_interval and \
|
||
len(detection_window_frames) >= detection_frame_count and \
|
||
total_recorded_frames < total_target_frames:
|
||
|
||
try:
|
||
print(f"\n==== 开始10秒间隔检测（总帧：{frame_count}，累计已录：{total_recorded_frames/video_fps:.1f}秒） ====")
|
||
# 从检测窗口中均匀抽取3帧（避免连续帧重复）
|
||
sample_step = max(1, len(detection_window_frames) // detection_frame_count)
|
||
sample_frames = detection_window_frames[::sample_step][:detection_frame_count]
|
||
print(f"📋 检测窗口共{len(detection_window_frames)}帧，均匀抽取{len(sample_frames)}帧进行判断")
|
||
|
||
# 统计“盖板不对齐”的帧数
|
||
noready_frame_count = 0
|
||
valid_detection = True # 标记检测是否有效（无无效帧）
|
||
for idx, sample_frame in enumerate(sample_frames):
|
||
# 调用模型获取分类结果
|
||
class_name = yolov11_cls_inference(cls_model_path, sample_frame, target_size)
|
||
|
||
# 校验分类结果有效性
|
||
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_frame_count += 1
|
||
print(f"✅ 抽取帧{idx+1}：分类结果={class_name}（符合条件）")
|
||
else:
|
||
print(f"❌ 抽取帧{idx+1}：分类结果={class_name}（不符合条件）")
|
||
|
||
# 检测通过条件：所有帧有效 + 3帧全为不对齐
|
||
if valid_detection and noready_frame_count == detection_frame_count:
|
||
print(f"\n✅ 本次检测通过（{noready_frame_count}/{detection_frame_count}帧均为盖板不对齐）")
|
||
# 保存检测通过的帧（用于录制起始，避免丢失检测阶段的画面）
|
||
confirmed_frames.extend(sample_frames)
|
||
|
||
# 检查磁盘空间（剩余<5GB则停止）
|
||
total_disk, used_disk, free_disk = shutil.disk_usage(output_video_dir)
|
||
if free_disk < 1024 * 1024 * 1024 * 5:
|
||
print(f"❌ 磁盘空间严重不足（仅剩 {free_disk / (1024 ** 3):.2f} GB），停止录制并退出。")
|
||
raise SystemExit(1)
|
||
|
||
# 生成当前分段的视频文件名（包含时间戳和分段号）
|
||
current_segment = total_recorded_frames // single_recording_frames + 1
|
||
timestamp = time.strftime("%Y%m%d_%H%M%S")
|
||
current_video_filepath = os.path.join(
|
||
output_video_dir, f"video_{timestamp}_part{current_segment}.mp4"
|
||
)
|
||
|
||
# 初始化视频写入器
|
||
video_writer = cv2.VideoWriter(
|
||
current_video_filepath, video_codec, video_fps, (frame_width, frame_height)
|
||
)
|
||
if not video_writer.isOpened():
|
||
print(f"⚠️ 视频写入器初始化失败（路径：{current_video_filepath}），跳过本次录制")
|
||
confirmed_frames = []
|
||
continue
|
||
|
||
# 写入检测阶段的确认帧（录制起始画面）
|
||
for frame in confirmed_frames:
|
||
video_writer.write(frame)
|
||
recorded_frames = len(confirmed_frames)
|
||
is_recording = True # 标记为正在录制
|
||
|
||
# 打印录制开始信息
|
||
print(f"\n📹 开始录制第{current_segment}段视频（目标10秒）")
|
||
print(f"📁 视频保存路径：{current_video_filepath}")
|
||
print(f"🔢 已写入检测阶段的确认帧：{recorded_frames}帧")
|
||
|
||
# 重置检测相关的临时状态
|
||
confirmed_frames = []
|
||
|
||
# 检测未通过（未满足“3帧全不对齐”或有无效帧）
|
||
else:
|
||
if valid_detection:
|
||
print(f"\n❌ 本次检测未通过（仅{noready_frame_count}/{detection_frame_count}帧为盖板不对齐，需全部符合）")
|
||
else:
|
||
print(f"\n❌ 本次检测未通过（存在无效分类结果）")
|
||
# 重置检测临时状态
|
||
confirmed_frames = []
|
||
|
||
# 更新检测时间，清空检测窗口（准备下一次检测）
|
||
last_detection_time = time.time()
|
||
detection_window_frames = []
|
||
|
||
# 捕获模型调用异常（不终止程序，仅重置检测状态）
|
||
except Exception as e:
|
||
print(f"\n⚠️ 分类模型调用异常: {str(e)}（总帧：{frame_count}）")
|
||
confirmed_frames = []
|
||
last_detection_time = time.time()
|
||
detection_window_frames = []
|
||
continue
|
||
|
||
# -------------------------- 2. 正在录制时：执行写入逻辑 --------------------------
|
||
if is_recording and video_writer is not None:
|
||
# 写入当前帧（已预处理为旋转180度的BGR格式）
|
||
video_writer.write(rotated_bgr)
|
||
recorded_frames += 1
|
||
|
||
# 检查当前分段是否录制完成（达到10秒的帧数）
|
||
if recorded_frames >= single_recording_frames:
|
||
# 释放当前视频写入器（完成当前分段）
|
||
video_writer.release()
|
||
video_writer = None
|
||
is_recording = False # 标记为未录制
|
||
|
||
# 更新累计录制进度
|
||
total_recorded_frames += recorded_frames
|
||
actual_recording_time = recorded_frames / video_fps
|
||
|
||
# 打印分段完成信息
|
||
print(f"\n✅ 第{current_segment}段视频录制完成")
|
||
print(f"🔢 实际录制：{recorded_frames}帧 ≈ {actual_recording_time:.1f}秒")
|
||
print(f"📊 累计录制：{total_recorded_frames/video_fps:.1f}/{total_target_duration}秒")
|
||
|
||
# 检查是否达到总目标（60秒）
|
||
if total_recorded_frames >= total_target_frames:
|
||
print(f"\n🎉 已完成累计{total_target_duration}秒的录制目标！")
|
||
# 重置累计状态（如需重复录制，保留此逻辑；如需单次录制，可添加break退出）
|
||
total_recorded_frames = 0
|
||
current_segment = 0
|
||
|
||
# 重置录制相关的临时状态，准备下一次检测
|
||
recorded_frames = 0
|
||
last_detection_time = time.time()
|
||
detection_window_frames = []
|
||
|
||
# 捕获用户中断（Ctrl+C）
|
||
except KeyboardInterrupt:
|
||
print(f"\n\n👋 用户中断程序")
|
||
# 中断前保存当前录制的视频
|
||
if video_writer is not None:
|
||
video_writer.release()
|
||
print(f"⚠️ 已保存当前录制的视频：{current_video_filepath}")
|
||
print(f"📊 中断时累计录制：{total_recorded_frames/video_fps:.1f}秒")
|
||
break
|
||
|
||
# 最终释放资源（无论内层循环因何退出）
|
||
finally:
|
||
cap.release() # 释放摄像头连接
|
||
if video_writer is not None:
|
||
video_writer.release() # 释放视频写入器
|
||
print(f"\n🔌 视频流已关闭，累计已录：{total_recorded_frames/video_fps:.1f}秒")
|
||
|
||
# 程序结束
|
||
print("\n📋 程序结束") |