127 lines
6.4 KiB
Python
127 lines
6.4 KiB
Python
import io
|
|
import logging
|
|
import time
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from .builder import ITEMS
|
|
from .clothing import Clothing
|
|
from PIL import Image
|
|
|
|
from ..utils.conversion_image import rgb_to_rgba
|
|
from ..utils.upload_image import upload_png_mask
|
|
from ...utils.generate_uuid import generate_uuid
|
|
|
|
|
|
@ITEMS.register_module()
|
|
class Shoes(Clothing):
|
|
# TODO location of shoes has little mismatch
|
|
def __init__(self, **kwargs):
|
|
pipeline = [
|
|
dict(type='LoadImageFromFile', path=kwargs['path'], color=kwargs['color']),
|
|
dict(type='KeypointDetection'),
|
|
dict(type='ContourDetection'),
|
|
dict(type='Painting'),
|
|
dict(type='Scaling'),
|
|
dict(type='Split'),
|
|
# dict(type='ImageShow', key=['image', 'mask', 'pattern_image']),
|
|
]
|
|
kwargs.update(pipeline=pipeline)
|
|
super(Shoes, self).__init__(**kwargs)
|
|
|
|
def organize(self, layer):
|
|
left_shoe_mask, right_shoe_mask = self.cut()
|
|
|
|
left_layer = dict(name=f'{type(self).__name__.lower()}_left',
|
|
image=self.result['shoes_left'],
|
|
image_url=self.result['left_image_url'],
|
|
mask_url=self.result['left_mask_url'],
|
|
sacle=self.result['scale'],
|
|
clothes_keypoint=self.result['clothes_keypoint'],
|
|
position=self.calculate_start_point(self.result['keypoint'],
|
|
self.result['scale'],
|
|
self.result['clothes_keypoint'],
|
|
self.result['body_point'],
|
|
'left'))
|
|
layer.insert(left_layer)
|
|
|
|
right_layer = dict(name=f'{type(self).__name__.lower()}_right',
|
|
image=self.result['shoes_right'],
|
|
image_url=self.result['right_image_url'],
|
|
mask_url=self.result['right_mask_url'],
|
|
sacle=self.result['scale'],
|
|
clothes_keypoint=self.result['clothes_keypoint'],
|
|
position=self.calculate_start_point(self.result['keypoint'],
|
|
self.result['scale'],
|
|
self.result['clothes_keypoint'],
|
|
self.result['body_point'],
|
|
'right'))
|
|
|
|
layer.insert(right_layer)
|
|
|
|
def cut(self):
|
|
"""
|
|
Cut shoes mask into two pieces
|
|
Returns:
|
|
"""
|
|
contour, _ = cv2.findContours(self.result['mask'], cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
contours = sorted(contour, key=cv2.contourArea, reverse=True)
|
|
|
|
bounding_boxes = [cv2.boundingRect(c) for c in contours[:2]]
|
|
(contours, bounding_boxes) = zip(*sorted(zip(contours[:2], bounding_boxes), key=lambda x: x[1][0], reverse=False))
|
|
|
|
epsilon_left = 0.001 * cv2.arcLength(contours[0], True)
|
|
|
|
approx_left = cv2.approxPolyDP(contours[0], epsilon_left, True)
|
|
mask_left = np.zeros(self.result['final_image'].shape[:2], np.uint8)
|
|
cv2.drawContours(mask_left, [approx_left], -1, 255, -1)
|
|
item_mask_left = cv2.GaussianBlur(mask_left, (5, 5), 0)
|
|
|
|
rgba_image = rgb_to_rgba((self.result['final_image'].shape[0], self.result['final_image'].shape[1]), self.result['final_image'], item_mask_left)
|
|
result_image = np.zeros_like(rgba_image)
|
|
result_image[self.result['front_mask'] != 0] = rgba_image[self.result['front_mask'] != 0]
|
|
result_left_image_pil = Image.fromarray(result_image, 'RGBA')
|
|
result_left_image_pil = result_left_image_pil.resize((int(result_left_image_pil.width * self.result["scale"]), int(result_left_image_pil.height * self.result["scale"])), Image.LANCZOS)
|
|
self.result['shoes_left'], self.result["left_image_url"], self.result["left_mask_url"] = upload_png_mask(result_left_image_pil, f"{generate_uuid()}")
|
|
|
|
epsilon_right = 0.001 * cv2.arcLength(contours[1], True)
|
|
approx_right = cv2.approxPolyDP(contours[1], epsilon_right, True)
|
|
mask_right = np.zeros(self.result['final_image'].shape[:2], np.uint8)
|
|
cv2.drawContours(mask_right, [approx_right], -1, 255, -1)
|
|
item_mask_right = cv2.GaussianBlur(mask_right, (5, 5), 0)
|
|
|
|
rgba_image = rgb_to_rgba((self.result['final_image'].shape[0], self.result['final_image'].shape[1]), self.result['final_image'], item_mask_right)
|
|
result_image = np.zeros_like(rgba_image)
|
|
result_image[self.result['front_mask'] != 0] = rgba_image[self.result['front_mask'] != 0]
|
|
result_right_image_pil = Image.fromarray(result_image, 'RGBA')
|
|
result_right_image_pil = result_right_image_pil.resize((int(result_right_image_pil.width * self.result["scale"]), int(result_right_image_pil.height * self.result["scale"])), Image.LANCZOS)
|
|
self.result['shoes_right'], self.result["right_image_url"], self.result["right_mask_url"] = upload_png_mask(result_right_image_pil, f"{generate_uuid()}")
|
|
|
|
return item_mask_left, item_mask_right
|
|
|
|
@staticmethod
|
|
def calculate_start_point(keypoint_type, scale, clothes_point, body_point, location):
|
|
"""
|
|
left shoes align left
|
|
right shoes align right
|
|
Args:
|
|
keypoint_type: string, "toe"
|
|
scale: float
|
|
clothes_point: dict{'left': [x1, y1, z1], 'right': [x2, y2, z2]}
|
|
body_point: dict, containing keypoint data of body figure
|
|
location: string, indicates whether the start point belongs to right or left shoe
|
|
|
|
Returns:
|
|
start_point: tuple (x', y')
|
|
x' = y_body - y1 * scale
|
|
y' = x_body - x1 * scale
|
|
"""
|
|
if location not in ['left', 'right']:
|
|
raise KeyError(f'location value must be left or right but got {location}')
|
|
side_indicator = f'{keypoint_type}_{location}'
|
|
# clothes_point = {k: tuple(map(lambda x: int(scale * x), v[0: 2])) for k, v in clothes_point.items()}
|
|
start_point = (body_point[side_indicator][1] - int(int(clothes_point[side_indicator].split("_")[1]) * scale),
|
|
body_point[side_indicator][0] - int(int(clothes_point[side_indicator].split("_")[0]) * scale))
|
|
return start_point
|