Spaces:
Runtime error
Runtime error
| import numpy as np | |
| import os | |
| import time | |
| import sys | |
| import torch | |
| import zipfile | |
| 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 | |
| # Use GPU if available | |
| if torch.cuda.is_available(): | |
| device = torch.device("cuda") | |
| else: | |
| device = torch.device("cpu") | |
| print(f"####Using {device}.#####") | |
| # 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/") | |
| os.system("mkdir ./saved_models/u2net") | |
| os.system("mkdir ./saved_models/u2netp") | |
| os.system("mkdir ./pose") | |
| os.system("mkdir ./checkpoints") | |
| # Get pose model | |
| if not os.path.exists("./pose/pose_deploy_linevec.prototxt"): | |
| os.system("wget -O ./pose/pose_deploy_linevec.prototxt https://github.com/hasibzunair/fifa-demo/releases/download/v1.0/pose_deploy_linevec.prototxt") | |
| if not os.path.exists("./pose/pose_iter_440000.caffemodel"): | |
| 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 | |
| if not os.path.exists("lip_final.pth"): | |
| os.system("wget https://github.com/hasibzunair/fifa-demo/releases/download/v1.0/lip_final.pth") | |
| # Get U-2-Net weights | |
| if not os.path.exists("saved_models/u2netp/u2netp.pth"): | |
| os.system("wget -P saved_models/u2netp/ https://github.com/hasibzunair/fifa-demo/releases/download/v1.0/u2netp.pth") | |
| if not os.path.exists("saved_models/u2net/u2net.pth"): | |
| os.system("wget -P saved_models/u2net/ https://github.com/hasibzunair/fifa-demo/releases/download/v1.0/u2net.pth") | |
| # Get model checkpoints | |
| if not os.path.exists("./checkpoints/decavtonfifapretrain/"): | |
| os.system("wget -O ./checkpoints/decavtonfifapretrain.zip https://github.com/hasibzunair/vton-demo/releases/download/v1.0/decavtonfifapretrain.zip") | |
| with zipfile.ZipFile('./checkpoints/decavtonfifapretrain.zip', 'r') as zip_ref: | |
| zip_ref.extractall('./checkpoints/') | |
| print("########################Setup done!########################") | |
| # Load U-2-Net model | |
| print(f"####Using {device}.#####") | |
| u2net = u2net_load.model(model_name = 'u2netp') | |
| def composite_background(img_mask, person_image_path, tryon_image_path): | |
| """Put background back on the person image after tryon.""" | |
| person = np.array(Image.open(person_image_path)) | |
| # tryon image | |
| tryon = np.array(Image.open(tryon_image_path)) | |
| # persom image mask from rembg | |
| p_mask = np.array(img_mask) | |
| # make binary mask | |
| p_mask = np.where(p_mask>0, 1, 0) | |
| # invert mask | |
| p_mask_inv = np.logical_not(p_mask) | |
| # make bg without person | |
| background = person * np.stack((p_mask_inv, p_mask_inv, p_mask_inv), axis=2) | |
| # make tryon image without background | |
| tryon_nobg = tryon * np.stack((p_mask, p_mask, p_mask), axis=2) | |
| tryon_nobg = tryon_nobg.astype("uint8") | |
| # composite | |
| tryon_with_bg = np.add(tryon_nobg, background) | |
| tryon_with_bg_pil = Image.fromarray(np.uint8(tryon_with_bg)).convert('RGB') | |
| tryon_with_bg_pil.save("results/test/try-on/tryon_with_bg.png") | |
| # Main inference function | |
| def inference(clothing_image, person_image, retrieve_bg): | |
| """ | |
| Do try-on! | |
| """ | |
| remove_bg = "no" | |
| # 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") | |
| # Get binary from person image, used in def_composite_background | |
| img_mask = remove(img, alpha_matting=True, alpha_matting_erode_size=15, only_mask=True) | |
| 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("python 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 try-on with background added back on the person image | |
| if retrieve_bg == "yes": | |
| composite_background(img_mask, 'Data_preprocessing/test_img/person.png', | |
| 'results/test/try-on/person.png') | |
| return os.path.join("results/test/try-on/tryon_with_bg.png") | |
| # Return only try-on result | |
| else: | |
| return os.path.join("results/test/try-on/person.png") | |
| title = "Image based Virtual Try-On" | |
| description = "This is a demo for an image based virtual try-on system. 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. This demo currently uses a temporary GPU. You can always run the demo locally, of course on a machine with a GPU!" | |
| article = "<p style='text-align: center'><a href='will_be_added' 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" | |
| # todos | |
| # train model with background removed then add feature, also add remove_bg in inferene() | |
| # add gr.inputs.Radio(choices=["yes","no"], label="Remove background from the person image?", type='index') in inputs | |
| gr.Interface( | |
| inference, | |
| [gr.inputs.Image(type='file', label="Clothing Image"), | |
| gr.inputs.Image(type='file', label="Person Image"), | |
| gr.inputs.Radio(choices=["yes","no"], default="no", label="Retrieve original background from the person image?")], | |
| gr.outputs.Image(type="file", label="Predicted Output"), | |
| examples=[["./sample_images/1/cloth.jpg", "./sample_images/1/person.jpg", "yes"], | |
| ["./sample_images/2/cloth.jpg", "./sample_images/2/person.jpg", "no"]], | |
| title=title, | |
| description=description, | |
| article=article, | |
| allow_flagging=False, | |
| analytics_enabled=False, | |
| thumbnail=thumbnail, | |
| ).launch(debug=True, enable_queue=True) |