Rebuttal version

Author: kmbae

config.py

@@ -28,6 +28,8 @@ def add_argument_group(name):
 data_arg.add_argument('--a_grayscale', type=str2bool, default=False)
 data_arg.add_argument('--b_grayscale', type=str2bool, default=False)
 data_arg.add_argument('--num_worker', type=int, default=12)
+data_arg.add_argument('--dataset_A1', type=str, default='../data/edges2handbags')
+data_arg.add_argument('--dataset_A2', type=str, default='../data/edges2shoes')
 
 # Training / test parameters
 train_arg = add_argument_group('Training')

data_loader.py

@@ -6,6 +6,8 @@
 
 import torch
 from torchvision import transforms
+#from skimage import feature
+#from skimage.color import rgb2gray
 
 PIX2PIX_DATASETS = [
     'facades', 'cityscapes', 'maps', 'edges2shoes', 'edges2handbags']
@@ -67,9 +69,12 @@ def __init__(self, root, scale_size, data_type, skip_pix2pix_processing=False):
 
     def __getitem__(self, index):
         image = Image.open(self.paths[index]).convert('RGB')
+        edges = Image.open(self.paths[index].replace('/A','/B1')).convert('RGB')
+        #edges = image.filter(ImageFilter.FIND_EDGES)
+
         #if self.data_type=='B':
         #    image = image.filter(ImageFilter.MinFilter(3))
-        return self.transform(image)
+        return {'image':self.transform(image), 'edges':self.transform(edges)}
 
     def __len__(self):
         return len(self.paths)

main.py

@@ -23,11 +23,11 @@ def main(config):
         batch_size = config.sample_per_image
 
     a_data_loader, b_data_loader = get_loader(
-            data_path, batch_size, config.input_scale_size,
+            config.dataset_A1, batch_size, config.input_scale_size,
             config.num_worker, config.skip_pix2pix_processing)
 
     a1_data_loader, b1_data_loader = get_loader(
-            '../data/edges2handbags', batch_size, config.input_scale_size,
+            config.dataset_A2, batch_size, config.input_scale_size,
             config.num_worker, config.skip_pix2pix_processing)
 
     trainer = Trainer(config, a_data_loader, b_data_loader, a1_data_loader, b1_data_loader)

test.py

