import spaces import gradio as gr import apply_net import os import sys import cv2 sys.path.append('./') import numpy as np import argparse import torch import torchvision import pytorch_lightning from torch import autocast from torchvision import transforms from pytorch_lightning import seed_everything from einops import rearrange from functools import partial from omegaconf import OmegaConf from PIL import Image from typing import List import matplotlib.pyplot as plt from torchvision.transforms.functional import to_pil_image from utils_mask import get_mask_location from preprocess.humanparsing.run_parsing import Parsing from preprocess.openpose.run_openpose import OpenPose from ldm.util import instantiate_from_config, get_obj_from_str from ldm.models.diffusion.ddim import DDIMSampler from detectron2.data.detection_utils import convert_PIL_to_numpy,_apply_exif_orientation if __name__ == "__main__": parser = argparse.ArgumentParser(description="Script for demo model") parser.add_argument("-b", "--base", type=str, default=r"configs/test_vitonhd.yaml") parser.add_argument("-c", "--ckpt", type=str, default=r"ckpt/hitonhd.ckpt") parser.add_argument("-s", "--seed", type=str, default=42) parser.add_argument("-d", "--ddim", type=str, default=16) opt = parser.parse_args() seed_everything(opt.seed) config = OmegaConf.load(f"{opt.base}") device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") model = instantiate_from_config(config.model) model.load_state_dict(torch.hub.load_state_dict_from_url("https://huggingface.co/basso4/hitonhd/resolve/main/hitonhd.ckpt")["state_dict"], strict=False) model.cuda() model.eval() model = model.to(device) sampler = DDIMSampler(model) # model = instantiate_from_config(config.model) # model.load_state_dict(torch.load(opt.ckpt, map_location="cpu")["state_dict"], strict=False) # model.cuda() # model.eval() # device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") # model = model.to(device) # sampler = DDIMSampler(model) precision_scope = autocast @spaces.GPU def start_tryon(dict,garm_img): #load human image human_img = dict['background'].convert("RGB").resize((768,1024)) #mask tensor_transfrom = transforms.Compose( [ transforms.ToTensor(), transforms.Normalize([0.5], [0.5]), ] ) parsing_model = Parsing(0) openose_model = OpenPose(0) openose_model.preprocessor.body_estimation.model.to(device) keypoints = openose_model(human_img.resize((384,512))) model_parse, _ = parsing_model(human_img.resize((384,512))) mask, mask_gray = get_mask_location('hd', "upper_body", model_parse, keypoints) mask_cv = mask mask = mask.resize((768, 1024)) mask_gray = (1-transforms.ToTensor()(mask)) * tensor_transfrom(human_img) mask_gray = to_pil_image((mask_gray+1.0)/2.0) #densepose human_img_arg = _apply_exif_orientation(human_img.resize((384,512))) human_img_arg = convert_PIL_to_numpy(human_img_arg, format="BGR") args = apply_net.create_argument_parser().parse_args(('show', './configs/configs_densepose/densepose_rcnn_R_50_FPN_s1x.yaml', 'https://dl.fbaipublicfiles.com/densepose/densepose_rcnn_R_50_FPN_s1x/165712039/model_final_162be9.pkl', 'dp_segm', '-v', '--opts', 'MODEL.DEVICE', 'cuda')) # verbosity = getattr(args, "verbosity", None) pose_img = args.func(args,human_img_arg) pose_img = pose_img[:,:,::-1] pose_img = Image.fromarray(pose_img).resize((768,1024)) #preprocessing image human_img = human_img.convert("RGB").resize((512, 512)) human_img = torchvision.transforms.ToTensor()(human_img) garm_img = garm_img.convert("RGB").resize((224, 224)) garm_img = torchvision.transforms.ToTensor()(garm_img) mask = mask.convert("L").resize((512,512)) mask = torchvision.transforms.ToTensor()(mask) mask = 1-mask pose_img = pose_img.convert("RGB").resize((512, 512)) pose_img = torchvision.transforms.ToTensor()(pose_img) #Normalize human_img = torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))(human_img) garm_img = torchvision.transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))(garm_img) pose_img = torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))(pose_img) #create inpaint & hint inpaint = human_img * mask hint = torchvision.transforms.Resize((512, 512))(garm_img) hint = torch.cat((hint, pose_img), dim=0) # {"human_img": human_img, # [3, 512, 512] # "inpaint_image": inpaint, # [3, 512, 512] # "inpaint_mask": mask, # [1, 512, 512] # "garm_img": garm_img, # [3, 224, 224] # "hint": hint, # [6, 512, 512] # } with torch.no_grad(): with precision_scope("cuda"): #loading data inpaint = inpaint.unsqueeze(0).to(torch.float16).to(device) reference = garm_img.unsqueeze(0).to(torch.float16).to(device) mask = mask.unsqueeze(0).to(torch.float16).to(device) hint = hint.unsqueeze(0).to(torch.float16).to(device) truth = human_img.unsqueeze(0).to(torch.float16).to(device) #data preprocessing encoder_posterior_inpaint = model.first_stage_model.encode(inpaint) z_inpaint = model.scale_factor * (encoder_posterior_inpaint.sample()).detach() mask_resize = torchvision.transforms.Resize([z_inpaint.shape[-2],z_inpaint.shape[-1]])(mask) test_model_kwargs = {} test_model_kwargs['inpaint_image'] = z_inpaint test_model_kwargs['inpaint_mask'] = mask_resize shape = (model.channels, model.image_size, model.image_size) #predict samples, _ = sampler.sample(S=opt.ddim, batch_size=1, shape=shape, pose=hint, conditioning=reference, verbose=False, eta=0, test_model_kwargs=test_model_kwargs) samples = 1. / model.scale_factor * samples x_samples = model.first_stage_model.decode(samples[:,:4,:,:]) x_samples_ddim = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0) x_samples_ddim = x_samples_ddim.cpu().permute(0, 2, 3, 1).numpy() x_checked_image=x_samples_ddim x_checked_image_torch = torch.from_numpy(x_checked_image).permute(0, 3, 1, 2) x_checked_image_torch = torch.nn.functional.interpolate(x_checked_image_torch.float(), size=[512,384]) #apply seamlessClone technique here #img_base dict = dict['background'].convert("RGB").resize((384, 512)) dict = np.array(dict) dict = cv2.cvtColor(dict, cv2.COLOR_RGB2BGR) #img_output img_cv = rearrange(x_checked_image_torch[0], 'c h w -> h w c').cpu().numpy() img_cv = (img_cv * 255).astype(np.uint8) img_cv = cv2.cvtColor(img_cv, cv2.COLOR_RGB2BGR) #mask mask_cv = mask_cv.convert("L").resize((384,512)) mask_cv = np.array(mask_cv) mask_cv = 255-mask_cv img_C = cv2.seamlessClone(dict, img_cv, mask_cv, (192,256), cv2.NORMAL_CLONE) img_C = cv2.cvtColor(img_C, cv2.COLOR_BGR2RGB) return img_C, mask_gray example_path = os.path.join(os.path.dirname(__file__), 'example') garm_list = os.listdir(os.path.join(example_path,"cloth")) garm_list_path = [os.path.join(example_path,"cloth",garm) for garm in garm_list] human_list = os.listdir(os.path.join(example_path,"human")) human_list_path = [os.path.join(example_path,"human",human) for human in human_list] human_ex_list = [] for ex_human in human_list_path: ex_dict= {} ex_dict['background'] = ex_human ex_dict['layers'] = None ex_dict['composite'] = None human_ex_list.append(ex_dict) ##default human image_blocks = gr.Blocks().queue() with image_blocks as demo: gr.Markdown("## FPT_VTON 👕👔👚") gr.Markdown("Virtual Try-on with your image and garment image") with gr.Row(): with gr.Column(): imgs = gr.ImageEditor(sources='upload', type="pil", label='Human Picture or use Examples below', interactive=True) example = gr.Examples( inputs=imgs, examples_per_page=10, examples=human_list_path ) with gr.Column(): garm_img = gr.Image(label="Garment", sources='upload', type="pil") example = gr.Examples( inputs=garm_img, examples_per_page=8, examples=garm_list_path ) with gr.Column(): image_out_c = gr.Image(label="Output", elem_id="output-img",show_download_button=True) try_button = gr.Button(value="Try-on") # with gr.Column(): # image_out_c = gr.Image(label="Output", elem_id="output-img",show_download_button=False) with gr.Column(): masked_img = gr.Image(label="Masked image output", elem_id="masked_img", show_download_button=True) try_button.click(fn=start_tryon, inputs=[imgs,garm_img], outputs=[image_out_c,masked_img], api_name='tryon') image_blocks.launch()