116 lines
6.9 KiB
Python
116 lines
6.9 KiB
Python
import logging
|
|
import cv2
|
|
import numpy as np
|
|
from cv2 import cvtColor, COLOR_BGR2RGBA
|
|
from app.service.utils.generate_uuid import generate_uuid
|
|
from ..builder import PIPELINES
|
|
from PIL import Image
|
|
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
|
|
"""
|
|
|
|
# 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'], result['mask'])
|
|
result_front_image = np.zeros_like(rgba_image)
|
|
result_front_image[front_mask != 0] = rgba_image[front_mask != 0]
|
|
|
|
result_front_image_pil = Image.fromarray(cvtColor(result_front_image, COLOR_BGR2RGBA))
|
|
front_new_size = (int(result_front_image_pil.width * result["scale"] * result["resize_scale"]), int(result_front_image_pil.height * result["scale"] * result["resize_scale"]))
|
|
result_front_image_pil = result_front_image_pil.resize(front_new_size, Image.LANCZOS)
|
|
front_mask = cv2.resize(front_mask, front_new_size)
|
|
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)
|
|
result_back_image[back_mask != 0] = rgba_image[back_mask != 0]
|
|
|
|
result_back_image_pil = Image.fromarray(cvtColor(result_back_image, COLOR_BGR2RGBA))
|
|
back_new_size = (int(result_back_image_pil.width * result["scale"] * result["resize_scale"]), int(result_back_image_pil.height * result["scale"] * result["resize_scale"]))
|
|
result_back_image_pil = result_back_image_pil.resize(back_new_size, Image.LANCZOS)
|
|
back_mask = cv2.resize(back_mask, back_new_size)
|
|
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
|
|
return result
|
|
except Exception as e:
|
|
logging.warning(f"split runtime exception : {e} image_id : {result['image_id']}")
|
|
|
|
# @ RunTime
|
|
# 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']
|
|
#
|
|
# rgba_image = rgb_to_rgba((result['final_image'].shape[0], result['final_image'].shape[1]), result['final_image'], result['mask'])
|
|
# result_front_image = np.zeros_like(rgba_image)
|
|
# result_front_image[front_mask != 0] = rgba_image[front_mask != 0]
|
|
#
|
|
# result_front_image_pil = Image.fromarray(cvtColor(result_front_image, COLOR_BGR2RGBA))
|
|
# front_new_size = (int(result_front_image_pil.width * result["scale"] * result["resize_scale"]), int(result_front_image_pil.height * result["scale"] * result["resize_scale"]))
|
|
# result_front_image_pil = result_front_image_pil.resize(front_new_size, Image.LANCZOS)
|
|
# front_mask = cv2.resize(front_mask, front_new_size)
|
|
# 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)
|
|
# result_back_image[result['back_mask'] != 0] = rgba_image[result['back_mask'] != 0]
|
|
#
|
|
# result_back_image_pil = Image.fromarray(cvtColor(result_back_image, COLOR_BGR2RGBA))
|
|
# back_new_size = (int(result_back_image_pil.width * result["scale"] * result["resize_scale"]), int(result_back_image_pil.height * result["scale"] * result["resize_scale"]))
|
|
# result_back_image_pil = result_back_image_pil.resize(back_new_size, Image.LANCZOS)
|
|
# back_mask = cv2.resize(result['back_mask'], back_new_size)
|
|
# 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
|
|
# return result
|
|
# except Exception as e:
|
|
# logging.warning(f"split runtime exception : {e} image_id : {result['image_id']}")
|