app.py

import gradio as gr
import asyncio
import os
import subprocess
import shutil
import logging
import json
from pathlib import Path
import yaml
import coloredlogs
import re

from custom_utils.setup_env import run_setup_script



logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
coloredlogs.install(level='INFO', fmt='%(levelname)s: %(message)s',
                    level_styles={'info': {'color': 'green'},
                                  'error':{'color':'red'}
                                  }
                    )

# === Загрузка конфига ===
with open("GarmentCode/gradio_config.yml", "r") as f:
    config = yaml.safe_load(f)

# === Вспомогательные функции ===

def organize_images_into_structure(source_folder, new_folder):
    os.makedirs(new_folder, exist_ok=True)
    for file_name in os.listdir(source_folder):
        source_file = os.path.join(source_folder, file_name)
        if os.path.isfile(source_file) and file_name.lower().endswith(('.png', '.jpg', '.jpeg')):
            folder_name = os.path.splitext(file_name)[0]
            subfolder_path = os.path.join(new_folder, folder_name, "outputs")
            os.makedirs(subfolder_path, exist_ok=True)
            destination_file = os.path.join(subfolder_path, file_name)
            shutil.copy(source_file, destination_file)


def update_body_model(selected_key):
    return config['body_model_files'].get(selected_key, config['defaults']['body_config'])


def get_design_choices():
    design_dir = config['assets_paths']['design_params'] # "assets/design_params/"
    design_choices = []
    for fname in os.listdir(design_dir):
        if fname.endswith(".yaml"):
            design_choices.append(os.path.splitext(fname)[0])
    return design_choices

def get_sim_props_choices():
    sim_props_dir=config['assets_paths']['sim_props']
    sim_props_choices_yaml, sim_props_choices_json=[],[]
    for fname in os.listdir(sim_props_dir):
        if fname.endswith(".yaml"):
            sim_props_choices_yaml.append(os.path.splitext(fname)[0])
        elif fname.endswith(".json"):
            sim_props_choices_json.append(os.path.splitext(fname)[0])
    return sim_props_choices_yaml, sim_props_choices_json

def extract_dataset_path(stdout: str) -> str:
    # Ищем строку с "Result saved to:"
    match = re.search(r"Result saved to: (.+)", stdout)
    if match:
        path = match.group(1).strip()
        # Убираем подкаталог default_body, если он есть в конце пути
        if path.endswith('/default_body') or path.endswith('\\default_body'):
            path = os.path.dirname(path)
        return path
    return None

def get_png_path(folder: str):
    basename = os.path.basename(folder)
    png_path = os.path.join(folder, f"{basename}_pattern.png")

    if os.path.isfile(png_path):
        logging.info(f"PNG FILE LOCATED AT {png_path}")
        return png_path
    return None


def zip_folder(folder_path: str) -> str:
    folder_path_obj = Path(folder_path)
    zip_path = str(folder_path_obj) + '.zip'
    shutil.make_archive(str(folder_path_obj), 'zip', folder_path)
    return zip_path

def show_button_if_both_uploaded(front_img, back_img):
    return gr.update(visible=front_img is not None and back_img is not None)


def update_submenu(clothing_type):
    if clothing_type == "Top":
        return (
            gr.Dropdown(visible=True, label="Front Collar",
                        choices=["CircleNeckHalf", "CurvyNeckHalf", "VNeckHalf", "SquareNeckHalf",
                                 "TrapezoidNeckHalf", "CircleArcNeckHalf", "Bezier2NeckHalf"],
                        interactive=True),
            gr.Dropdown(visible=True, label="Back Collar",
                        choices=["CircleNeckHalf", "CurvyNeckHalf", "VNeckHalf", "SquareNeckHalf",
                                 "TrapezoidNeckHalf", "CircleArcNeckHalf", "Bezier2NeckHalf"],
                        interactive=True)
        )
    elif clothing_type == "Down":
        return (
            gr.Dropdown(visible=True, label="Waist Style",
                        choices=["High Waist", "Low Waist", "Medium Waist"],
                        interactive=True),
            gr.Dropdown(visible=False, interactive=True)
        )
    elif clothing_type == "Full body":
        return (
            gr.Dropdown(visible=True, label="Overall Style",
                        choices=["Slim Fit", "Loose Fit", "Regular Fit"],
                        interactive=True),
            gr.Dropdown(visible=False, interactive=True)
        )
    else:
        return (
            gr.Dropdown(visible=False, interactive=True),
            gr.Dropdown(visible=False, interactive=True)
        )


# === Симуляция/сэмплирование ===

# TODO: recreate fitting. Is it strictly necessary?

async def run_dataset_fitting(design_file: str, size: int, name: str, replicate: bool) -> str:
    return

