171 lines
5.6 KiB
Python
171 lines
5.6 KiB
Python
import cv2
|
|
import time
|
|
import os
|
|
import numpy as np
|
|
from PIL import Image
|
|
from skimage.metrics import structural_similarity as ssim
|
|
|
|
# ================== 配置区域 ==================
|
|
RTSP_URL = "rtsp://admin:XJ123456@192.168.10.50:554/streaming/channels/101"
|
|
|
|
SAVE_INTERVAL = 10 # 每 20 帧尝试一次处理
|
|
DISPLAY_STREAM = False # 是否显示画面
|
|
|
|
# --- 灰图过滤配置 ---
|
|
GRAY_LOWER = 70
|
|
GRAY_UPPER = 230
|
|
GRAY_RATIO_THRESHOLD = 0.7 # 灰色像素占比超过此值视为灰图
|
|
|
|
# --- SSIM 去重配置 ---
|
|
SSIM_THRESHOLD = 0.9 # 相似度超过此值视为重复图片
|
|
|
|
# --- 保存目录 ---
|
|
OUTPUT_DIR = "camera02_save"
|
|
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
|
|
|
# ================== 辅助函数 ==================
|
|
|
|
def is_large_gray(image):
|
|
"""
|
|
检测图像是否为大面积灰图
|
|
image: PIL Image 对象 (RGB)
|
|
"""
|
|
arr = np.array(image)
|
|
# 检查 R, G, B 三个通道是否都在灰色范围内
|
|
gray_mask = (
|
|
(arr[:, :, 0] >= GRAY_LOWER) & (arr[:, :, 0] <= GRAY_UPPER) &
|
|
(arr[:, :, 1] >= GRAY_LOWER) & (arr[:, :, 1] <= GRAY_UPPER) &
|
|
(arr[:, :, 2] >= GRAY_LOWER) & (arr[:, :, 2] <= GRAY_UPPER)
|
|
)
|
|
# 计算灰色像素占比
|
|
return np.mean(gray_mask) > GRAY_RATIO_THRESHOLD
|
|
|
|
def open_camera(url):
|
|
cap = cv2.VideoCapture(url)
|
|
if not cap.isOpened():
|
|
return None
|
|
cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
|
|
cap.set(cv2.CAP_PROP_OPEN_TIMEOUT_MSEC, 5000)
|
|
cap.set(cv2.CAP_PROP_READ_TIMEOUT_MSEC, 5000)
|
|
return cap
|
|
|
|
def get_valid_frame(cap):
|
|
"""安全读取一帧"""
|
|
if cap is None or not cap.isOpened():
|
|
return None, False
|
|
|
|
ret, frame = cap.read()
|
|
if not ret or frame is None or frame.size == 0:
|
|
return None, False
|
|
return frame, True
|
|
|
|
# ================== 主程序 ==================
|
|
|
|
print(f"✅ 正在连接摄像头: {RTSP_URL} ...")
|
|
cap = open_camera(RTSP_URL)
|
|
|
|
if cap is None:
|
|
print(f"❌ 连接失败！请检查 IP、账号密码或网络。")
|
|
exit(1)
|
|
|
|
print("📡 摄像头已连接，开始采集...")
|
|
print(f" - 保存目录: {os.path.abspath(OUTPUT_DIR)}")
|
|
print(f" - 灰度阈值: [{GRAY_LOWER}, {GRAY_UPPER}], 占比 > {GRAY_RATIO_THRESHOLD}")
|
|
print(f" - SSIM 去重阈值: > {SSIM_THRESHOLD}")
|
|
print(" - 按 Ctrl+C 停止\n")
|
|
|
|
frame_count = 0
|
|
saved_count = 0
|
|
last_gray_frame = None # 用于存储上一帧的灰度图以计算 SSIM
|
|
|
|
try:
|
|
while True:
|
|
frame, ret = get_valid_frame(cap)
|
|
|
|
if not ret:
|
|
print("⚠️ 读取帧失败，尝试重连...")
|
|
time.sleep(2)
|
|
cap.release()
|
|
cap = open_camera(RTSP_URL)
|
|
if cap is None:
|
|
print("❌ 重连失败，退出程序。")
|
|
break
|
|
last_gray_frame = None # 重连后重置对比帧
|
|
continue
|
|
|
|
frame_count += 1
|
|
|
|
# 间隔采样：不是指定间隔的帧直接跳过（可选显示）
|
|
if frame_count % SAVE_INTERVAL != 0:
|
|
if DISPLAY_STREAM:
|
|
cv2.imshow("Camera Stream", frame)
|
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
|
break
|
|
continue
|
|
|
|
# --- 步骤 1: 灰图检测 ---
|
|
try:
|
|
# 转 RGB 供 PIL 使用
|
|
img_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
pil_img = Image.fromarray(img_rgb)
|
|
|
|
if is_large_gray(pil_img):
|
|
print(f"⏭️ 跳过：检测到灰图 (帧 {frame_count})")
|
|
continue
|
|
except Exception as e:
|
|
print(f"⚠️ 灰图检测异常: {e}")
|
|
continue
|
|
|
|
# --- 步骤 2: SSIM 相似性去重 ---
|
|
try:
|
|
gray_curr = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
|
|
|
if last_gray_frame is not None:
|
|
# 计算结构相似性
|
|
# resize 确保尺寸一致（防止分辨率动态变化导致报错）
|
|
if last_gray_frame.shape != gray_curr.shape:
|
|
last_gray_frame = cv2.resize(last_gray_frame, (gray_curr.shape[1], gray_curr.shape[0]))
|
|
|
|
sim_score = ssim(last_gray_frame, gray_curr)
|
|
|
|
if sim_score > SSIM_THRESHOLD:
|
|
print(f"⏭️ 跳过：画面重复 (SSIM={sim_score:.3f}, 帧 {frame_count})")
|
|
continue
|
|
|
|
# 更新参考帧
|
|
last_gray_frame = gray_curr.copy()
|
|
|
|
except Exception as e:
|
|
print(f"⚠️ SSIM 计算异常: {e}")
|
|
# 如果计算出错，可以选择跳过或强制保存，这里选择跳过以防崩溃
|
|
continue
|
|
|
|
# --- 步骤 3: 保存图片 ---
|
|
ts = time.strftime("%Y%m%d_%H%M%S")
|
|
ms = int((time.time() % 1) * 1000)
|
|
filename = f"img_{ts}_{ms:03d}.png"
|
|
save_path = os.path.join(OUTPUT_DIR, filename)
|
|
|
|
# 如果需要旋转/翻转，在这里操作 (例如翻转 180 度)
|
|
# frame = cv2.flip(frame, -1)
|
|
|
|
cv2.imwrite(save_path, frame)
|
|
saved_count += 1
|
|
print(f"✅ [{saved_count}] 已保存有效图片: {filename}")
|
|
|
|
if DISPLAY_STREAM:
|
|
cv2.putText(frame, f"Saved: {filename}", (10, 30),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
|
|
cv2.imshow("Camera Stream", frame)
|
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
|
break
|
|
|
|
except KeyboardInterrupt:
|
|
print("\n🛑 用户手动停止")
|
|
|
|
finally:
|
|
if cap:
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|
|
print(f"🔚 程序结束。共保存 {saved_count} 张图片。")
|