JustClothify / app.py
Brasd99's picture
Bug fix
6819f9c
raw
history blame
1.76 kB
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()