feat design 提速测试

fix

app/service/design_test/pipeline/split.py

@@ -0,0 +1,74 @@
+import io
+import logging
+
+import cv2
+import numpy as np
+from PIL import Image
+from cv2 import cvtColor, COLOR_BGR2RGBA
+
+from app.core.config import AIDA_CLOTHING
+from app.service.design.utils.conversion_image import rgb_to_rgba
+from app.service.design.utils.upload_image import upload_png_mask
+from app.service.utils.generate_uuid import generate_uuid
+from app.service.utils.oss_client import oss_upload_image
+
+
+class Split(object):
+    def __init__(self, minio_client):
+        self.minio_client = minio_client
+
+    def __call__(self, result):
+        try:
+
+            if result['name'] in ('outwear', 'dress', 'blouse', 'skirt', 'trousers', 'tops', 'bottoms'):
+                front_mask = result['front_mask']
+                back_mask = result['back_mask']
+                rgba_image = rgb_to_rgba(result['final_image'], front_mask + back_mask)
+                new_size = (int(rgba_image.shape[1] * result["scale"] * result["resize_scale"][0]), int(rgba_image.shape[0] * result["scale"] * result["resize_scale"][1]))
+                rgba_image = cv2.resize(rgba_image, new_size)
+                result_front_image = np.zeros_like(rgba_image)
+                front_mask = cv2.resize(front_mask, new_size)
+                result_front_image[front_mask != 0] = rgba_image[front_mask != 0]
+                result_front_image_pil = Image.fromarray(cvtColor(result_front_image, COLOR_BGR2RGBA))
+                result['front_image'], result["front_image_url"], _ = upload_png_mask(self.minio_client, result_front_image_pil, f'{generate_uuid()}', mask=None)
+
+                height, width = front_mask.shape
+                mask_image = np.zeros((height, width, 3))
+                mask_image[front_mask != 0] = [0, 0, 255]
+
+                if result["name"] in ('blouse', 'dress', 'outwear', 'tops'):
+                    result_back_image = np.zeros_like(rgba_image)
+                    back_mask = cv2.resize(back_mask, new_size)
+                    result_back_image[back_mask != 0] = rgba_image[back_mask != 0]
+                    result_back_image_pil = Image.fromarray(cvtColor(result_back_image, COLOR_BGR2RGBA))
+                    result['back_image'], result["back_image_url"], _ = upload_png_mask(self.minio_client, result_back_image_pil, f'{generate_uuid()}', mask=None)
+                    mask_image[back_mask != 0] = [0, 255, 0]
+
+                    rbga_mask = rgb_to_rgba(mask_image, front_mask + back_mask)
+                    mask_pil = Image.fromarray(cvtColor(rbga_mask.astype(np.uint8), COLOR_BGR2RGBA))
+                    image_data = io.BytesIO()
+                    mask_pil.save(image_data, format='PNG')
+                    image_data.seek(0)
+                    image_bytes = image_data.read()
+                    req = oss_upload_image(oss_client=self.minio_client, bucket=AIDA_CLOTHING, object_name=f"mask/mask_{generate_uuid()}.png", image_bytes=image_bytes)
+                    result['mask_url'] = req.bucket_name + "/" + req.object_name
+                else:
+                    rbga_mask = rgb_to_rgba(mask_image, front_mask)
+                    mask_pil = Image.fromarray(cvtColor(rbga_mask.astype(np.uint8), COLOR_BGR2RGBA))
+                    image_data = io.BytesIO()
+                    mask_pil.save(image_data, format='PNG')
+                    image_data.seek(0)
+                    image_bytes = image_data.read()
+                    req = oss_upload_image(oss_client=self.minio_client, bucket=AIDA_CLOTHING, object_name=f"mask/mask_{generate_uuid()}.png", image_bytes=image_bytes)
+                    result['mask_url'] = req.bucket_name + "/" + req.object_name
+                    result['back_image'] = None
+                    result["back_image_url"] = None
+                    # result["back_mask_url"] = None
+                    # result['back_mask_image'] = None
+            # 创建中间图层
+            result_pattern_image_rgba = rgb_to_rgba(result['pattern_image'], result['mask'])
+            result_pattern_image_pil = Image.fromarray(cvtColor(result_pattern_image_rgba, COLOR_BGR2RGBA))
+            result['pattern_image'], result['pattern_image_url'], _ = upload_png_mask(self.minio_client, result_pattern_image_pil, f'{generate_uuid()}')
+            return result
+        except Exception as e:
+            logging.warning(f"split runtime exception : {e} image_id : {result['image_id']}")