import gradio as gr
import os
import wget
import subprocess
subprocess.call(['pip', 'install', 'git+https://github.com/facebookresearch/detectron2@main#subdirectory=projects/DensePose'])
from helpers.processor import TextureProcessor

def image_processing(person_img, model_img):
    return texture_processor.extract(person_img, model_img)

def load_model(current_path):
    data_path = os.path.join(current_path, 'data')
    if not os.path.isdir(data_path):
        os.mkdir(data_path)
        url = 'https://raw.githubusercontent.com/facebookresearch/detectron2/main/projects/DensePose/configs/densepose_rcnn_R_50_FPN_WC1M_s1x.yaml'
        wget.download(url, os.path.join(data_path, 'config.yaml'))
        url = 'https://dl.fbaipublicfiles.com/densepose/densepose_rcnn_R_50_FPN_WC1M_s1x/217144516/model_final_48a9d9.pkl'
        wget.download(url, os.path.join(data_path, 'weights.pkl'))
        url = 'https://raw.githubusercontent.com/facebookresearch/detectron2/main/projects/DensePose/configs/Base-DensePose-RCNN-FPN.yaml'
        wget.download(url, os.path.join(data_path, 'Base-DensePose-RCNN-FPN.yaml'))

current_path = os.getcwd()
load_model(current_path)
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)

inputs = [
    gr.inputs.Image(label='Person Image', type='numpy'),
    gr.inputs.Image(label='Model Image (with clothes)', type='numpy')
]

outputs = gr.outputs.Image(label='Result Image', type='numpy')

title = 'JustClothify'
description = '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.'

gr.Interface(
    fn=image_processing, 
    inputs=inputs, 
    outputs=outputs, 
    theme='soft',
    title=title,
    description=description).launch()