60 lines
2.1 KiB
Python
60 lines
2.1 KiB
Python
|
from .builder import ITEMS
|
||
|
|
from .clothing import Clothing
|
||
|
|
|
||
|
|
|
||
|
|
@ITEMS.register_module()
|
||
|
|
class Hairstyle(Clothing):
|
||
|
|
def __init__(self, **kwargs):
|
||
|
|
pipeline = [
|
||
|
|
dict(type='LoadImageFromFile', path=kwargs['path']),
|
||
|
|
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(Hairstyle, self).__init__(**kwargs)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def calculate_start_point(keypoint_type, scale, clothes_point, body_point):
|
||
|
|
"""
|
||
|
|
align up
|
||
|
|
Args:
|
||
|
|
keypoint_type: string, "head_point"
|
||
|
|
scale: float
|
||
|
|
clothes_point: dict{'left': [x1, y1, z1], 'right': [x2, y2, z2]}
|
||
|
|
body_point: dict, containing keypoint data of body figure
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
start_point: tuple (x', y')
|
||
|
|
x' = y_body - y1 * scale
|
||
|
|
y' = x_body - x1 * scale
|
||
|
|
"""
|
||
|
|
side_indicator = f'{keypoint_type}_up'
|
||
|
|
# clothes_point = {k: tuple(map(lambda x: int(scale * x), v[0: 2])) for k, v in clothes_point.items()}
|
||
|
|
# logging.info(clothes_point[side_indicator])
|
||
|
|
|
||
|
|
start_point = (
|
||
|
|
int(body_point[side_indicator][1] - int(clothes_point[side_indicator].split("_")[1] * scale)),
|
||
|
|
int(body_point[side_indicator][0] - int(clothes_point[side_indicator].split("_")[0] * scale))
|
||
|
|
)
|
||
|
|
return start_point
|
||
|
|
|
||
|
|
|
||
|
|
@ITEMS.register_module()
|
||
|
|
class Earring(Clothing):
|
||
|
|
def __init__(self, **kwargs):
|
||
|
|
pipeline = [
|
||
|
|
dict(type='LoadImageFromFile', path=kwargs['path']),
|
||
|
|
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(Earring, self).__init__(**kwargs)
|