async def run_pattern_sampler(size: int, dataset_name: str, replicate: bool, pattern_config_name: str, body_file:str):
    pattern_config_path = f"assets/design_params/{pattern_config_name}.yaml"
    body_path = getattr(body_file, 'name', body_file) if body_file else "mean_female"

    cmd = [
        "python3", "GarmentCode/pattern_sampler.py",
        "-s", str(size),
        "--name", dataset_name,
        "--pattern-config", pattern_config_path,
        "--body", body_path
    ]
    if replicate:
        cmd.append("--replicate")

    proc = await asyncio.create_subprocess_exec(
        *cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE
    )
    stdout, stderr = await proc.communicate()
    logging.info(stdout.decode())
    log = stdout.decode()
    if proc.returncode != 0:
        logging.error(f"Ошибка при генерации датасета:{stderr.decode()}")
        return f"Ошибка при генерации датасета:\n{stderr.decode()}", gr.update(visible=False), None
    dataset_path = (extract_dataset_path(log))
    zip_folder_path = f"{zip_folder(dataset_path)}"
    return stdout.decode(), gr.update(visible=True), dataset_path, zip_folder_path


async def run_pattern_data_sim(dataset_path: str, clothing_spec: str, default_body: bool = False):
    dataset_path = dataset_path.replace("datasets/", "", 1)
    clothing_spec=clothing_spec+".yaml"

    cmd = [
        "python3", "GarmentCode/pattern_data_sim.py",
        "--data", dataset_path,
        "-c", clothing_spec
    ]
    # Добавляем флаг только если default_body=True
    if default_body:
        cmd.append("--default_body")

    proc = await asyncio.create_subprocess_exec(
        *cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE
    )
    stdout, stderr = await proc.communicate()
    if proc.returncode != 0:
        logging.error(f"Ошибка при симуляции: {stderr.decode()}")
        return f"Ошибка при симуляции:\n{stderr.decode()}"
    dataset_sim_path = "datasets_sim/" + dataset_path
    zip_file_path = f"{zip_folder(dataset_sim_path)}"
    return zip_file_path


async def simulation(pattern_spec_path:str, sim_props_path:str):
    cmd = [
        "python3", "GarmentCode/test_garment_sim.py",
        "-p", pattern_spec_path,
        "-s", sim_props_path,
    ]


async def create_pattern(
        body_config_file,  # gr.File или None
        design_config_file,  # gr.File или None
        body_to_use=None,  # str или None
        design_to_use=None  # str или None
):
    # Если файл не выбран - используем дефолт из конфига
    body_config_path = getattr(body_config_file, 'name', body_config_file) if body_config_file else config['defaults'][
        'body_config']
    design_config_path = getattr(design_config_file, 'name', design_config_file) if design_config_file else \
    config['defaults']['design_config']
    logging.info("CREATION OF PATTERN FILES...")
    cmd = [
        "python3", "test_garmentcode.py",
        "--body-config", body_config_path,
        "--design-config", design_config_path,
    ]
    if body_to_use:
        cmd += ["--body-to-use", body_to_use]
    if design_to_use:
        cmd += ["--design-to-use", design_to_use]

    proc = await asyncio.create_subprocess_exec(
        *cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE
    )
    stdout, stderr = await proc.communicate()
    if proc.returncode != 0:
        logging.error(f"Pattern creation failed:\n{stderr.decode()}")
        return f"Pattern creation failed:\n{stderr.decode()}"

    result_str = stdout.decode()
    folder = None
    for line in result_str.splitlines():
        if "saved to" in line:
            folder = line.split("saved to")[-1].strip()
            break
    if not folder or not os.path.isdir(folder):
        logging.error("Ошибка: не удалось найти папку с результатом.")
        return "Ошибка: не удалось найти папку с результатом.", None, None

    png_file = get_png_path(folder)
    if not png_file:
        logging.error("Ошибка: PNG-файл не найден.")
        return "Ошибка: PNG-файл не найден.", None, None

    zip_path = zip_folder(folder)
    logging.info("PATTERN CREATION ENDED SUCCESSFULLY")
    return result_str, zip_path, png_file


