版本变更V3.35将图像的处理统一更换到新表当中

inventory-backend/app/utils/ai_vision.py

@@ -9,6 +9,7 @@ import time
 import numpy as np
 from PIL import Image
 import onnxruntime as ort
+import cv2
 
 # ============================================================================
 # 全局模型单例（项目启动时加载一次）
@@ -48,33 +49,115 @@ IMAGENET_STD = [0.229, 0.224, 0.225]
 # 模型输入尺寸
 INPUT_SIZE = 224
 
+# ============================================================================
+# 背景去除配置：HSV 色彩空间阈值
+# ============================================================================
+# OpenCV HSV: H∈[0,180], S∈[0,255], V∈[0,255]
+# 注意：OpenCV 中 H 通道范围是 0-180（是 OpenCV 自己的标准，和美术的 0-360 对应）
 
-def _center_crop_and_resize(image: Image.Image) -> Image.Image:
+# 绿色背景阈值（工业绿幕常用色）
+# H: 35~85 对应绿色谱（浅绿到深绿）
+# S: 低饱和度（35）到高饱和度（255）
+# V: 明暗均可（30~255）
+BG_GREEN_LOWER = np.array([35, 35, 30])
+BG_GREEN_UPPER = np.array([90, 255, 255])
+
+# 白色/浅色背景阈值（高明度、低饱和度区域）
+# H: 不限制（0~180），只看 S 和 V
+# S: 很低的饱和度（0~35）→ 接近纯灰/白色
+# V: 高明度（180~255）
+BG_WHITE_LOWER = np.array([0, 0, 180])
+BG_WHITE_UPPER = np.array([180, 40, 255])
+
+# 中性灰填充色（BGR → 转换后 RGB 也是 128,128,128）
+NEUTRAL_GRAY_BGR = (128, 128, 128)
+
+
+def _remove_background(image: Image.Image) -> Image.Image:
     """
-    CLIP 官方预处理：中心裁剪抗干扰
-    - 将图片最短边缩放到 224
-    - 从正中间切取 224x224 区域
+    利用 OpenCV HSV 色彩空间识别并替换背景为中性灰
+
+    支持两种背景类型：
+    1. 工业绿幕/绿色背景（H: 35~90）
+    2. 白色/浅色背景（高亮度、低饱和度）
+
+    逻辑：
+    - 将 PIL Image 转为 OpenCV 格式 (RGB → BGR)
+    - 转 HSV，分别生成绿色掩码和白色掩码
+    - 合并掩码后，按掩码将背景区域替换为中性灰
+    - 还原为 PIL Image (BGR → RGB) 返回
+
+    参数:
+        image: PIL Image (RGB, uint8)
+
+    返回:
+        处理后的 PIL Image (RGB, uint8)
+    """
+    # PIL (RGB) → OpenCV (BGR)
+    img_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+
+    # 转入 HSV 色彩空间
+    hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV)
+
+    # 生成掩码 1：绿色背景
+    mask_green = cv2.inRange(hsv, BG_GREEN_LOWER, BG_GREEN_UPPER)
+
+    # 生成掩码 2：白色/浅色背景
+    mask_white = cv2.inRange(hsv, BG_WHITE_LOWER, BG_WHITE_UPPER)
+
+    # 合并掩码（任意一种背景都替换）
+    mask_combined = cv2.bitwise_or(mask_green, mask_white)
+
+    # 形态学处理：消除噪点（小面积背景噪点填平）
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
+    mask_combined = cv2.morphologyEx(mask_combined, cv2.MORPH_CLOSE, kernel)
+    mask_combined = cv2.morphologyEx(mask_combined, cv2.MORPH_OPEN, kernel)
+
+    # 背景替换：将掩码区域填充为中性灰
+    # 其中 mask_combined=255 的区域为背景，替换为 NEUTRAL_GRAY_BGR
+    img_bgr_no_bg = img_bgr.copy()
+    img_bgr_no_bg[mask_combined > 0] = NEUTRAL_GRAY_BGR
+
+    # OpenCV (BGR) → PIL (RGB)
+    result = Image.fromarray(cv2.cvtColor(img_bgr_no_bg, cv2.COLOR_BGR2RGB))
+
+    return result
+
+
+def _letterbox_image(image: Image.Image, size: int = 224) -> Image.Image:
+    """
+    Letterbox 预处理：等比例缩放 + 灰色填充，保持内容不变形
+
+    - 将原图最长边缩放到 224
+    - 短边按相同比例缩放
+    - 不足部分用 RGB(128,128,128) 灰色填充至 224x224
+
+    参数:
+        image: PIL Image 对象
+        size: 目标尺寸，默认 224
+
+    返回:
+        224x224 PIL Image
     """
     w, h = image.size
 
-    # 计算缩放后的目标尺寸
-    if w < h:
-        new_w = INPUT_SIZE
-        new_h = int(h * INPUT_SIZE / w)
-    else:
-        new_h = INPUT_SIZE
-        new_w = int(w * INPUT_SIZE / h)
+    # 计算缩放比例，使最长边等于 size
+    scale = size / max(w, h)
+    new_w = int(w * scale)
+    new_h = int(h * scale)
 
-    # 缩放
-    image = image.resize((new_w, new_h), Image.BILINEAR)
+    # 等比例缩放
+    resized = image.resize((new_w, new_h), Image.LANCZOS)
 
-    # 中心裁剪
-    left = (new_w - INPUT_SIZE) // 2
-    top = (new_h - INPUT_SIZE) // 2
-    right = left + INPUT_SIZE
-    bottom = top + INPUT_SIZE
+    # 创建灰色画布
+    canvas = Image.new('RGB', (size, size), (128, 128, 128))
 
-    return image.crop((left, top, right, bottom))
+    # 将缩放后的图片粘贴到画布正中央
+    paste_x = (size - new_w) // 2
+    paste_y = (size - new_h) // 2
+    canvas.paste(resized, (paste_x, paste_y))
+
+    return canvas
 
 
 def _normalize(image_np: np.ndarray) -> np.ndarray:
@@ -111,8 +194,12 @@ def get_image_embedding(image_path: str) -> list:
         load_clip_model()
 
     # 1. 图片预处理
+    #    Step 1: 背景去除（HSV 色彩空间，绿色/白色背景 → 中性灰替换）
     image = Image.open(image_path).convert('RGB')
-    image = _center_crop_and_resize(image)
+    image = _remove_background(image)
+
+    #    Step 2: Letterbox 等比例缩放（保持内容不变形）
+    image = _letterbox_image(image, INPUT_SIZE)
     input_data = _normalize(np.array(image))
     input_data = np.expand_dims(input_data, axis=0)  # [1, 3, 224, 224]