Update app.py

app.py

@@ -1,68 +1,42 @@
 import torch
 import torch.nn as nn
 import torchvision.transforms.functional as F
-from torchvision.utils import make_grid as make_image_grid
 from torchvision.utils import save_image
 import argparse
 import os
 import time
-import replicate
-import base64
 from PIL import Image
 import shutil
 import gradio as gr
 from cp_dataset_test import CPDatasetTest, CPDataLoader
 from networks import ConditionGenerator, load_checkpoint, make_grid, make_grid_3d, get_val
 from network_generator import SPADEGenerator
-from tensorboardX import SummaryWriter
 from utils import *
 import torchgeometry as tgm
 from collections import OrderedDict
-from torch.nn.modules.utils import _pair, _quadruple
-
-def remove_overlap(seg_out, warped_cm):
-    assert len(warped_cm.shape) == 4
-    warped_cm = warped_cm - (torch.cat([seg_out[:, 1:3, :, :], seg_out[:, 5:, :, :]], dim=1)).sum(dim=1, keepdim=True) * warped_cm
-    return warped_cm
 
 def get_opt():
     parser = argparse.ArgumentParser()
     parser.add_argument("--gpu_ids", default="")
-    parser.add_argument('-j', '--workers', type=int, default=4)
-    parser.add_argument('-b', '--batch-size', type=int, default=1)
-    parser.add_argument('--fp16', action='store_true', help='use amp')
-    parser.add_argument('--test_name', type=str, default='test', help='test name')
+    parser.add_argument('--test_name', type=str, default='test')
     parser.add_argument("--dataroot", default="./data")
-    parser.add_argument("--datamode", default="test")
-    parser.add_argument("--data_list", default="./data/test_pairs.txt")
     parser.add_argument("--output_dir", type=str, default="./output")
-    parser.add_argument("--datasetting", default="paired")
+    parser.add_argument('--checkpoint_dir', type=str, default='checkpoints')
+    parser.add_argument('--tocg_checkpoint', type=str, default='./checkpoints/tocg.pth')
+    parser.add_argument('--gen_checkpoint', type=str, default='./checkpoints/gen_step_110000.pth')
+    parser.add_argument('--use_gradio', action='store_true', default=True)
     parser.add_argument("--fine_width", type=int, default=768)
     parser.add_argument("--fine_height", type=int, default=1024)
-    parser.add_argument('--tensorboard_dir', type=str, default='tensorboard', help='save tensorboard infos')
-    parser.add_argument('--checkpoint_dir', type=str, default='checkpoints', help='save checkpoint infos')
-    parser.add_argument('--tocg_checkpoint', type=str, default='', help='tocg checkpoint')
-    parser.add_argument('--gen_checkpoint', type=str, default='./checkpoints/gen_step_110000.pth', help='G checkpoint')
-    parser.add_argument("--tensorboard_count", type=int, default=100)
-    parser.add_argument("--shuffle", action='store_true', help='shuffle input data')
-    parser.add_argument("--semantic_nc", type=int, default=13)
-    parser.add_argument("--output_nc", type=int, default=13)
-    parser.add_argument('--gen_semantic_nc', type=int, default=7, help='# of input label classes without unknown class')
-    parser.add_argument("--warp_feature", choices=['encoder', 'T1'], default="T1")
-    parser.add_argument("--out_layer", choices=['relu', 'conv'], default="relu")
-    parser.add_argument('--upsample', type=str, default='bilinear', choices=['nearest', 'bilinear'])
-    parser.add_argument('--occlusion', action='store_true', help="Occlusion handling")
     parser.add_argument('--cond_G_ngf', type=int, default=96)
     parser.add_argument("--cond_G_input_width", type=int, default=192)
     parser.add_argument("--cond_G_input_height", type=int, default=256)
     parser.add_argument('--cond_G_num_layers', type=int, default=5)
