AiDA_Python/app/service/design_test/item.py

import time
from concurrent.futures import ThreadPoolExecutor
from pprint import pprint

import cv2

from app.core.config import PRIORITY_DICT
from app.service.design.utils.synthesis_item import synthesis, synthesis_single
from app.service.design_test.pipeline import LoadImage, KeyPoint, Segmentation, Color, PrintPainting, Scaling, Split, LoadBodyImage, ContourDetection


class BaseItem:
    def __init__(self, data, basic):
        self.result = data.copy()
        self.result['name'] = data['type'].lower()
        self.result.pop("type")
        self.result.update(basic)


class TopItem(BaseItem):
    def __init__(self, data, basic, minio_client):
        super().__init__(data, basic)
        self.top_pipeline = [
            LoadImage(minio_client),
            KeyPoint(),
            Segmentation(minio_client),
            Color(minio_client),
            PrintPainting(minio_client),
            Scaling(),
            Split(minio_client)
        ]

    def process(self):
        for item in self.top_pipeline:
            self.result = item(self.result)
        return self.result


class BottomItem(BaseItem):
    def __init__(self, data, basic, minio_client):
        super().__init__(data, basic)
        self.bottom_pipeline = [
            LoadImage(minio_client),
            KeyPoint(),
            ContourDetection(),
            # Segmentation(),
            Color(minio_client),
            PrintPainting(minio_client),
            Scaling(),
            Split(minio_client)
        ]

    def process(self):
        for item in self.bottom_pipeline:
            self.result = item(self.result)
        return self.result


class BodyItem(BaseItem):
    def __init__(self, data, basic, minio_client):
        super().__init__(data, basic)
        self.top_pipeline = [
            LoadBodyImage(minio_client),
        ]

    def process(self):
        for item in self.top_pipeline:
            self.result = item(self.result)
        return self.result


def process_item(item, basic, minio_client):
    if item['type'] == "Body":
        body_server = BodyItem(data=item, basic=basic, minio_client=minio_client)
        item_data = body_server.process()
    elif item['type'].lower() in ['blouse', 'outwear', 'dress', 'tops']:
        top_server = TopItem(data=item, basic=basic, minio_client=minio_client)
        item_data = top_server.process()
    else:
        bottom_server = BottomItem(data=item, basic=basic, minio_client=minio_client)
        item_data = bottom_server.process()
    return item_data


def calculate_start_point(keypoint_type, scale, clothes_point, body_point, offset, resize_scale):
    """
    Align left
    Args:
        keypoint_type: string, "waistband" | "shoulder" | "ear_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 + offset
        y' = x_body - x1 * scale + offset

    """
    side_indicator = f'{keypoint_type}_left'
    start_point = (
        int(body_point[side_indicator][1] + offset[1] - int(clothes_point[side_indicator][0]) * scale),  # y
        int(body_point[side_indicator][0] + offset[0] - int(clothes_point[side_indicator][1]) * scale)  # x
    )
    return start_point


# 服装图层给数据组装
def organize_clothing(layer):
    # 起始坐标
    start_point = calculate_start_point(layer['keypoint'], layer['scale'], layer['clothes_keypoint'], layer['body_point_test'], layer["offset"], layer["resize_scale"])
    # 前片数据
    front_layer = dict(priority=layer['priority'] if layer.get("layer_order", False) else PRIORITY_DICT.get(f'{layer["name"].lower()}_front', None),
                       name=f'{layer["name"].lower()}_front',
                       image=layer["front_image"],
                       # mask_image=layer['front_mask_image'],
                       image_url=layer['front_image_url'],
                       mask_url=layer['mask_url'],
                       sacle=layer['scale'],
                       clothes_keypoint=layer['clothes_keypoint'],
                       position=start_point,
                       resize_scale=layer["resize_scale"],
                       mask=cv2.resize(layer['mask'], layer["front_image"].size),
                       gradient_string=layer['gradient_string'] if 'gradient_string' in layer.keys() else "",
                       pattern_image_url=layer['pattern_image_url'],
                       pattern_image=layer['pattern_image']

                       )
    # 后片数据
    back_layer = dict(priority=-layer.get("priority", 0) if layer.get("layer_order", False) else PRIORITY_DICT.get(f'{layer["name"].lower()}_back', None),
                      name=f'{layer["name"].lower()}_back',
                      image=layer["back_image"],
                      # mask_image=layer['back_mask_image'],
                      image_url=layer['back_image_url'],
                      mask_url=layer['mask_url'],
                      sacle=layer['scale'],
                      clothes_keypoint=layer['clothes_keypoint'],
                      position=start_point,
                      resize_scale=layer["resize_scale"],
                      mask=cv2.resize(layer['mask'], layer["front_image"].size),
                      gradient_string=layer['gradient_string'] if 'gradient_string' in layer.keys() else "",
                      pattern_image_url=layer['pattern_image_url'],
                      )
    return front_layer, back_layer


