feat

fix   design 用户自定义mask

app/service/design/items/pipelines/segmentation.py

@@ -1,9 +1,11 @@
 import logging
 import os
 
+import cv2
 import numpy as np
 
 from app.core.config import SEG_CACHE_PATH
+from app.service.utils.oss_client import oss_get_image
 from ..builder import PIPELINES
 from ...utils.design_ensemble import get_seg_result
 
@@ -15,19 +17,33 @@ class Segmentation(object):
 
     # @ClassCallRunTime
     def __call__(self, result):
-        # 本地查询seg 缓存是否存在
+        if "seg_mask_url" in result.keys() and result['seg_mask_url'] != "":
-        _, seg_result = self.load_seg_result(result["image_id"])
+            seg_mask = oss_get_image(bucket=result['seg_mask_url'].split('/')[0], object_name=result['seg_mask_url'][result['seg_mask_url'].find('/') + 1:], data_type="cv2")
-        result['seg_result'] = seg_result
+            seg_mask = cv2.resize(seg_mask, (result['img_shape'][1], result['img_shape'][0]))
-        if not _:
+            # 转换颜色空间为 RGB（OpenCV 默认是 BGR）
-            # 推理获得seg 结果
+            image_rgb = cv2.cvtColor(seg_mask, cv2.COLOR_BGR2RGB)
-            seg_result = get_seg_result(result["image_id"], result['image'])[0]
+
-            self.save_seg_result(seg_result, result['image_id'])
+            front_color = np.array([255, 0, 0], dtype=np.uint8)
-        # 处理前片后片
+            back_color = np.array([0, 255, 0], dtype=np.uint8)
-        temp_front = seg_result == 1.0
+
-        result['front_mask'] = (255 * (temp_front + 0).astype(np.uint8))
+            # 创建红色和绿色掩码
-        temp_back = seg_result == 2.0
+            result['front_mask'] = cv2.inRange(image_rgb, front_color, front_color)
-        result['back_mask'] = (255 * (temp_back + 0).astype(np.uint8))
+            result['back_mask'] = cv2.inRange(image_rgb, back_color, back_color)
-        result['mask'] = result['front_mask'] + result['back_mask']
+            result['mask'] = result['front_mask'] + result['back_mask']
+        else:
+            # 本地查询seg 缓存是否存在
+            _, seg_result = self.load_seg_result(result["image_id"])
+            result['seg_result'] = seg_result
+            if not _:
+                # 推理获得seg 结果
+                seg_result = get_seg_result(result["image_id"], result['image'])[0]
+                self.save_seg_result(seg_result, result['image_id'])
+            # 处理前片后片
+            temp_front = seg_result == 1.0
+            result['front_mask'] = (255 * (temp_front + 0).astype(np.uint8))
+            temp_back = seg_result == 2.0
+            result['back_mask'] = (255 * (temp_back + 0).astype(np.uint8))
+            result['mask'] = result['front_mask'] + result['back_mask']
         return result
 
     @staticmethod