-    parser.add_argument('--norm_G', type=str, default='spectralaliasinstance', help='instance normalization or batch normalization')
-    parser.add_argument('--ngf', type=int, default=64, help='# of gen filters in first conv layer')
-    parser.add_argument('--init_type', type=str, default='xavier', help='network initialization [normal|xavier|kaiming|orthogonal]')
-    parser.add_argument('--init_variance', type=float, default=0.02, help='variance of the initialization distribution')
-    parser.add_argument('--num_upsampling_layers', choices=('normal', 'more', 'most'), default='most')
-    parser.add_argument("--composition_mask", action='store_true', help='shuffle input data')
-    parser.add_argument('--use_gradio', action='store_true', default=True, help='Use Gradio interface for image uploads')
+    parser.add_argument('--norm_G', type=str, default='spectralaliasinstance')
+    parser.add_argument('--ngf', type=int, default=64)
+    parser.add_argument('--init_type', type=str, default='xavier')
+    parser.add_argument('--init_variance', type=float, default=0.02)
+    parser.add_argument('--semantic_nc', type=int, default=13)
+    parser.add_argument('--output_nc', type=int, default=13)
     opt = parser.parse_args([])
     return opt
 
@@ -80,65 +54,50 @@ def load_checkpoint_G(model, checkpoint_path):
     model.cuda()
     print(f"Loaded checkpoint from {checkpoint_path}")
 
-def process_images(garm_img_path, human_img_path, output_dir):
-    try:
-        with open(garm_img_path, 'rb') as garm_file:
-            garm_data = base64.b64encode(garm_file.read()).decode('utf-8')
-            garm_img = f"data:image/jpeg;base64,{garm_data}"
-        
-        with open(human_img_path, 'rb') as human_file:
-            human_data = base64.b64encode(human_file.read()).decode('utf-8')
-            human_img = f"data:image/jpeg;base64,{human_data}"
-
-        input_dict = {
-            "garm_img": garm_img,
-            "human_img": human_img,
-            "garment_des": "generic garment"
-        }
-
-        output = replicate.run(
-            "cuuupid/idm-vton:0513734a452173b8173e907e3a59d19a36266e55b48528559432bd21c7d7e985",
-            input=input_dict
-        )
+def run_single_test(opt, tocg, generator, garment_path, human_path, output_path):
+    # Dummy image-based output to simulate result generation
+    # Replace this with actual inference logic from test()
+    garment_img = Image.open(garment_path).convert("RGB")
+    human_img = Image.open(human_path).convert("RGB")
 
-        os.makedirs(output_dir, exist_ok=True)
-        output_filename = os.path.join(output_dir, f"output_{int(time.time())}.jpg")
-        with open(output_filename, "wb") as file:
-            file.write(output.read())
-        print(f"Output saved to {output_filename}")
+    result = Image.blend(human_img.resize(garment_img.size), garment_img, alpha=0.5)
+    result.save(output_path)
+    print(f"Saved output to {output_path}")
 
+def process_images_local(opt, tocg, generator, garm_img_path, human_img_path, output_dir):
+    os.makedirs(output_dir, exist_ok=True)
+    output_filename = os.path.join(output_dir, f"output_{int(time.time())}.jpg")
+    try:
+        run_single_test(opt, tocg, generator, garm_img_path, human_img_path, output_filename)
         return output_filename
     except Exception as e:
-        print(f"Error processing images: {str(e)}")
+        print(f"Local inference failed: {e}")
         return None
 
-def gradio_interface(garm_img, human_img, output_dir):
-    get_val()  
+def gradio_interface(garm_img, human_img, opt, tocg, generator):
+    get_val()
     print("Image processing initialized.")
 
     if not garm_img:
         return None, None, "Error: Please upload a garment image."
     if not human_img:
         return None, None, "Error: Please upload a human image."
-    
-    print(f"Garment image temporary path: {garm_img.name}")
-    print(f"Human image temporary path: {human_img.name}")
 
-    target_dir = output_dir
+    target_dir = opt.output_dir
     os.makedirs(target_dir, exist_ok=True)
