predict single image

7 files changed, 31 insertions, 37 deletions

diff --git a/README.md b/README.md
index e2dbfbb..3d3eafe 100644
--- a/README.md
+++ b/README.md
@@ -149,7 +149,6 @@ This will combine each pair of images (A,B) into a single image file, ready for
 
 ## TODO
 - add reflection and other padding layers.
-- add more preprocessing options.
 
 ## Related Projects:
 [CycleGAN](https://github.com/junyanz/CycleGAN): Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks  
diff --git a/models/models.py b/models/models.py
index efcd898..d5bb9d8 100644
--- a/models/models.py
+++ b/models/models.py
@@ -13,7 +13,7 @@ def create_model(opt):
     elif opt.model == 'test':
         assert(opt.dataset_mode == 'single')
         from .test_model import TestModel
-        model = OneDirectionTestModel()
+        model = TestModel()
     else:
         raise ValueError("Model [%s] not recognized." % opt.model)
     model.initialize(opt)
diff --git a/models/pix2pix_model.py b/models/pix2pix_model.py
index 4581d33..e44529b 100644
--- a/models/pix2pix_model.py
+++ b/models/pix2pix_model.py
@@ -8,6 +8,7 @@ from util.image_pool import ImagePool
 from .base_model import BaseModel
 from . import networks
 
+
 class Pix2PixModel(BaseModel):
     def name(self):
         return 'Pix2PixModel'
@@ -23,12 +24,12 @@ class Pix2PixModel(BaseModel):
 
         # load/define networks
         self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf,
-                                    opt.which_model_netG, opt.norm, opt.use_dropout, self.gpu_ids)
+                                      opt.which_model_netG, opt.norm, opt.use_dropout, self.gpu_ids)
         if self.isTrain:
             use_sigmoid = opt.no_lsgan
             self.netD = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf,
-                                         opt.which_model_netD,
-                                         opt.n_layers_D, opt.norm, use_sigmoid, self.gpu_ids)
+                                          opt.which_model_netD,
+                                          opt.n_layers_D, opt.norm, use_sigmoid, self.gpu_ids)
         if not self.isTrain or opt.continue_train:
             self.load_network(self.netG, 'G', opt.which_epoch)
             if self.isTrain:
@@ -71,7 +72,7 @@ class Pix2PixModel(BaseModel):
         self.fake_B = self.netG.forward(self.real_A)
         self.real_B = Variable(self.input_B, volatile=True)
 
-    #get image paths
+    # get image paths
     def get_image_paths(self):
         return self.image_paths
 
@@ -83,7 +84,7 @@ class Pix2PixModel(BaseModel):
         self.loss_D_fake = self.criterionGAN(self.pred_fake, False)
 
         # Real
-        real_AB = torch.cat((self.real_A, self.real_B), 1)#.detach()
+        real_AB = torch.cat((self.real_A, self.real_B), 1)
         self.pred_real = self.netD.forward(real_AB)
         self.loss_D_real = self.criterionGAN(self.pred_real, True)
 
@@ -118,10 +119,10 @@ class Pix2PixModel(BaseModel):
 
     def get_current_errors(self):
         return OrderedDict([('G_GAN', self.loss_G_GAN.data[0]),
-                ('G_L1', self.loss_G_L1.data[0]),
-                ('D_real', self.loss_D_real.data[0]),
-                ('D_fake', self.loss_D_fake.data[0])
-        ])
+                            ('G_L1', self.loss_G_L1.data[0]),
+                            ('D_real', self.loss_D_real.data[0]),
+                            ('D_fake', self.loss_D_fake.data[0])
+                            ])
 
     def get_current_visuals(self):
         real_A = util.tensor2im(self.real_A.data)
diff --git a/models/test_model.py b/models/test_model.py
index a356263..65aa088 100644
--- a/models/test_model.py
+++ b/models/test_model.py
@@ -10,37 +10,32 @@ class TestModel(BaseModel):
         return 'TestModel'
 
     def initialize(self, opt):