# 模特图层给数据组装
def organize_body(layer):
    body_layer = dict(priority=0,
                      name=layer["name"].lower(),
                      image=layer['body_image'],
                      image_url=layer['body_path'],
                      mask_image=None,
                      mask_url=None,
                      sacle=1,
                      # mask=layer['body_mask'],
                      position=(0, 0))
    return body_layer


def process_layer(item, layers):
    if item['name'] == "mannequin":
        body_layer = organize_body(item)
        layers.append(body_layer)
        return item['body_image'].size
    else:
        front_layer, back_layer = organize_clothing(item)
        layers.append(front_layer)
        layers.append(back_layer)


def process_object(object_data):
    basic = object_data['basic']
    items_response = {'layers': []}

    if basic['single_overall'] == "overall":
        item_results = [process_item(item, basic) for item in object_data['items']]
        layers = []
        futures = []
        body_size = None
        for item in item_results:
            futures = [process_layer(item, layers)]
        for future in futures:
            if future is not None:
                body_size = future
        layers = sorted(layers, key=lambda s: s.get("priority", float('inf')))

        layers, new_size = update_base_size_priority(layers, body_size)

        for lay in layers:
            items_response['layers'].append({
                'image_category': lay['name'],
                'position': lay['position'],
                'priority': lay.get("priority", None),
                'resize_scale': lay['resize_scale'] if "resize_scale" in lay.keys() else None,
                'image_size': lay['image'] if lay['image'] is None else lay['image'].size,
                'gradient_string': lay['gradient_string'] if 'gradient_string' in lay.keys() else "",
                'mask_url': lay['mask_url'],
                'image_url': lay['image_url'] if 'image_url' in lay.keys() else None,
                'pattern_image_url': lay['pattern_image_url'] if 'pattern_image_url' in lay.keys() else None,

                # 'image': lay['image'],
                # 'mask_image': lay['mask_image'],
            })
        items_response['synthesis_url'] = synthesis(layers, new_size, basic)
    else:
        item_results = process_item(object_data['items'][0], basic)
        items_response['layers'].append({
            'image_category': f"{item_results['name']}_front",
            'image_size': item_results['back_image'].size if item_results['back_image'] else None,
            'position': None,
            'priority': 0,
            'image_url': item_results['front_image_url'],
            'mask_url': item_results['mask_url'],
            "gradient_string": item_results['gradient_string'] if 'gradient_string' in item_results.keys() else "",
            'pattern_image_url': item_results['pattern_image_url'] if 'pattern_image_url' in item_results.keys() else None,

        })
        items_response['layers'].append({
            'image_category': f"{item_results['name']}_back",
            'image_size': item_results['front_image'].size if item_results['front_image'] else None,
            'position': None,
            'priority': 0,
            'image_url': item_results['back_image_url'],
            'mask_url': item_results['mask_url'],
            "gradient_string": item_results['gradient_string'] if 'gradient_string' in item_results.keys() else "",
            'pattern_image_url': item_results['pattern_image_url'] if 'pattern_image_url' in item_results.keys() else None,

        })
        items_response['synthesis_url'] = synthesis_single(item_results['front_image'], item_results['back_image'])
    return items_response


def update_base_size_priority(layers, size):
    # 计算透明背景图片的宽度
    min_x = min(info['position'][1] for info in layers)
    x_list = []
    for info in layers:
        if info['image'] is not None:
            x_list.append(info['position'][1] + info['image'].width)
    max_x = max(x_list)
    new_width = max_x - min_x
    new_height = 700
    # 更新坐标
    for info in layers:
        info['adaptive_position'] = (info['position'][0], info['position'][1] - min_x)
    return layers, (new_width, new_height)


