Files
zjsh_yolov11/image/image.py

145 lines
5.3 KiB
Python
Raw Normal View History

2025-08-13 14:49:06 +08:00
import cv2
import time
import os
import numpy as np
from PIL import Image
from skimage.metrics import structural_similarity as ssim
# ================== 配置参数 ==================
url = "rtsp://admin:XJ123456@192.168.1.51:554/streaming/channels/101"
save_interval = 15 # 每隔 N 帧处理一次(可调)
SSIM_THRESHOLD = 0.9 # SSIM 相似度阈值,>0.9 认为太像
output_dir = os.path.join("userdata", "image") # 固定路径userdata/image
# 灰色判断参数
GRAY_LOWER = 70
GRAY_UPPER = 230
GRAY_RATIO_THRESHOLD = 0.7
# 创建输出目录
if not os.path.exists(output_dir):
os.makedirs(output_dir)
print(f"已创建目录: {output_dir}")
def is_large_gray(image, gray_lower=GRAY_LOWER, gray_upper=GRAY_UPPER, ratio_thresh=GRAY_RATIO_THRESHOLD):
"""
判断图片是否大面积为灰色R/G/B 都在 [gray_lower, gray_upper] 区间
"""
img_array = np.array(image)
if len(img_array.shape) != 3 or img_array.shape[2] != 3:
return True # 非三通道图视为无效/灰色
h, w, _ = img_array.shape
total = h * w
gray_mask = (
(img_array[:, :, 0] >= gray_lower) & (img_array[:, :, 0] <= gray_upper) &
(img_array[:, :, 1] >= gray_lower) & (img_array[:, :, 1] <= gray_upper) &
(img_array[:, :, 2] >= gray_lower) & (img_array[:, :, 2] <= gray_upper)
)
gray_pixels = np.sum(gray_mask)
gray_ratio = gray_pixels / total
return gray_ratio > ratio_thresh
max_retry_seconds = 10 # 最大重试时间为10秒
retry_interval_seconds = 1 # 每隔1秒尝试重新连接一次
while True: # 外层循环用于处理重新连接逻辑
cap = cv2.VideoCapture(url)
start_time = time.time() # 记录开始尝试连接的时间
while not cap.isOpened(): # 如果无法打开摄像头,则进入重试逻辑
if time.time() - start_time >= max_retry_seconds:
print(f"已尝试重新连接 {max_retry_seconds} 秒,但仍无法获取视频流。")
exit()
print("无法打开摄像头,正在尝试重新连接...")
time.sleep(retry_interval_seconds) # 等待一段时间后再次尝试
cap = cv2.VideoCapture(url)
print("✅ 开始读取视频流...")
frame_count = 0
last_gray = None # 用于 SSIM 去重
try:
while True:
ret, frame = cap.read()
if not ret:
print("读取帧失败,可能是流中断或摄像头断开")
cap.release() # 释放资源以便重新连接
break # 跳出内层循环尝试重新连接
frame_count += 1
# 仅在指定间隔处理保存逻辑
if frame_count % save_interval != 0:
cv2.imshow('Camera Stream (Live)', frame)
if cv2.waitKey(1) == ord('q'):
raise KeyboardInterrupt
continue
print(f"处理帧 {frame_count}")
# 转为 PIL 图像(用于后续判断和旋转)
rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
pil_image = Image.fromarray(rgb_frame)
# STEP 1: 判断是否为大面积灰色(优先级最高)
if is_large_gray(pil_image):
print(f"跳过:大面积灰色图像 (frame_{frame_count})")
cv2.imshow('Camera Stream (Live)', frame)
if cv2.waitKey(1) == ord('q'):
raise KeyboardInterrupt
continue
# STEP 2: 判断是否为重复帧(基于 SSIM
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
if last_gray is not None:
try:
similarity = ssim(gray, last_gray)
if similarity > SSIM_THRESHOLD:
print(f"跳过:与上一帧太相似 (SSIM={similarity:.3f})")
cv2.imshow('Camera Stream (Live)', frame)
if cv2.waitKey(1) == ord('q'):
raise KeyboardInterrupt
continue
except Exception as e:
print(f"SSIM 计算异常: {e}")
# 更新 last_gray 用于下一帧比较
last_gray = gray.copy()
# STEP 3: 旋转 180 度
rotated_pil = pil_image.rotate(180, expand=False)
# 生成文件名(时间戳 + 毫秒防重),使用 .png 扩展名
timestamp = time.strftime("%Y%m%d_%H%M%S")
ms = int((time.time() % 1) * 1000)
filename = f"frame_{timestamp}_{ms:03d}.png" # ✅ 确保是 .png
filepath = os.path.join(output_dir, filename)
# 保存图像为 PNG 格式(无损)
try:
rotated_pil.save(filepath, format='PNG') # ✅ PNG 不需要 quality 参数
print(f"已保存: {filepath}")
except Exception as e:
print(f"保存失败 {filename}: {e}")
# 显示画面
cv2.imshow('Camera Stream (Live)', frame)
if cv2.waitKey(1) == ord('q'):
raise KeyboardInterrupt
except KeyboardInterrupt:
print("\n用户中断")
break # 跳出外层循环并退出程序
finally:
cap.release()
cv2.destroyAllWindows()
print(f"视频流已关闭,共处理 {frame_count} 帧。")
print("程序结束")