from multiprocessing import set_start_method #set_start_method("fork") import sys #sys.path.insert(0, "../HR-VITON-main") from test_generator import * import re import inspect from dataclasses import dataclass, field from tqdm import tqdm import pandas as pd import os import torch import pandas as pd import gradio as gr import streamlit as st from io import BytesIO #### pip install streamlit-image-select from streamlit_image_select import image_select demo_image_dir = "demo_images_dir" assert os.path.exists(demo_image_dir) demo_images = list(map(lambda y: os.path.join(demo_image_dir, y) ,filter(lambda x: x.endswith(".png") or x.endswith(".jpeg") or x.endswith(".jpg") ,os.listdir(demo_image_dir)))) assert demo_images #https://github.com/jrieke/streamlit-image-select/issues/10 #.image-box { # border: 1px solid rgba(49, 51, 63, 0.2); # border-radius: 0.25rem; # padding: calc(0.25rem + 1px); # height: 10rem; # min-width: 10rem; #} demo_images = list(map(lambda x: x.resize((256, 256)), map(Image.open, demo_images))) @dataclass class OPT: #### ConditionGenerator out_layer = None warp_feature = None #### SPADEGenerator semantic_nc = None fine_height = None fine_width = None ngf = None num_upsampling_layers = None norm_G = None gen_semantic_nc = None #### weight load tocg_checkpoint = None gen_checkpoint = None cuda = False data_list = None datamode = None dataroot = None batch_size = None shuffle = False workers = None clothmask_composition = None occlusion = False datasetting = None opt = OPT() opt.out_layer = "relu" opt.warp_feature = "T1" input1_nc = 4 # cloth + cloth-mask nc = 13 input2_nc = nc + 3 # parse_agnostic + densepose output_nc = nc tocg = ConditionGenerator(opt, input1_nc=input1_nc, input2_nc=input2_nc, output_nc=output_nc, ngf=96, norm_layer=nn.BatchNorm2d) #### SPADEResBlock from network_generator import SPADEResBlock opt.semantic_nc = 7 opt.fine_height = 1024 opt.fine_width = 768 opt.ngf = 64 opt.num_upsampling_layers = "most" opt.norm_G = "spectralaliasinstance" opt.gen_semantic_nc = 7 generator = SPADEGenerator(opt, 3+3+3) generator.print_network() #### https://drive.google.com/open?id=1XJTCdRBOPVgVTmqzhVGFAgMm2NLkw5uQ&authuser=0 opt.tocg_checkpoint = "mtviton.pth" #### https://drive.google.com/open?id=1T5_YDUhYSSKPC_nZMk2NeC-XXUFoYeNy&authuser=0 opt.gen_checkpoint = "gen.pth" opt.cuda = False load_checkpoint(tocg, opt.tocg_checkpoint,opt) load_checkpoint_G(generator, opt.gen_checkpoint,opt) #### def test scope tocg.eval() generator.eval() opt.data_list = "test_pairs.txt" opt.datamode = "test" opt.dataroot = "zalando-hd-resized" opt.batch_size = 1 opt.shuffle = False opt.workers = 1 opt.semantic_nc = 13 test_dataset = CPDatasetTest(opt) test_loader = CPDataLoader(opt, test_dataset) def construct_images(img_tensors, img_names = [None]): #for img_tensor, img_name in zip(img_tensors, img_names): for img_tensor, img_name in zip(img_tensors, img_names): tensor = (img_tensor.clone() + 1) * 0.5 * 255 tensor = tensor.cpu().clamp(0, 255) try: array = tensor.numpy().astype('uint8') except: array = tensor.detach().numpy().astype('uint8') if array.shape[0] == 1: array = array.squeeze(0) elif array.shape[0] == 3: array = array.swapaxes(0, 1).swapaxes(1, 2) im = Image.fromarray(array) return im def single_pred_slim_func(opt, inputs, tocg = tocg, generator = generator): gauss = tgm.image.GaussianBlur((15, 15), (3, 3)) if opt.cuda: gauss = gauss.cuda() # Model if opt.cuda: tocg.cuda() tocg.eval() generator.eval() num = 0 iter_start_time = time.time() with torch.no_grad(): for inputs in [inputs]: if opt.cuda : #pose_map = inputs['pose'].cuda() pre_clothes_mask = inputs['cloth_mask'][opt.datasetting].cuda() #label = inputs['parse'] parse_agnostic = inputs['parse_agnostic'] agnostic = inputs['agnostic'].cuda() clothes = inputs['cloth'][opt.datasetting].cuda() # target cloth densepose = inputs['densepose'].cuda() #im = inputs['image'] #input_label, input_parse_agnostic = label.cuda(), parse_agnostic.cuda() input_parse_agnostic = parse_agnostic.cuda() pre_clothes_mask = torch.FloatTensor((pre_clothes_mask.detach().cpu().numpy() > 0.5).astype(np.float)).cuda() else : #pose_map = inputs['pose'] pre_clothes_mask = inputs['cloth_mask'][opt.datasetting] #label = inputs['parse'] parse_agnostic = inputs['parse_agnostic'] agnostic = inputs['agnostic'] clothes = inputs['cloth'][opt.datasetting] # target cloth densepose = inputs['densepose'] #im = inputs['image'] #input_label, input_parse_agnostic = label, parse_agnostic input_parse_agnostic = parse_agnostic pre_clothes_mask = torch.FloatTensor((pre_clothes_mask.detach().cpu().numpy() > 0.5).astype(np.float)) # down #pose_map_down = F.interpolate(pose_map, size=(256, 192), mode='bilinear') pre_clothes_mask_down = F.interpolate(pre_clothes_mask, size=(256, 192), mode='nearest') #input_label_down = F.interpolate(input_label, size=(256, 192), mode='bilinear') input_parse_agnostic_down = F.interpolate(input_parse_agnostic, size=(256, 192), mode='nearest') #agnostic_down = F.interpolate(agnostic, size=(256, 192), mode='nearest') clothes_down = F.interpolate(clothes, size=(256, 192), mode='bilinear') densepose_down = F.interpolate(densepose, size=(256, 192), mode='bilinear') shape = pre_clothes_mask.shape # multi-task inputs input1 = torch.cat([clothes_down, pre_clothes_mask_down], 1) input2 = torch.cat([input_parse_agnostic_down, densepose_down], 1) # forward flow_list, fake_segmap, warped_cloth_paired, warped_clothmask_paired = tocg(opt,input1, input2) # warped cloth mask one hot if opt.cuda : warped_cm_onehot = torch.FloatTensor((warped_clothmask_paired.detach().cpu().numpy() > 0.5).astype(np.float)).cuda() else : warped_cm_onehot = torch.FloatTensor((warped_clothmask_paired.detach().cpu().numpy() > 0.5).astype(np.float)) if opt.clothmask_composition != 'no_composition': if opt.clothmask_composition == 'detach': cloth_mask = torch.ones_like(fake_segmap) cloth_mask[:,3:4, :, :] = warped_cm_onehot fake_segmap = fake_segmap * cloth_mask if opt.clothmask_composition == 'warp_grad': cloth_mask = torch.ones_like(fake_segmap) cloth_mask[:,3:4, :, :] = warped_clothmask_paired fake_segmap = fake_segmap * cloth_mask # make generator input parse map fake_parse_gauss = gauss(F.interpolate(fake_segmap, size=(opt.fine_height, opt.fine_width), mode='bilinear')) fake_parse = fake_parse_gauss.argmax(dim=1)[:, None] if opt.cuda : old_parse = torch.FloatTensor(fake_parse.size(0), 13, opt.fine_height, opt.fine_width).zero_().cuda() else: old_parse = torch.FloatTensor(fake_parse.size(0), 13, opt.fine_height, opt.fine_width).zero_() old_parse.scatter_(1, fake_parse, 1.0) labels = { 0: ['background', [0]], 1: ['paste', [2, 4, 7, 8, 9, 10, 11]], 2: ['upper', [3]], 3: ['hair', [1]], 4: ['left_arm', [5]], 5: ['right_arm', [6]], 6: ['noise', [12]] } if opt.cuda : parse = torch.FloatTensor(fake_parse.size(0), 7, opt.fine_height, opt.fine_width).zero_().cuda() else: parse = torch.FloatTensor(fake_parse.size(0), 7, opt.fine_height, opt.fine_width).zero_() for i in range(len(labels)): for label in labels[i][1]: parse[:, i] += old_parse[:, label] # warped cloth N, _, iH, iW = clothes.shape flow = F.interpolate(flow_list[-1].permute(0, 3, 1, 2), size=(iH, iW), mode='bilinear').permute(0, 2, 3, 1) flow_norm = torch.cat([flow[:, :, :, 0:1] / ((96 - 1.0) / 2.0), flow[:, :, :, 1:2] / ((128 - 1.0) / 2.0)], 3) grid = make_grid(N, iH, iW,opt) warped_grid = grid + flow_norm warped_cloth = F.grid_sample(clothes, warped_grid, padding_mode='border') warped_clothmask = F.grid_sample(pre_clothes_mask, warped_grid, padding_mode='border') if opt.occlusion: warped_clothmask = remove_overlap(F.softmax(fake_parse_gauss, dim=1), warped_clothmask) warped_cloth = warped_cloth * warped_clothmask + torch.ones_like(warped_cloth) * (1-warped_clothmask) output = generator(torch.cat((agnostic, densepose, warped_cloth), dim=1), parse) # save output return output #save_images(output, unpaired_names, output_dir) #num += shape[0] #print(num) opt.clothmask_composition = "warp_grad" opt.occlusion = False opt.datasetting = "unpaired" def read_img_and_trans(dataset ,opt ,img_path): if type(img_path) in [type("")]: im = Image.open(img_path) else: im = img_path im = transforms.Resize(opt.fine_width, interpolation=2)(im) im = dataset.transform(im) return im import sys sys.path.insert(0, "fashion-eye-try-on") import os from PIL import Image import gradio as gr from cloth_segmentation import generate_cloth_mask def generate_cloth_mask_and_display(cloth_img): path = 'fashion-eye-try-on/cloth/cloth.jpg' if os.path.exists(path): os.remove(path) cloth_img.save(path) try: # os.system('.\cloth_segmentation\generate_cloth_mask.py') generate_cloth_mask() except Exception as e: print(e) return cloth_mask_img = Image.open("fashion-eye-try-on/cloth_mask/cloth.jpg") return cloth_mask_img def take_human_feature_from_dataset(dataset, idx): inputs_upper = list(torch.utils.data.DataLoader( [dataset[idx]], batch_size=1))[0] return { "parse_agnostic": inputs_upper["parse_agnostic"], "agnostic": inputs_upper["agnostic"], "densepose": inputs_upper["densepose"], } def take_all_feature_with_dataset(cloth_img_path, idx, opt = opt, dataset = test_dataset, only_show_human = False): if type(cloth_img_path) != type(""): assert hasattr(cloth_img_path, "save") cloth_img_path.save("tmp_cloth.jpg") cloth_img_path = "tmp_cloth.jpg" assert type(cloth_img_path) == type("") inputs_upper_dict = take_human_feature_from_dataset(dataset, idx) if only_show_human: return Image.fromarray((inputs_upper_dict["densepose"][0].numpy().transpose((1, 2, 0)) * 255).astype(np.uint8)) cloth_readed = read_img_and_trans(dataset, opt, cloth_img_path ) #assert ((cloth_readed - inputs_upper["cloth"][opt.datasetting][0]) ** 2).sum().numpy() < 1e-15 cloth_input = { opt.datasetting: cloth_readed[None,:] } mask_img = generate_cloth_mask_and_display( Image.open( cloth_img_path ) ) cloth_mask_input = { opt.datasetting: torch.Tensor((np.asarray(mask_img) / 255))[None, None, :] } inputs_upper_dict["cloth"] = cloth_input inputs_upper_dict["cloth_mask"] = cloth_mask_input return inputs_upper_dict def pred_func(cloth_img, pidx ): idx = int(pidx) im = cloth_img #### truly input inputs_upper_dict = take_all_feature_with_dataset( im, idx, only_show_human = False) output_slim = single_pred_slim_func(opt, inputs_upper_dict) output_img = construct_images(output_slim) return output_img option = st.selectbox( "Choose cloth image or Upload cloth image", ("Choose", "Upload", ) ) if type(option) != type(""): option = "Choose" img = None uploaded_file = None if option == "Upload": # To read file as bytes: uploaded_file = st.file_uploader("Upload img") if uploaded_file is not None: bytes_data = uploaded_file.getvalue() img = Image.open(BytesIO(bytes_data)) cloth_img = img.convert("RGB").resize((256 + 128, 512)) st.image(cloth_img) uploaded_file = st.selectbox( "Have Choose the image", ("Wait", "Have Done") ) else: img = image_select("Choose img", demo_images) #img = Image.open(img) cloth_img = img.convert("RGB").resize((256 + 128, 512)) st.image(cloth_img) uploaded_file = st.selectbox( "Have Choose the image", ("Wait", "Have Done") ) if img is not None and (uploaded_file is not "Wait" and uploaded_file is not None): cloth_img = img.convert("RGB").resize((768, 1024)) #pidx = 44 pidx_index_list = [44, 84, 67] poeses = [] for idx in range(len(pidx_index_list)): poeses.append( take_all_feature_with_dataset( cloth_img, pidx_index_list[idx], only_show_human = True) ) col1, col2, col3 = st.columns(3) with col1: st.header("Pose 0") pose_img = poeses[0] st.image(pose_img) b = pred_func(cloth_img, pidx_index_list[0]) st.image(b) with col2: st.header("Pose 1") pose_img = poeses[1] st.image(pose_img) b = pred_func(cloth_img, pidx_index_list[1]) st.image(b) with col3: st.header("Pose 2") pose_img = poeses[2] st.image(pose_img) b = pred_func(cloth_img, pidx_index_list[2]) st.image(b)