+        assert(not opt.isTrain)
         BaseModel.initialize(self, opt)
+        self.input_A = self.Tensor(opt.batchSize, opt.input_nc, opt.fineSize, opt.fineSize)
 
-        nb = opt.batchSize
-        size = opt.fineSize
-        self.input_A = self.Tensor(nb, opt.input_nc, size, size)
-
-        assert(not self.isTrain)
-        self.netG_A = networks.define_G(opt.input_nc, opt.output_nc,
-                                        opt.ngf, opt.which_model_netG,
-                                        opt.norm, opt.use_dropout,
-                                        self.gpu_ids)
+        self.netG = networks.define_G(opt.input_nc, opt.output_nc,
+                                      opt.ngf, opt.which_model_netG,
+                                      opt.norm, opt.use_dropout,
+                                      self.gpu_ids)
         which_epoch = opt.which_epoch
-        #AtoB = self.opt.which_direction == 'AtoB'
-        #which_network = 'G_A' if AtoB else 'G_B'
-        self.load_network(self.netG_A, 'G', which_epoch)
+        self.load_network(self.netG, 'G', which_epoch)
 
         print('---------- Networks initialized -------------')
-        networks.print_network(self.netG_A)
+        networks.print_network(self.netG)
         print('-----------------------------------------------')
 
     def set_input(self, input):
-        AtoB = self.opt.which_direction == 'AtoB'
-        input_A = input['A' if AtoB else 'B']
+        # we need to use single_dataset mode
+        input_A = input['A']
         self.input_A.resize_(input_A.size()).copy_(input_A)
-        self.image_paths = input['A_paths' if AtoB else 'B_paths']
+        self.image_paths = input['A_paths']
 
     def test(self):
         self.real_A = Variable(self.input_A)
-        self.fake_B = self.netG_A.forward(self.real_A)
+        self.fake_B = self.netG.forward(self.real_A)
 
-    #get image paths
+    # get image paths
     def get_image_paths(self):
         return self.image_paths
 
diff --git a/scripts/test_single.sh b/scripts/test_single.sh
new file mode 100644
index 0000000..6157b29
--- /dev/null
+++ b/scripts/test_single.sh
@@ -0,0 +1 @@
+python test.py --dataroot ./datasets/facades/testB/ --name facades_pix2pix --model test --which_model_netG unet_256 --which_direction BtoA --dataset_mode single --use_dropout
diff --git a/test.py b/test.py
index d2910ab..f019d10 100644
--- a/test.py
+++ b/test.py
@@ -1,18 +1,17 @@
 import time
 import os
 from options.test_options import TestOptions
-opt = TestOptions().parse()  # set CUDA_VISIBLE_DEVICES before import torch
-
 from data.data_loader import CreateDataLoader
 from models.models import create_model
 from util.visualizer import Visualizer
 from pdb import set_trace as st
 from util import html
 
-opt.nThreads = 1   # test code only supports nThreads=1
-opt.batchSize = 1  #test code only supports batchSize=1
-opt.serial_batches = True # no shuffle
-opt.no_flip = True
+opt = TestOptions().parse()
+opt.nThreads = 1   # test code only supports nThreads = 1
+opt.batchSize = 1  # test code only supports batchSize = 1
+opt.serial_batches = True  # no shuffle
+opt.no_flip = True  # no flip
 
 data_loader = CreateDataLoader(opt)
 dataset = data_loader.load_data()
diff --git a/train.py b/train.py
index 12c9dbc..0834d37 100644
--- a/train.py
+++ b/train.py
@@ -1,11 +1,10 @@
 import time
 from options.train_options import TrainOptions
-opt = TrainOptions().parse()  # set CUDA_VISIBLE_DEVICES before import torch
-
 from data.data_loader import CreateDataLoader
 from models.models import create_model
 from util.visualizer import Visualizer
 
+opt = TrainOptions().parse()
 data_loader = CreateDataLoader(opt)
 dataset = data_loader.load_data()
 dataset_size = len(data_loader)