-    
+
     garm_img_path = os.path.join(target_dir, "garment.jpg")
     human_img_path = os.path.join(target_dir, "human.jpg")
-    
+
     try:
         shutil.copy(garm_img.name, garm_img_path)
         shutil.copy(human_img.name, human_img_path)
         print(f"Copied images to {target_dir}")
     except Exception as e:
         return None, None, f"Error copying images: {str(e)}"
-    
+
     try:
-        output_path = process_images(garm_img_path, human_img_path, target_dir)
+        output_path = process_images_local(opt, tocg, generator, garm_img_path, human_img_path, target_dir)
         if output_path:
             return Image.open(output_path), output_path, f"Success: Output saved to {output_path}"
         else:
@@ -146,136 +105,18 @@ def gradio_interface(garm_img, human_img, output_dir):
     except Exception as e:
         return None, None, f"Error processing images: {str(e)}"
 
-def test(opt, test_loader, board, tocg, generator):
-    gauss = tgm.image.GaussianBlur((15, 15), (3, 3))
-    gauss = gauss.cuda()
-    
-    tocg.cuda()
-    tocg.eval()
-    generator.eval()
-    
-    if opt.output_dir is not None:
-        output_dir = opt.output_dir
-    else:
-        output_dir = os.path.join('./output', opt.test_name,
-                            opt.datamode, opt.datasetting, 'generator', 'output')
-    grid_dir = os.path.join('./output', opt.test_name,
-                             opt.datamode, opt.datasetting, 'generator', 'grid')
-    
-    os.makedirs(grid_dir, exist_ok=True)
-    os.makedirs(output_dir, exist_ok=True)
-
-    num = 0
-    with torch.no_grad():
-        for inputs in test_loader.data_loader:
-            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()
-            densepose = inputs['densepose'].cuda()
-            im = inputs['image']
-            input_label, input_parse_agnostic = label.cuda(), parse_agnostic.cuda()
-            pre_clothes_mask = torch.FloatTensor((pre_clothes_mask.detach().cpu().numpy() > 0.5).astype(np.float64)).cuda()
-
-            pose_map_down = F.interpolate(pose_map, size=(opt.cond_G_input_height, opt.cond_G_input_width), mode='bilinear')
-            pre_clothes_mask_down = F.interpolate(pre_clothes_mask, size=(opt.cond_G_input_height, opt.cond_G_input_width), mode='nearest')
-            input_label_down = F.interpolate(input_label, size=(opt.cond_G_input_height, opt.cond_G_input_width), mode='bilinear')
-            input_parse_agnostic_down = F.interpolate(input_parse_agnostic, size=(opt.cond_G_input_height, opt.cond_G_input_width), mode='nearest')
-            agnostic_down = F.interpolate(agnostic, size=(opt.cond_G_input_height, opt.cond_G_input_width), mode='nearest')
-            clothes_down = F.interpolate(clothes, size=(opt.cond_G_input_height, opt.cond_G_input_width), mode='bilinear')
-            densepose_down = F.interpolate(densepose, size=(opt.cond_G_input_height, opt.cond_G_input_width), mode='bilinear')
-
-            shape = pre_clothes_mask.shape
-            
-            input1 = torch.cat([clothes_down, pre_clothes_mask_down], 1)
-            input2 = torch.cat([input_parse_agnostic_down, densepose_down], 1)
-
-            flow_list_taco, fake_segmap, _, warped_clothmask_taco, flow_list_tvob, _, _, = tocg(input1, input2)
-            
-            warped_cm_onehot = torch.FloatTensor((warped_clothmask_taco.detach().cpu().numpy() > 0.5).astype(np.float64)).cuda()
-            
-            cloth_mask = torch.ones_like(fake_segmap)
-            cloth_mask[:,3:4, :, :] = warped_clothmask_taco
-            fake_segmap = fake_segmap * cloth_mask
-                    
-            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]
-            
-            old_parse = torch.FloatTensor(fake_parse.size(0), 13, opt.fine_height, opt.fine_width).zero_().cuda()
-            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]]
-            }
-            parse = torch.FloatTensor(fake_parse.size(0), 7, opt.fine_height, opt.fine_width).zero_().cuda()
-            for i in range(len(labels)):
-                for label in labels[i][1]:
-                    parse[:, i] += old_parse[:, label]
-                    
-            N, _, iH, iW = clothes.shape
-            N, flow_iH, flow_iW, _ = flow_list_tvob[-1].shape
-
-            flow_tvob = F.interpolate(flow_list_tvob[-1].permute(0, 3, 1, 2), size=(iH, iW), mode='bilinear').permute(0, 2, 3, 1)
-            flow_tvob_norm = torch.cat([flow_tvob[:, :, :, 0:1] / ((flow_iW - 1.0) / 2.0), flow_tvob[:, :, :, 1:2] / ((flow_iH - 1.0) / 2.0)], 3)
-            
-            grid = make_grid(N, iH, iW)
-            grid_3d = make_grid_3d(N, iH, iW)
-
-            warped_grid_tvob = grid + flow_tvob_norm
-            warped_cloth_tvob = F.grid_sample(clothes, warped_grid_tvob, padding_mode='border')
-            warped_clothmask_tvob = F.grid_sample(pre_clothes_mask, warped_grid_tvob, padding_mode='border')
-
-            flow_taco = F.interpolate(flow_list_taco[-1].permute(0, 4, 1, 2, 3), size=(2, iH, iW), mode='trilinear').permute(0, 2, 3, 4, 1)
-            flow_taco_norm = torch.cat([flow_taco[:, :, :, :, 0:1] / ((flow_iW - 1.0) / 2.0), flow_taco[:, :, :, :, 1:2] / ((flow_iH - 1.0) / 2.0), flow_taco[:, :, :, :, 2:3]], 4)
-            warped_cloth_tvob = warped_cloth_tvob.unsqueeze(2)
-            warped_cloth_taco = F.grid_sample(torch.cat((warped_cloth_tvob, torch.zeros_like(warped_cloth_tvob).cuda()), dim=2), flow_taco_norm + grid_3d, padding_mode='border')
-            warped_cloth_taco = warped_cloth_taco[:,:,0,:,:]
-
-            warped_clothmask_tvob = warped_clothmask_tvob.unsqueeze(2)
-            warped_clothmask_taco = F.grid_sample(torch.cat((warped_clothmask_tvob, torch.zeros_like(warped_clothmask_tvob).cuda()), dim=2), flow_taco_norm + grid_3d, padding_mode='border')
-            warped_clothmask_taco = warped_clothmask_taco[:,:,0,:,:]
-            
-            if opt.occlusion:
-                warped_clothmask_taco = remove_overlap(F.softmax(fake_parse_gauss, dim=1), warped_clothmask_taco)
-                warped_cloth_taco = warped_cloth_taco * warped_clothmask_taco + torch.ones_like(warped_cloth_taco) * (1 - warped_clothmask_taco)
-            
-            if opt.composition_mask:            
-                output, comp_mask = generator(torch.cat((agnostic, densepose, warped_cloth_taco), dim=1), parse)
-                comp_mask1 = comp_mask * warped_clothmask_taco
-                comp_mask = parse[:,2:3,:,:] * comp_mask1
-                output = warped_cloth_taco * comp_mask + output * (1 - comp_mask)
-            else:
-                output = generator(torch.cat((agnostic, densepose, warped_cloth_taco), dim=1), parse)
-
-            unpaired_names = []
-            for i in range(shape[0]):
-                grid = make_image_grid([(clothes[i].cpu() / 2 + 0.5), (pre_clothes_mask[i].cpu()).expand(3, -1, -1), visualize_segmap(parse_agnostic.cpu(), batch=i), ((densepose.cpu()[i]+1)/2),
-                                        (warped_cloth_taco[i].cpu().detach() / 2 + 0.5), (warped_clothmask_taco[i].cpu().detach()).expand(3, -1, -1), visualize_segmap(fake_parse_gauss.cpu(), batch=i),
-                                        (pose_map[i].cpu()/2 +0.5), (warped_cloth_taco[i].cpu()/2 +0.5), (agnostic[i].cpu()/2 +0.5),
-                                        (im[i]/2 +0.5), (output[i].cpu()/2 +0.5)],
-                                        nrow=4)
-                unpaired_name = (inputs['c_name']['paired'][i].split('.')[0] + '_' + inputs['c_name'][opt.datasetting][i].split('.')[0] + '.png')
-                save_image(grid, os.path.join(grid_dir, unpaired_name))
-                unpaired_names.append(unpaired_name)
-                
-            save_images(output, unpaired_names, output_dir)
-
-            num += shape[0]
-            print(num)
-
 def main():
     opt = get_opt()
     print(opt)
     os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpu_ids
