Spaces:
Runtime error
Runtime error
| import numpy as np | |
| import os | |
| import time | |
| import sys | |
| import gradio as gr | |
| import u2net_load | |
| import u2net_run | |
| from rembg import remove | |
| from PIL import Image, ImageOps | |
| from predict_pose import generate_pose_keypoints | |
| # Make directories | |
| os.system("mkdir ./Data_preprocessing") | |
| os.system("mkdir ./Data_preprocessing/test_color") | |
| os.system("mkdir ./Data_preprocessing/test_colormask") | |
| os.system("mkdir ./Data_preprocessing/test_edge") | |
| os.system("mkdir ./Data_preprocessing/test_img") | |
| os.system("mkdir ./Data_preprocessing/test_label") | |
| os.system("mkdir ./Data_preprocessing/test_mask") | |
| os.system("mkdir ./Data_preprocessing/test_pose") | |
| os.system("mkdir ./inputs") | |
| os.system("mkdir ./inputs/img") | |
| os.system("mkdir ./inputs/cloth") | |
| os.system("mkdir ./saved_models/u2net") | |
| os.system("mkdir ./saved_models/u2netp") | |
| os.system("mkdir ./pose") | |
| os.system("mkdir ./checkpoints") | |
| # Get pose model | |
| os.system("wget -O ./pose/pose_deploy_linevec.prototxt https://github.com/hasibzunair/fifa-demo/releases/download/v1.0/pose_deploy_linevec.prototxt") | |
| os.system("wget -O ./pose/pose_iter_440000.caffemodel https://github.com/hasibzunair/fifa-demo/releases/download/v1.0/pose_iter_440000.caffemodel") | |
| # For segmentation mask generation | |
| os.system("wget https://github.com/hasibzunair/fifa-demo/releases/download/v1.0/lip_final.pth") | |
| # Get U-2-Net weights | |
| os.system("wget -P saved_models/u2netp/ https://github.com/hasibzunair/fifa-demo/releases/download/v1.0/u2netp.pth") | |
| os.system("wget -P saved_models/u2net/ https://github.com/hasibzunair/fifa-demo/releases/download/v1.0/u2net.pth") | |
| # Get model checkpoints | |
| os.system("wget -O ./checkpoints/decavtonfifapretrain.zip https://github.com/hasibzunair/vton-demo/releases/download/v1.0/decavtonfifapretrain.zip") | |
| os.system("unzip ./checkpoints/decavtonfifapretrain.zip -d ./checkpoints/") | |
| print("########################Setup done!########################") | |
| # Load U-2-Net model | |
| u2net = u2net_load.model(model_name = 'u2netp') | |
| # Main inference function | |
| def inference(clothing_image, person_image, remove_bg): | |
| """ | |
| Do try-on! | |
| """ | |
| # Read cloth and person images | |
| cloth = Image.open(clothing_image) # cloth | |
| person = Image.open(person_image) # person | |
| # Save cloth and person images in "input" folder | |
| cloth.save(os.path.join("inputs/cloth/cloth.png")) | |
| person.save(os.path.join("inputs/img/person.png")) | |
| ############## Clothing image pre-processing | |
| cloth_name = 'cloth.png' | |
| cloth_path = os.path.join('inputs/cloth', sorted(os.listdir('inputs/cloth'))[0]) | |
| cloth = Image.open(cloth_path) | |
| # Resize cloth image | |
| cloth = ImageOps.fit(cloth, (192, 256), Image.BICUBIC).convert("RGB") | |
| # Save resized cloth image | |
| cloth.save(os.path.join('Data_preprocessing/test_color', cloth_name)) | |
| # 1. Get binary mask for clothing image | |
| u2net_run.infer(u2net, 'Data_preprocessing/test_color', 'Data_preprocessing/test_edge') | |
| ############## Person image pre-processing | |
| start_time = time.time() | |
| # Person image | |
| img_name = 'person.png' | |
| img_path = os.path.join('inputs/img', sorted(os.listdir('inputs/img'))[0]) | |
| img = Image.open(img_path) | |
| if remove_bg == "yes": | |
| # Remove background | |
| img = remove(img, alpha_matting=True, alpha_matting_erode_size=15) | |
| print("Removing background from person image..") | |
| img = ImageOps.fit(img, (192, 256), Image.BICUBIC).convert("RGB") | |
| img_path = os.path.join('Data_preprocessing/test_img', img_name) | |
| img.save(img_path) | |
| resize_time = time.time() | |
| print('Resized image in {}s'.format(resize_time-start_time)) | |
| # 2. Get parsed person image (test_label), uses person image | |
| os.system("python3 Self-Correction-Human-Parsing-for-ACGPN/simple_extractor.py --dataset 'lip' --model-restore 'lip_final.pth' --input-dir 'Data_preprocessing/test_img' --output-dir 'Data_preprocessing/test_label'") | |
| parse_time = time.time() | |
| print('Parsing generated in {}s'.format(parse_time-resize_time)) | |
| # 3. Get pose map from person image | |
| pose_path = os.path.join('Data_preprocessing/test_pose', img_name.replace('.png', '_keypoints.json')) | |
| generate_pose_keypoints(img_path, pose_path) | |
| pose_time = time.time() | |
| print('Pose map generated in {}s'.format(pose_time-parse_time)) | |
| os.system("rm -rf Data_preprocessing/test_pairs.txt") | |
| # Format: person, cloth image | |
| with open('Data_preprocessing/test_pairs.txt','w') as f: | |
| f.write('person.png cloth.png') | |
| # Do try-on | |
| os.system("python test.py --name decavtonfifapretrain") | |
| tryon_image = Image.open("results/test/try-on/person.png") | |
| print("Size of image is: ", tryon_image.size) | |
| return os.path.join("results/test/try-on/person.png") | |
| title = "Virtual Dressing Room" | |
| description = "This is a demo for image based virtual try-on. It generates a synthetic image of a person wearing a target clothing item. To use it, simply upload your clothing item and person images, or click one of the examples to load them. Read more at the links below." | |
| article = "<p style='text-align: center'><a href='https://arxiv.org/abs/1512.03385' target='_blank'>Fill in Fabrics: Body-Aware Self-Supervised Inpainting for Image-Based Virtual Try-On (Under Review!)</a> | <a href='https://github.com/dktunited/fifa_demo' target='_blank'>Github</a></p>" | |
| thumbnail = None # "./pathtothumbnail.png" | |
| gr.Interface( | |
| inference, | |
| [gr.inputs.Image(type='filepath', label="Clothing Image"), | |
| gr.inputs.Image(type='filepath', label="Person Image"), | |
| gr.inputs.Radio(choices=["yes","no"], default="no", label="Remove background from person image?")], | |
| gr.outputs.Image(type="filepath", label="Predicted Output"), | |
| examples=[["./examples/1/cloth.jpg", "./examples/1/person.jpg"], | |
| ["./examples/2/cloth.jpg", "./examples/2/person.jpg"]], | |
| title=title, | |
| description=description, | |
| article=article, | |
| allow_flagging=False, | |
| analytics_enabled=False, | |
| thumbnail=thumbnail, | |
| ).launch(debug=True, enable_queue=True) |