更新charge振捣判断

2026-03-10 16:51:57 +08:00
parent 235101b4d8
commit 5d79686ba0
15 changed files with 268873 additions and 0 deletions
--- a/image_01.py
+++ b/image_01.py
@ -0,0 +1,170 @@
+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} 张图片。")