Files
zjsh_yolov11/image/image_upload.py
2025-08-13 14:49:06 +08:00

165 lines
6.1 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import cv2
import time
import os
import numpy as np
from PIL import Image
from skimage.metrics import structural_similarity as ssim
import threading
import transport_client
from transport_client import send_images # 导入上传函数
# ================== 配置数 ==================
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
# 上传配置(新增)
UPLOAD_SERVER_URL = "http://www.xj-robot.com:6000/upload"
UPLOAD_SITE_NAME = "FactoryA"
UPLOAD_LINE_ID = "Line1"
UPLOAD_PURPOSE = "DET"
# 创建输出目录
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}")
# ✅ 新增:异步上传
def upload_task():
send_images(
folder_path=output_dir,
server_url=UPLOAD_SERVER_URL,
site_name=UPLOAD_SITE_NAME,
line_id=UPLOAD_LINE_ID,
purpose=UPLOAD_PURPOSE
)
threading.Thread(target=upload_task, daemon=True).start()
print(f"📤 已提交上传任务: {filename}")
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("程序结束")