From fbaa6da92a60e432d2af1022d1626593e94b6eb3 Mon Sep 17 00:00:00 2001
From: zhouchengrong <Zchengrong@yeah.net>
Date: Wed, 28 Aug 2024 11:45:33 +0800
Subject: [PATCH] =?UTF-8?q?feat=20fix=20=20=20design=20=E7=94=A8=E6=88=B7?=
 =?UTF-8?q?=E8=87=AA=E5=AE=9A=E4=B9=89mask?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../design/items/pipelines/segmentation.py    | 42 +++++++++++++------
 1 file changed, 29 insertions(+), 13 deletions(-)

diff --git a/app/service/design/items/pipelines/segmentation.py b/app/service/design/items/pipelines/segmentation.py
index fba6082..ac38b34 100644
--- a/app/service/design/items/pipelines/segmentation.py
+++ b/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 缓存是否存在
-        _, 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']
+        if "seg_mask_url" in result.keys() and result['seg_mask_url'] != "":
+            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")
+            seg_mask = cv2.resize(seg_mask, (result['img_shape'][1], result['img_shape'][0]))
+            # 转换颜色空间为 RGB（OpenCV 默认是 BGR）
+            image_rgb = cv2.cvtColor(seg_mask, cv2.COLOR_BGR2RGB)
+
+            front_color = np.array([255, 0, 0], dtype=np.uint8)
+            back_color = np.array([0, 255, 0], dtype=np.uint8)
+
+            # 创建红色和绿色掩码
+            result['front_mask'] = cv2.inRange(image_rgb, front_color, front_color)
+            result['back_mask'] = cv2.inRange(image_rgb, back_color, back_color)
+            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