AiDA_Python/app/service/design/items/pipelines/split.py

import logging

import cv2
import numpy as np
from PIL import Image
from cv2 import cvtColor, COLOR_BGR2RGBA

from app.service.utils.decorator import ClassCallRunTime
from app.service.utils.generate_uuid import generate_uuid
from ..builder import PIPELINES
from ...utils.conversion_image import rgb_to_rgba
from ...utils.upload_image import upload_png_mask


@PIPELINES.register_module()
class Split(object):
    """
    Split image into front and back layer according to the segmentation result
    """

    @ClassCallRunTime
    # KNet
    def __call__(self, result):
        try:
            if 'mask' not in result.keys():
                raise KeyError(f'Cannot find mask in result dict, please check ContourDetection is included in process pipelines.')
            if 'seg_result' not in result.keys():  # 没过seg模型
                result['front_mask'] = result['mask'].copy()
                result['back_mask'] = np.zeros_like(result['mask'])
            else:
                temp_front = result['seg_result'] == 1
                result['front_mask'] = (result['mask'] * (temp_front + 0).astype(np.uint8))
                temp_back = result['seg_result'] == 2
                result['back_mask'] = (result['mask'] * (temp_back + 0).astype(np.uint8))

            if result['name'] in ('outwear', 'dress', 'blouse', 'skirt', 'trousers', 'tops', 'bottoms'):
                if len(result['front_mask'].shape) > 2:
                    front_mask = result['front_mask'][0]
                else:
                    front_mask = result['front_mask']

                if len(result['back_mask'].shape) > 2:
                    back_mask = result['back_mask'][0]
                else:
                    back_mask = result['back_mask']

                # rgba_image = rgb_to_rgba((result['final_image'].shape[0], result['final_image'].shape[1]), result['final_image'], front_mask + 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"], result["front_mask_url"] = upload_png_mask(result_front_image_pil, f'{generate_uuid()}', mask=front_mask)
                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"], result["back_mask_url"] = upload_png_mask(result_back_image_pil, f'{generate_uuid()}', mask=back_mask)
                else:
                    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'], result['pattern_image_url'], _ = upload_png_mask(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']}")