import os
import subprocess
from typing import Dict
import json
import numpy as np
import wget
import gradio as gr
from helpers.processor import TextureProcessor

subprocess.call(['pip', 'install', 'git+https://github.com/facebookresearch/detectron2@main#subdirectory=projects/DensePose'])

def image_processing(person_img: np.ndarray, model_img: np.ndarray) -> np.ndarray:
    return texture_processor.extract(person_img, model_img)

def load_model(current_path: str, config: Dict) -> None:
    data_path = os.path.join(current_path, 'data')
    if not os.path.isdir(data_path):
        os.mkdir(data_path)
        for filename, url in config.items():
            wget.download(url, os.path.join(data_path, filename))

with open("config.json", "r") as f:
    config = json.load(f)

current_path = os.getcwd()
load_model(current_path, config)
densepose_config = os.path.join(current_path, 'data', 'config.yaml')
densepose_weights = os.path.join(current_path, 'data', 'weights.pkl')

texture_processor = TextureProcessor(densepose_config, densepose_weights)

title = '<h1 style="text-align:center">JustClothify</h1>'

with gr.Blocks(theme='soft', title='JustClothify') as blocks:
    gr.HTML(title)
    gr.Markdown('Upload an image of a person and an image of a model with clothes, the system will generate an image of a person wearing these clothes.')
    with gr.Row():
        person_image = gr.inputs.Image(label='Person Image', type='numpy')
        model_image = gr.inputs.Image(label='Model Image (with clothes)', type='numpy')
    process_button = gr.Button('Process')
    outputs = gr.outputs.Image(label='Result Image', type='numpy')

    process_button.click(fn=image_processing, inputs=[person_image, model_image], outputs=outputs)

blocks.launch()