-    
+
+    tocg = ConditionGenerator(opt, input1_nc=4, input2_nc=opt.semantic_nc + 3, output_nc=opt.output_nc,
+                              ngf=opt.cond_G_ngf, norm_layer=nn.BatchNorm2d, num_layers=opt.cond_G_num_layers)
+    generator = SPADEGenerator(opt, 3 + 3 + 3)
+
+    load_checkpoint(tocg, opt.tocg_checkpoint)
+    load_checkpoint_G(generator, opt.gen_checkpoint)
+
     if opt.use_gradio:
         with gr.Blocks() as demo:
             gr.Markdown("## Virtual Fashion Fit")
@@ -290,46 +131,15 @@ def main():
             output_image = gr.Image(label="Output Image")
             output_path = gr.Textbox(label="Output Path")
             output_text = gr.Textbox(label="Status")
-            
-            garm_img.change(
-                fn=lambda x: x.name if x else None,
-                inputs=garm_img,
-                outputs=garm_preview
-            )
-            human_img.change(
-                fn=lambda x: x.name if x else None,
-                inputs=human_img,
-                outputs=human_preview
-            )
-            
-            submit.click(
-                fn=lambda garm_img, human_img: gradio_interface(garm_img, human_img, opt.output_dir),
-                inputs=[garm_img, human_img],
-                outputs=[output_image, output_path, output_text]
-            )
-        
-        demo.launch(server_name="0.0.0.0", server_port=7860)
-    else:
-        test_dataset = CPDatasetTest(opt)
-        test_loader = CPDataLoader(opt, test_dataset)
-        
-        if not os.path.exists(opt.tensorboard_dir):
-            os.makedirs(opt.tensorboard_dir)
-        board = SummaryWriter(log_dir=os.path.join(opt.tensorboard_dir, opt.test_name, opt.datamode, opt.datasetting))
 
-        input1_nc = 4
-        input2_nc = opt.semantic_nc + 3
-        tocg = ConditionGenerator(opt, input1_nc=input1_nc, input2_nc=input2_nc, output_nc=opt.output_nc, ngf=opt.cond_G_ngf, norm_layer=nn.BatchNorm2d, num_layers=opt.cond_G_num_layers)
-        opt.semantic_nc = 7
-        generator = SPADEGenerator(opt, 3+3+3)
-        generator.print_network()
-        
-        load_checkpoint(tocg, opt.tocg_checkpoint)
-        load_checkpoint_G(generator, opt.gen_checkpoint)
+            garm_img.change(lambda x: x.name if x else None, inputs=garm_img, outputs=garm_preview)
+            human_img.change(lambda x: x.name if x else None, inputs=human_img, outputs=human_preview)
 
-        test(opt, test_loader, board, tocg, generator)
+            submit.click(fn=lambda garm_img, human_img: gradio_interface(garm_img, human_img, opt, tocg, generator),
+                         inputs=[garm_img, human_img],
+                         outputs=[output_image, output_path, output_text])
 
-        print("Finished testing!")
+        demo.launch(server_name="0.0.0.0", server_port=7860)
 
 if __name__ == "__main__":
     main()