@@ -0,0 +1,76 @@
+from models import *
+import torch
+from torch import nn
+import glob
+from torchvision import transforms
+import tqdm
+from config import get_config
+from PIL import Image
+import argparse
+import os
+
+# Usage example
+# run test.py --load=../logs/edges2shoes_2018-09-07_17-05-19 --iter=10000 --con=../validation_image/handbag_picking --sty=../validation_image/valid_x_A1
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--load", type=str, help="Saved file dir")
+parser.add_argument("--iter", type=str, help="Number of iteration")
+parser.add_argument("--con", type=str, help="Dir of content images")
+parser.add_argument("--sty", type=str, help="Dir of style images")
+args = parser.parse_args()
+
+scale_size = 64
+transform = transforms.Compose([
+    transforms.Scale(scale_size),
+    transforms.ToTensor(),
+    #transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+    ])
+ToPILImage = transforms.ToPILImage()
+
+if __name__=='__main__':
+    config, unparsed = get_config()
+    conv_dims, deconv_dims = [64, 128, 256, 512], [256, 128, 64]
+    a_channel = 3
+    b_channel = 3
+    num_gpu = torch.cuda.device_count()
+    G = GeneratorCNN_g(a_channel+b_channel, b_channel, conv_dims, deconv_dims, num_gpu)
+    F = GeneratorCNN(a_channel, a_channel, conv_dims, deconv_dims, num_gpu)
+
+    G = nn.DataParallel(G.cuda(),device_ids=range(torch.cuda.device_count()))
+    F = nn.DataParallel(F.cuda(),device_ids=range(torch.cuda.device_count()))
+    print('Loading model')
+    G.load_state_dict(torch.load(args.load + '/G_{}.pth'.format(args.iter)))
+    F.load_state_dict(torch.load(args.load + '/F_{}.pth'.format(args.iter)))
+    print('Model loaded')
+
+    if not os.path.exists('./results'):
+        os.mkdir('./results')
+
+    list_con = os.listdir(args.con)
+    list_sty = os.listdir(args.sty)
+    img_con = []
+    img_sty = []
+    for i in (list_con):
+        img_con.append(transform(Image.open(args.con+'/'+i)))
+    for j in (list_sty):
+        img_sty.append(transform(Image.open(args.sty+'/'+j)))
+
+    with torch.no_grad():
+        G.eval()
+        F.eval()
+        for i, con_tmp in tqdm.tqdm(enumerate(img_con)):
+            con_tmp = torch.unsqueeze(con_tmp.cuda(),0)
+            if not os.path.exists('./results/{}'.format(list_con[i].split('.')[0])):
+                os.mkdir('./results/{}'.format(list_con[i].split('.')[0]))
+            for j, sty_tmp in enumerate(img_sty):
+                sty_tmp = torch.unsqueeze(sty_tmp.cuda(),0)
+                img_out = G(F(con_tmp), sty_tmp)
+                img_out = ToPILImage(img_out.data[0].cpu())
+                img_out.save('./results/{}/{}_{}.jpg'.format(list_con[i].split('.')[0],
+                    list_con[i].split('.')[0],list_sty[j].split('.')[0]))
+
+
+
+
+
+

trainer.py

@@ -17,7 +17,7 @@
 import ipdb
 from tensorboardX import SummaryWriter
 from datetime import datetime
-
+import torchvision
 tmp = datetime.now()
 
 writer = SummaryWriter('../runs/' + str(tmp))
@@ -194,13 +194,35 @@ def train(self):
             self.G.parameters(),
             lr=self.lr, betas=(self.beta1, self.beta2))
 
-        A_loader, B_loader = iter(self.a_data_loader), iter(self.b_data_loader)
-        valid_x_A, valid_x_B = torch.Tensor(np.load('../valid_x_A1.npy')), torch.Tensor(np.load('../valid_x_B1.npy'))
-        valid_x_A, valid_x_B = self._get_variable(valid_x_A), self._get_variable(valid_x_B)
+        #A_loader, B_loader = iter(self.a_data_loader), iter(self.b_data_loader)
+        A_loader = iter(self.a_data_loader)
+        #valid_x_A, valid_x_B = torch.Tensor(np.load('../valid_x_A2.npy')), torch.Tensor(np.load('../valid_x_B2.npy'))
         #self._get_variable(A_loader.next()), self._get_variable(B_loader.next())
-        A1_loader, B1_loader = iter(self.a1_data_loader), iter(self.b1_data_loader)
-        valid_x_A1, valid_x_B1=torch.Tensor(np.load('../valid_x_A2.npy')), torch.Tensor(np.load('../valid_x_B2.npy'))
-        valid_x_A1, valid_x_B1 = self._get_variable(valid_x_A1), self._get_variable(valid_x_B1)
+        #A1_loader, B1_loader = iter(self.a1_data_loader), iter(self.b1_data_loader)
+        A1_loader = iter(self.a1_data_loader)
+        #valid_x_A1, valid_x_B1=torch.Tensor(np.load('../valid_x_A2_chair.npy')), torch.Tensor(np.load('../valid_x_B2_chair.npy'))
+        try:
+            valid_x_A, valid_x_B = torch.Tensor(np.load(self.config.dataset_A1+'_A.npy')), torch.Tensor(np.load(self.config.dataset_A1+'_B.npy'))
+            valid_x_A, valid_x_B = self._get_variable(valid_x_A), self._get_variable(valid_x_B)
+            valid_x_A1, valid_x_B1=torch.Tensor(np.load(self.config.dataset_A2+'_A.npy')), torch.Tensor(np.load(self.config.dataset_A2+'_B.npy'))
+            valid_x_A1, valid_x_B1 = self._get_variable(valid_x_A1), self._get_variable(valid_x_B1)
+        except:
+            print('Cannot load validation file. Creating new validation file')
+            x_A1 = A_loader.next()
+            x_A2 = A1_loader.next()
+            x_A1, x_B1 = x_A1['image'], x_A1['edges']
+            x_A2, x_B2 = x_A2['image'], x_A2['edges']
+            np.save(self.config.dataset_A1+'_A.npy',np.array(x_A1))
+            np.save(self.config.dataset_A1+'_B.npy',np.array(x_B1))
+            np.save(self.config.dataset_A2+'_A.npy',np.array(x_A2))
+            np.save(self.config.dataset_A2+'_B.npy',np.array(x_B2))
+
+            valid_x_A, valid_x_B = torch.Tensor(np.load(self.config.dataset_A1+'_A.npy')), torch.Tensor(np.load(self.config.dataset_A1+'_B.npy'))
+            valid_x_A, valid_x_B = self._get_variable(valid_x_A), self._get_variable(valid_x_B)
+            valid_x_A1, valid_x_B1=torch.Tensor(np.load(self.config.dataset_A2+'_A.npy')), torch.Tensor(np.load(self.config.dataset_A2+'_B.npy'))
+            valid_x_A1, valid_x_B1 = self._get_variable(valid_x_A1), self._get_variable(valid_x_B1)
+
+
         #self._get_variable(A1_loader.next()), self._get_variable(B1_loader.next())
         #ipdb.set_trace()
 
@@ -212,20 +234,23 @@ def train(self):
         for step in trange(self.start_step, self.max_step):
             try:
                 x_A1 = A_loader.next()
-                x_B1 = B_loader.next()
+                #x_B1 = B_loader.next()
             except StopIteration:
                 A_loader = iter(self.a_data_loader)
-                B_loader = iter(self.b_data_loader)
+                #B_loader = iter(self.b_data_loader)
                 x_A1 = A_loader.next()
-                x_B1 = B_loader.next()
+                #x_B1 = B_loader.next()
             try:
                 x_A2 = A1_loader.next()
-                x_B2 = B1_loader.next()
+                #x_B2 = B1_loader.next()
             except StopIteration:
                 A1_loader = iter(self.a1_data_loader)
-                B1_loader = iter(self.b1_data_loader)
+                #B1_loader = iter(self.b1_data_loader)
                 x_A2 = A1_loader.next()
-                x_B2 = B1_loader.next()
+                #x_B2 = B1_loader.next()
+
+            x_A1, x_B1 = x_A1['image'], x_A1['edges']
+            x_A2, x_B2 = x_A2['image'], x_A2['edges']
             if x_A1.size(0) != x_B1.size(0) or x_A2.size(0) != x_B2.size(0) or x_A1.size(0) != x_A2.size(0):
                 print("[!] Sampled dataset from A and B have different # of data. Try resampling...")
                 continue
@@ -422,17 +447,17 @@ def train(self):
 
                 self.generate_with_A(valid_x_A, valid_x_A1, self.model_dir, idx=step)
                 self.generate_with_B(valid_x_A1, valid_x_A, self.model_dir, idx=step)
-                writer.add_scalars('loss_G', {'l_g':l_g,'l_gan_A':l_gan_A,'l_const_A':l_const_A,
-                    'l_f':l_f, 'l_const_AB': l_const_AB},
+                writer.add_scalars('loss_G', {'l_g':l_g,'l_gan_A':l_gan_A,
+                    'l_f':l_f},
                     step)
+                writer.add_scalars('loss_F', {'l_const_A':l_const_A, 'l_const_AB': l_const_AB}, step)
                     #'l_const_B':l_const_B,'l_const_AB':l_const_AB,'l_const_BA':l_const_BA}, step)
                 writer.add_scalars('loss_D', {'l_d_A':l_d_A,'l_d_B':l_d_B}, step)
-
-            if step % self.save_step == self.save_step - 1:
+            if step % self.save_step == 0:
                 print("[*] Save models to {}...".format(self.model_dir))
 
-                torch.save(self.G.state_dict(), '{}/G_AB_{}.pth'.format(self.model_dir, step))
-                torch.save(self.F.state_dict(), '{}/G_BA_{}.pth'.format(self.model_dir, step))
+                torch.save(self.G.state_dict(), '{}/G_{}.pth'.format(self.model_dir, step))
+                torch.save(self.F.state_dict(), '{}/F_{}.pth'.format(self.model_dir, step))
 
                 torch.save(self.D_S.state_dict(), '{}/D_A_{}.pth'.format(self.model_dir, step))
                 torch.save(self.D_H.state_dict(), '{}/D_B_{}.pth'.format(self.model_dir, step))
@@ -447,6 +472,13 @@ def generate_with_A(self, inputs, input_ref, path, idx=None, tf_board=True):
         #x_ABA_path = '{}/{}_x_ABA.png'.format(path, idx)
 
         vutils.save_image(x_ABA.data, x_AB_path)
+        if not os.path.isdir('{}/{}_A1'.format(path, idx)):
+            os.makedirs('{}/{}_A1'.format(path, idx))
+        for i in range(x_ABA.size(0)):
+            tmp = x_ABA[i].detach().cpu()
+            tmp = torchvision.transforms.ToPILImage()(tmp)
+            tmp.save('{}/{}_A1/{}.png'.format(path, idx, i))
+
         print("[*] Samples saved: {}".format(x_AB_path))
         if tf_board:
             writer.add_image('x_A1f', x_AB[:16], idx)
@@ -468,6 +500,13 @@ def generate_with_B(self, inputs, input_ref, path, idx=None, tf_board=True):
         #x_BAB_path = '{}/{}_x_BAB.png'.format(path, idx)
 
         vutils.save_image(x_BAB.data, x_BA_path)
+        if not os.path.isdir('{}/{}_A2'.format(path, idx)):
+            os.makedirs('{}/{}_A2'.format(path, idx))
+        for i in range(x_BAB.size(0)):
+            tmp = x_BAB[i].detach().cpu()
+            tmp = torchvision.transforms.ToPILImage()(tmp)
+            tmp.save('{}/{}_A2/{}.png'.format(path, idx, i))
+
         print("[*] Samples saved: {}".format(x_BA_path))
         if tf_board:
             writer.add_image('x_A2f', x_BA[:16], idx)