56 lines
2.8 KiB
Python
56 lines
2.8 KiB
Python
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 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"], 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_image'], 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']}")
|