app.py

app.py

@@ -0,0 +1,157 @@
+import gradio as gr
+import torch
+import pickle
+import time
+from options.test_options import TestOptions
+from data.data_loader_test import CreateDataLoader
+from models.networks import ResUnetGenerator, load_checkpoint
+from models.afwm import AFWM
+import torch.nn as nn
+import os
+import numpy as np
+import torch
+import cv2
+import torch.nn.functional as F
+from torchvision import utils
+from util import flow_util
+
+def de_offset(s_grid):
+    [b,_,h,w] = s_grid.size()
+
+
+    x = torch.arange(w).view(1, -1).expand(h, -1).float()
+    y = torch.arange(h).view(-1, 1).expand(-1, w).float()
+    x = 2*x/(w-1)-1
+    y = 2*y/(h-1)-1
+    grid = torch.stack([x,y], dim=0).float().cuda()
+    grid = grid.unsqueeze(0).expand(b, -1, -1, -1)
+
+    offset = grid - s_grid
+
+    offset_x = offset[:,0,:,:] * (w-1) / 2
+    offset_y = offset[:,1,:,:] * (h-1) / 2
+
+    offset = torch.cat((offset_y,offset_x),0)
+    
+    return  offset
+
+def tryon(person,cloth,edge):
+
+    #save images in folders
+    cv2.imwrite('./data/test_ma_img/000001_0.jpg', person)
+    cv2.imwrite('./data/test_edge/000001_1.jpg', edge)
+    cv2.imwrite('./data/test_clothes/000001_1.jpg', cloth)
+
+    with open('opt.pkl', 'rb') as handle:
+        opt = pickle.load(handle)  
+    
+    f2c = flow_util.flow2color()
+    start_epoch, epoch_iter = 1, 0
+
+    data_loader = CreateDataLoader(opt)
+    dataset = data_loader.load_data()
+    dataset_size = len(data_loader)     #must be 1
+    print(dataset_size)
+
+    warp_model = AFWM(opt, 3)
+    print(warp_model)
+    warp_model.eval()
+    warp_model.cuda()
+    load_checkpoint(warp_model, opt.warp_checkpoint)
+
+    gen_model = ResUnetGenerator(7, 4, 5, ngf=64, norm_layer=nn.BatchNorm2d)
+    gen_model.eval()
+    gen_model.cuda()
+    load_checkpoint(gen_model, opt.gen_checkpoint)
+
+    total_steps = (start_epoch-1) * dataset_size + epoch_iter
+    step = 0
+    step_per_batch = dataset_size / opt.batchSize
+    
+    if not os.path.exists('our_t_results'):
+        os.mkdir('our_t_results')
+
+    for epoch in range(1,2):
+
+        for i, data in enumerate(dataset, start=epoch_iter):
+            iter_start_time = time.time()
+            total_steps += opt.batchSize
+            epoch_iter += opt.batchSize
+
+            real_image = data['image']
+            clothes = data['clothes']
+            ##edge is extracted from the clothes image with the built-in function in python
+            edge = data['edge']
+            edge = torch.FloatTensor((edge.detach().numpy() > 0.5).astype(np.int64))
+            clothes = clothes * edge
+
+            flow_out = warp_model(real_image.cuda(), clothes.cuda())
+            warped_cloth, last_flow, = flow_out
+            warped_edge = F.grid_sample(edge.cuda(), last_flow.permute(0, 2, 3, 1),
+                            mode='bilinear', padding_mode='zeros')
+
+            gen_inputs = torch.cat([real_image.cuda(), warped_cloth, warped_edge], 1)
+            gen_outputs = gen_model(gen_inputs)
+            p_rendered, m_composite = torch.split(gen_outputs, [3, 1], 1)
+            p_rendered = torch.tanh(p_rendered)
+            m_composite = torch.sigmoid(m_composite)
+            m_composite = m_composite * warped_edge
+            p_tryon = warped_cloth * m_composite + p_rendered * (1 - m_composite)
+
+            path = 'results/' + opt.name
+            os.makedirs(path, exist_ok=True)
+            #sub_path = path + '/PFAFN'
+            #os.makedirs(sub_path,exist_ok=True)
+            print(data['p_name'])
+
+            if step % 1 == 0:
+                
+                ## save try-on image only
+
+                utils.save_image(
+                    p_tryon,
+                    os.path.join('./our_t_results', data['p_name'][0]),
+                    nrow=int(1),
+                    normalize=True,
+                    value_range=(-1,1),
+                )
+                
+                ## save person image, garment, flow, warped garment, and try-on image
+                
+                #a = real_image.float().cuda()
+                #b = clothes.cuda()
+                #flow_offset = de_offset(last_flow)
+                #flow_color = f2c(flow_offset).cuda()
+                #c= warped_cloth.cuda()
+                #d = p_tryon
+                #combine = torch.cat([a[0],b[0], flow_color, c[0], d[0]], 2).squeeze()
+                #utils.save_image(
+                #    combine,
+                #    os.path.join('./im_gar_flow_wg', data['p_name'][0]),
+                #    nrow=int(1),
+                #    normalize=True,
+                #    range=(-1,1),
+                #)
+                
+
+            step += 1
+            if epoch_iter >= dataset_size:
+                break
+    
+    result_img = cv2.imread('./our_t_results/000001_0.jpg')
+    return result_img
+
+demo = gr.Interface(fn=tryon, 
+             inputs=[gr.inputs.Image(label="Person"),gr.inputs.Image(label="Cloth"),gr.inputs.Image(label="Edge")],
+             outputs="image"
+             )
+
+# def pp(inp1,inp2):
+#     return inp1+" hello "+inp2
+
+# demo2 = gr.Interface(fn=pp, 
+#              inputs=[gr.inputs.Textbox(lines=5, label="Input Text"),gr.inputs.Textbox(lines=5, label="Input Text2")],
+#              outputs=gr.outputs.Textbox(label="Generated Text"),
+            #  )
+             
+demo.launch()