def run():
    object = {"objects": [{"basic": {"body_point_test": {"waistband_right": [199, 239], "hand_point_right": [220, 308], "waistband_left": [113, 239], "hand_point_left": [92, 310], "shoulder_left": [99, 111], "shoulder_right": [214, 111]}, "layer_order": False, "scale_bag": 0.7, "scale_earrings": 0.16, "self_template": True, "single_overall": "single", "switch_category": "Outwear"}, "items": [
        {"color": "189 112 112", "icon": "none", "image_id": 116441, "offset": [1, 1], "path": "aida-sys-image/images/female/outwear/outwear_p3139.jpg", "print": {"element": {"element_angle_list": [], "element_path_list": [], "element_scale_list": [], "location": []}, "overall": {"location": [[0.0, 0.0]], "print_angle_list": [0.0, 0.0], "print_path_list": [], "print_scale_list": [0.0, 0.0]}, "single": {"location": [], "print_angle_list": [], "print_path_list": [], "print_scale_list": []}},
         "resize_scale": [1.0, 1.0], "type": "Outwear"}]}, {"basic": {"body_point_test": {"waistband_right": [199, 239], "hand_point_right": [220, 308], "waistband_left": [113, 239], "hand_point_left": [92, 310], "shoulder_left": [99, 111], "shoulder_right": [214, 111]}, "layer_order": False, "scale_bag": 0.7, "scale_earrings": 0.16, "self_template": True, "single_overall": "single", "switch_category": "Outwear"}, "items": [
        {"color": "189 112 112", "icon": "none", "image_id": 81518, "offset": [1, 1], "path": "aida-sys-image/images/female/outwear/0628000071.jpg", "print": {"element": {"element_angle_list": [], "element_path_list": [], "element_scale_list": [], "location": []}, "overall": {"location": [[0.0, 0.0]], "print_angle_list": [0.0, 0.0], "print_path_list": [], "print_scale_list": [0.0, 0.0]}, "single": {"location": [], "print_angle_list": [], "print_path_list": [], "print_scale_list": []}},
         "resize_scale": [1.0, 1.0], "type": "Outwear"}]}, {"basic": {"body_point_test": {"waistband_right": [199, 239], "hand_point_right": [220, 308], "waistband_left": [113, 239], "hand_point_left": [92, 310], "shoulder_left": [99, 111], "shoulder_right": [214, 111]}, "layer_order": False, "scale_bag": 0.7, "scale_earrings": 0.16, "self_template": True, "single_overall": "single", "switch_category": "Outwear"}, "items": [
        {"color": "189 112 112", "icon": "none", "image_id": 65687, "offset": [1, 1], "path": "aida-sys-image/images/female/outwear/outwear_746.jpg", "print": {"element": {"element_angle_list": [], "element_path_list": [], "element_scale_list": [], "location": []}, "overall": {"location": [[0.0, 0.0]], "print_angle_list": [0.0, 0.0], "print_path_list": [], "print_scale_list": [0.0, 0.0]}, "single": {"location": [], "print_angle_list": [], "print_path_list": [], "print_scale_list": []}},
         "resize_scale": [1.0, 1.0], "type": "Outwear"}]}, {"basic": {"body_point_test": {"waistband_right": [199, 239], "hand_point_right": [220, 308], "waistband_left": [113, 239], "hand_point_left": [92, 310], "shoulder_left": [99, 111], "shoulder_right": [214, 111]}, "layer_order": False, "scale_bag": 0.7, "scale_earrings": 0.16, "self_template": True, "single_overall": "single", "switch_category": "Outwear"}, "items": [
        {"color": "189 112 112", "icon": "none", "image_id": 90051, "offset": [1, 1], "path": "aida-sys-image/images/female/outwear/0628000864.jpg", "print": {"element": {"element_angle_list": [], "element_path_list": [], "element_scale_list": [], "location": []}, "overall": {"location": [[0.0, 0.0]], "print_angle_list": [0.0, 0.0], "print_path_list": [], "print_scale_list": [0.0, 0.0]}, "single": {"location": [], "print_angle_list": [], "print_path_list": [], "print_scale_list": []}},
         "resize_scale": [1.0, 1.0], "type": "Outwear"}]}, {"basic": {"body_point_test": {"waistband_right": [199, 239], "hand_point_right": [220, 308], "waistband_left": [113, 239], "hand_point_left": [92, 310], "shoulder_left": [99, 111], "shoulder_right": [214, 111]}, "layer_order": False, "scale_bag": 0.7, "scale_earrings": 0.16, "self_template": True, "single_overall": "single", "switch_category": "Outwear"}, "items": [
        {"color": "189 112 112", "icon": "none", "image_id": 67420, "offset": [1, 1], "path": "aida-sys-image/images/female/outwear/0825001648.jpg", "print": {"element": {"element_angle_list": [], "element_path_list": [], "element_scale_list": [], "location": []}, "overall": {"location": [[0.0, 0.0]], "print_angle_list": [0.0, 0.0], "print_path_list": [], "print_scale_list": [0.0, 0.0]}, "single": {"location": [], "print_angle_list": [], "print_path_list": [], "print_scale_list": []}},
         "resize_scale": [1.0, 1.0], "type": "Outwear"}]}, {"basic": {"body_point_test": {"waistband_right": [199, 239], "hand_point_right": [220, 308], "waistband_left": [113, 239], "hand_point_left": [92, 310], "shoulder_left": [99, 111], "shoulder_right": [214, 111]}, "layer_order": False, "scale_bag": 0.7, "scale_earrings": 0.16, "self_template": True, "single_overall": "single", "switch_category": "Outwear"}, "items": [
        {"color": "189 112 112", "icon": "none", "image_id": 90354, "offset": [1, 1], "path": "aida-sys-image/images/female/outwear/0628001300.jpg", "print": {"element": {"element_angle_list": [], "element_path_list": [], "element_scale_list": [], "location": []}, "overall": {"location": [[0.0, 0.0]], "print_angle_list": [0.0, 0.0], "print_path_list": [], "print_scale_list": [0.0, 0.0]}, "single": {"location": [], "print_angle_list": [], "print_path_list": [], "print_scale_list": []}},
         "resize_scale": [1.0, 1.0], "type": "Outwear"}]}, {"basic": {"body_point_test": {"waistband_right": [199, 239], "hand_point_right": [220, 308], "waistband_left": [113, 239], "hand_point_left": [92, 310], "shoulder_left": [99, 111], "shoulder_right": [214, 111]}, "layer_order": False, "scale_bag": 0.7, "scale_earrings": 0.16, "self_template": True, "single_overall": "single", "switch_category": "Outwear"}, "items": [
        {"color": "189 112 112", "icon": "none", "image_id": 67420, "offset": [1, 1], "path": "aida-sys-image/images/female/outwear/0825001648.jpg", "print": {"element": {"element_angle_list": [], "element_path_list": [], "element_scale_list": [], "location": []}, "overall": {"location": [[0.0, 0.0]], "print_angle_list": [0.0, 0.0], "print_path_list": [], "print_scale_list": [0.0, 0.0]}, "single": {"location": [], "print_angle_list": [], "print_path_list": [], "print_scale_list": []}},
         "resize_scale": [1.0, 1.0], "type": "Outwear"}]}, {"basic": {"body_point_test": {"waistband_right": [199, 239], "hand_point_right": [220, 308], "waistband_left": [113, 239], "hand_point_left": [92, 310], "shoulder_left": [99, 111], "shoulder_right": [214, 111]}, "layer_order": False, "scale_bag": 0.7, "scale_earrings": 0.16, "self_template": True, "single_overall": "single", "switch_category": "Outwear"}, "items": [
        {"color": "189 112 112", "icon": "none", "image_id": 101477, "offset": [1, 1], "path": "aida-sys-image/images/female/outwear/903000063.jpg", "print": {"element": {"element_angle_list": [], "element_path_list": [], "element_scale_list": [], "location": []}, "overall": {"location": [[0.0, 0.0]], "print_angle_list": [0.0, 0.0], "print_path_list": [], "print_scale_list": [0.0, 0.0]}, "single": {"location": [], "print_angle_list": [], "print_path_list": [], "print_scale_list": []}},
         "resize_scale": [1.0, 1.0], "type": "Outwear"}]}], "process_id": "3615898424593104"}

    object_result = {}
    with ThreadPoolExecutor() as executor:
        results = list(executor.map(process_object, object['objects']))
        for i, result in enumerate(results):
            object_result[i] = result

    pprint(object_result)


if __name__ == '__main__':
    start_time = time.time()
    run()
    print(time.time() - start_time)