优化

app/models/schp/mask.py

@@ -3,14 +3,14 @@ import numpy as np
 from PIL import Image
 
 
-def extract_color_region_simple(original_image, mask_image, target_color, margin=10):
+def extract_color_region_multiple(original_image, mask_image, target_colors, margin=10):
     """
-    简化版本：提取指定颜色区域，白色背景
+    支持多种颜色的版本：提取指定颜色区域，白色背景
 
     Args:
         original_image: 原始图片 PIL Image 对象
         mask_image: 色块图 PIL Image 对象
-        target_color: 目标颜色 RGB值，例如 [192, 0, 0] (红色)
+        target_colors: 目标颜色列表，例如 [[192, 0, 0], [0, 192, 0], [0, 0, 192]] (红绿蓝)
         margin: 裁剪边距，默认 10
 
     Returns:
@@ -29,23 +29,40 @@ def extract_color_region_simple(original_image, mask_image, target_color, margin
     original_rgb = np.array(original_image.convert('RGB'))
     mask_rgb = np.array(mask_image.convert('RGB'))
 
-    # 创建目标颜色掩码 - 精确匹配
+    # 创建多个目标颜色的组合掩码
+    combined_mask = np.zeros(mask_rgb.shape[:2], dtype=np.uint8)
+
+    found_colors = []
+    for i, target_color in enumerate(target_colors):
         target_color = np.array(target_color)
+        # 为每个颜色创建掩码
         color_mask = np.all(mask_rgb == target_color, axis=2).astype(np.uint8) * 255
 
+        # 统计该颜色的像素数量
+        color_pixels = np.sum(color_mask > 0)
+        if color_pixels > 0:
+            found_colors.append((target_color.tolist(), color_pixels))
+            # 将该颜色的掩码添加到组合掩码中
+            combined_mask = cv2.bitwise_or(combined_mask, color_mask)
+
+    if len(found_colors) == 0:
+        print(f"未找到任何目标颜色区域！")
+        print(f"目标颜色列表: {[color for color in target_colors]}")
+        return None
+
     # 形态学处理
     kernel = np.ones((3, 3), np.uint8)
-    color_mask = cv2.morphologyEx(color_mask, cv2.MORPH_CLOSE, kernel)
-    color_mask = cv2.morphologyEx(color_mask, cv2.MORPH_OPEN, kernel)
+    combined_mask = cv2.morphologyEx(combined_mask, cv2.MORPH_CLOSE, kernel)
+    combined_mask = cv2.morphologyEx(combined_mask, cv2.MORPH_OPEN, kernel)
 
     # 调整掩码尺寸
-    if color_mask.shape != original_rgb.shape[:2]:
-        color_mask = cv2.resize(color_mask, (original_rgb.shape[1], original_rgb.shape[0]))
+    if combined_mask.shape != original_rgb.shape[:2]:
+        combined_mask = cv2.resize(combined_mask, (original_rgb.shape[1], original_rgb.shape[0]))
 
     # 找边界框并裁剪
-    coords = np.column_stack(np.where(color_mask > 0))
+    coords = np.column_stack(np.where(combined_mask > 0))
     if len(coords) == 0:
-        print(f"未找到目标颜色区域 RGB{target_color.tolist()}！")
+        print(f"处理后未找到有效区域！")
         return None
 
     y_min, x_min = coords.min(axis=0)
@@ -59,7 +76,7 @@ def extract_color_region_simple(original_image, mask_image, target_color, margin
 
     # 裁剪
     cropped_original = original_rgb[y_min:y_max, x_min:x_max]
-    cropped_mask = color_mask[y_min:y_max, x_min:x_max]
+    cropped_mask = combined_mask[y_min:y_max, x_min:x_max]
 
     # 创建白色背景结果
     result = np.full_like(cropped_original, 255, dtype=np.uint8)
@@ -70,9 +87,11 @@ def extract_color_region_simple(original_image, mask_image, target_color, margin
     total_pixels = cropped_mask.shape[0] * cropped_mask.shape[1]
 
     print(f"提取完成")
-    print(f"目标颜色: RGB{target_color.tolist()}")
+    print(f"找到的目标颜色:")
+    for color, pixels in found_colors:
+        print(f"  RGB{color}: {pixels} 像素")
     print(f"裁剪区域尺寸: {result.shape[1]} x {result.shape[0]}")
-    print(f"目标颜色像素数量: {mask_pixels}")
+    print(f"总目标颜色像素数量: {mask_pixels}")
     print(f"占裁剪区域的比例: {mask_pixels / total_pixels * 100:.2f}%")
 
     # 转换为PIL Image并返回
@@ -89,16 +108,42 @@ if __name__ == "__main__":
     original_image = Image.open(original_image_path)
     mask_image = Image.open(mask_image_path)
 
-    # 示例: 提取红色区域
-    print("=== 提取红色区域 ===")
-    result_image_red = extract_color_region_simple(
+    # 示例1: 提取多种颜色区域
+    print("=== 提取多种颜色区域 ===")
+    target_colors = [
+        [192, 0, 0],  # 红色
+        [0, 192, 0],  # 绿色
+        [0, 0, 192],  # 蓝色
+        [192, 192, 0],  # 黄色
+    ]
+
+    result_image_multiple = extract_color_region_multiple(
+        original_image,
+        mask_image,
+        target_colors=target_colors
+    )
+
+    if result_image_multiple:
+        result_image_multiple.save(r"D:\work\image_search\extracted_multiple.png")
+        print("多颜色区域提取成功并已保存")
+    else:
+        print("多颜色区域提取失败")
+
+    # 示例2: 只提取红色和绿色
+    print("\n=== 提取红色和绿色区域 ===")
+    red_green_colors = [
+        [192, 0, 0],  # 红色
+        [0, 192, 0],  # 绿色
+    ]
+
+    result_image_rg = extract_color_region_multiple(
         original_image,
         mask_image,
-        target_color=[192, 0, 0]
+        target_colors=red_green_colors
     )
 
-    if result_image_red:
-        result_image_red.save(r"D:\work\image_search\extracted_red.png")
-        print("红色区域提取成功并已保存")
+    if result_image_rg:
+        result_image_rg.save(r"D:\work\image_search\extracted_red_green.png")
+        print("红绿区域提取成功并已保存")
     else:
-        print("红色区域提取失败")
+        print("红绿区域提取失败")

app/services/feature_extractor.py

@@ -8,7 +8,7 @@ from PIL import Image
 from io import BytesIO
 import logging
 
-from app.models.schp.mask import extract_color_region_simple
+from app.models.schp.mask import extract_color_region_multiple
 from app.models.schp.simple_extractor import HumanParsingModel, process_single_image, get_color_by_label
 
 
@@ -108,16 +108,19 @@ class FeatureExtractor:
         # model_name = "ViT-L/14@336px"
         # model, preprocess = self.init_model(device, model_name)
 
-        color = get_color_by_label(part)
+        colors = []
 
+        for item in part.split("|"):
+            color = get_color_by_label(item)
             if color:
+                colors.append(color)
+
+        if colors:
             mask_img = process_single_image(self.schp_model, img)
-            mask_img.save(r"D:\work\image_search\mask.png")
-            extract_img = extract_color_region_simple(img, mask_img, color)
+            extract_img = extract_color_region_multiple(img, mask_img, colors)
             if extract_img:
                 img = extract_img
 
-
         try:
             # 调整图像大小并添加填充
             image = self.resize_with_padding(img)