with gr.Blocks(title="3D Garment Generator", theme=gr.themes.Default(text_size="sm")) as demo:
    with gr.Tabs():
        with gr.Tab("Генератор одежды"):
            front_collar = gr.State()
            back_collar = gr.State()

            with gr.Row():
                with gr.Column(scale=1):
                    with gr.Accordion("Инференс",open=False):
                        front_img = gr.Image(label="Добавьте фотографию впереди")
                        back_img = gr.Image(label="Добавьте фотографию сзади")
                        create_pattern_btn = gr.Button("Создать паттерн", visible=False)

                    with gr.Accordion("1. Пресеты", open=False):
                        gr.Markdown("### Параметры генерации")
                        with gr.Accordion("Загрузить файлы", open=False):
                            body_config_file = gr.File(label="Body config (YAML)")
                            design_config_file = gr.File(label="Design config (YAML)")
                            system_config_file = gr.File(label="System config (JSON)")
                        body_to_use_dropdown = gr.Dropdown(
                            choices=list(config['body_model_files'].keys()),
                            label="Body key", value="mean_female")
                        design_to_use_dropdown = gr.Dropdown(
                            choices=["t-shirt"], label="Design key", value="t-shirt")
                        create_pattern_btn = gr.Button("Создать паттерн")
                        result_box = gr.Textbox(label="Результат")
                        image_result = gr.Image(label="Результат PNG")

                    with gr.Accordion("2. Параметры генерации", open=True):
                        clothing_type_dropdown = gr.Dropdown(
                            label="Тип одежды",
                            choices=["Top", "Down", "Full body"],
                            value="Top",
                            interactive=True
                        )
                        front_dropdown = gr.Dropdown(
                            visible=False,
                            label="Front Collar",
                            choices=[],
                            interactive=True
                        )
                        back_dropdown = gr.Dropdown(
                            visible=False,
                            label="Back Collar",
                            choices=[],
                            interactive=True
                        )
                        overall_length = gr.Slider(
                            1, 50, value=25, label="Длина одежды",
                            info="Выберите значение от 1 до 50",
                            interactive=True
                        )
                        clothing_type_dropdown.change(
                            fn=update_submenu,
                            inputs=clothing_type_dropdown,
                            outputs=[front_dropdown, back_dropdown]
                        )
                    with gr.Accordion("3. Дополнительные инструменты", open=False):
                        with gr.Accordion("Генерация датасета", open=False):
                            dataset_size = gr.Slider(1, 5000, value=1, label="Размер датасета",
                                                        info="Выберите значение от 1 до 5000", interactive=True)
                            replicate_checkbox = gr.Checkbox(label="Репликация", value=False)
                            dataset_name_box = gr.Textbox(label="Название датасета", max_lines=1)

                            pattern_config_choices = get_design_choices()
                            pattern_config_file = gr.Radio(
                                label="Config",
                                choices=pattern_config_choices,
                                value=pattern_config_choices[0] if pattern_config_choices else None
                            )
                            body_file = gr.File(label="Body config (YAML)")
                            pattern_sampler_run_button = gr.Button("Сгенерировать паттерны", size="sm")
                            pattern_result_box = gr.Textbox(label="Лог генерации")
                            # sim_props_gen_choices_yaml = gr.Radio(label="Properties config",
                            #                                 choices=get_sim_props_choices()[0],
                            #                                 interactive=True,
                            #                                 value=pattern_config_choices[
                            #                                     0] if pattern_config_choices else None
                            #
                            #                                 )
                            sim_props_gen_choices_json = gr.Radio(label="Clothing config",
                                                             choices=get_sim_props_choices()[0],
                                                             interactive=True,
                                                             value=pattern_config_choices[
                                                                 0] if pattern_config_choices else None

                                                             )
                            dataset_path_box = gr.Textbox(visible=False)
                            default_body_checkbox = gr.Checkbox(label="default body",value=False)
                            dataset_gen_button = gr.Button("Генерация датасета",
                                                                    variant="primary", visible=False)

                    ## Button callbacks ##
                with gr.Column(scale=3):
                    viewer = gr.Model3D(
                        config['body_model_files']['custom_female'],
                        display_mode="solid", label="3D Рендер",
                        height=500, container=True, camera_position=(90, None, None),
                        interactive=False, clear_color=(0.85, 0.85, 0.85, 1)
                    )
                    fit_button = gr.Button("Запустить симуляцию", variant="primary", size="sm")
                    with gr.Row():
                        save_render_btn = gr.Button("Сохранить рендер")
                        save_obj_btn = gr.Button("Сохранить OBJ")
                        clear_button = gr.Button("Очистить")
                    zip_file = gr.File(label="Download Zip File", file_count="multiple",
                                       file_types=[".zip"], height=50, interactive=False)

            create_pattern_btn.click(
                fn=create_pattern,
                inputs=[body_config_file, design_config_file, body_to_use_dropdown, design_to_use_dropdown],
                outputs=[result_box, zip_file, image_result]
            )
            body_to_use_dropdown.change(
                fn=update_body_model,
                inputs=body_to_use_dropdown,
                outputs=viewer
            )
            pattern_sampler_run_button.click(
                fn=run_pattern_sampler,
                inputs=[dataset_size, dataset_name_box, replicate_checkbox, pattern_config_file,
                        body_file],
                outputs=[pattern_result_box, dataset_gen_button,  dataset_path_box, zip_file]
            )
            dataset_gen_button.click(
                fn=run_pattern_data_sim,
                inputs=[dataset_path_box, sim_props_gen_choices_json, default_body_checkbox],
                outputs=[zip_file],
            )
            front_img.change(
                fn=show_button_if_both_uploaded,
                inputs=[front_img, back_img],
                outputs=create_pattern_btn
            )
            back_img.change(
                fn=show_button_if_both_uploaded,
                inputs=[front_img, back_img],
                outputs=create_pattern_btn
            )

        with gr.Tab("SewFormer"):
            gr.Markdown("Photo to pattern")
        with gr.Tab("Image to Pattern"):
            gr.Markdown("JPG to PBR")

if __name__ == "__main__":
    run_setup_script()
    demo.launch(server_port=7860)