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