Bug fixing & add cuda options (NVIDIA#28)

* Fix an incorrect number.

* Add CUDA options

Author: suquark

demo.py

@@ -19,6 +19,7 @@
 parser.add_argument('--style_image_path', default='./images/style1.png')
 parser.add_argument('--style_seg_path', default=[])
 parser.add_argument('--output_image_path', default='./results/example1.png')
+parser.add_argument('--cuda', type=bool, default=True, help='Enable CUDA.')
 args = parser.parse_args()
 
 # Load model
@@ -28,7 +29,7 @@
 except:
     print("Fail to load PhotoWCT models. PhotoWCT submodule not updated?")
     exit()
-    
+
 p_wct.cuda(0)
 
 process_stylization.stylization(
@@ -38,4 +39,5 @@
     content_seg_path=args.content_seg_path,
     style_seg_path=args.style_seg_path,
     output_image_path=args.output_image_path,
+    cuda=args.cuda,
 )

models.py

@@ -163,7 +163,7 @@ def forward_multiple(self, x):
         out = self.conv3_1(out)
         out = self.relu3_1(out)
 
-        if self.level < 3: return out, out2, out1
+        if self.level < 4: return out, out2, out1
 
         out3 = out
 

process_stylization.py

@@ -32,7 +32,8 @@ def __exit__(self, exc_type, exc_value, exc_tb):
         print(self.msg % (time.time() - self.start_time))
 
 
-def stylization(p_wct, content_image_path, style_image_path, content_seg_path, style_seg_path, output_image_path):
+def stylization(p_wct, content_image_path, style_image_path, content_seg_path, style_seg_path, output_image_path,
+                cuda):
     # Load image
     cont_img = Image.open(content_image_path).convert('RGB')
     styl_img = Image.open(style_image_path).convert('RGB')
@@ -45,8 +46,14 @@ def stylization(p_wct, content_image_path, style_image_path, content_seg_path, s
 
     cont_img = transforms.ToTensor()(cont_img).unsqueeze(0)
     styl_img = transforms.ToTensor()(styl_img).unsqueeze(0)
-    cont_img = Variable(cont_img.cuda(0), volatile=True)
-    styl_img = Variable(styl_img.cuda(0), volatile=True)
+    
+    if cuda:
+        cont_img = cont_img.cuda(0)
+        styl_img = styl_img.cuda(0)
+        p_wct.cuda(0)
+    
+    cont_img = Variable(cont_img, volatile=True)
+    styl_img = Variable(styl_img, volatile=True)
 
     cont_seg = np.asarray(cont_seg)
     styl_seg = np.asarray(styl_seg)
@@ -59,6 +66,10 @@ def stylization(p_wct, content_image_path, style_image_path, content_seg_path, s
         out_img = p_pro.process(output_image_path, content_image_path)
     out_img.save(output_image_path)
 
+    if not cuda:
+        print("NotImplemented: The CPU version of smooth filter has not been implemented currently.")
+        return
+    
     with Timer("Elapsed time in post processing: %f"):
         out_img = smooth_filter(output_image_path, content_image_path, f_radius=15, f_edge=1e-1)
     out_img.save(output_image_path)

process_stylization_examples.py

@@ -15,6 +15,7 @@
 parser = argparse.ArgumentParser(description='Photorealistic Image Stylization')
 parser.add_argument('--model', default='./PhotoWCTModels/photo_wct.pth',
                     help='Path to the PhotoWCT model. These are provided by the PhotoWCT submodule, please use `git submodule update --init --recursive` to pull.')
+parser.add_argument('--cuda', type=bool, default=True, help='Enable CUDA.')
 args = parser.parse_args()
 
 folder = 'examples'
@@ -29,7 +30,6 @@
 # Load model
 p_wct = PhotoWCT()
 p_wct.load_state_dict(torch.load(args.model))
-p_wct.cuda(0)
 
 for f in cont_img_list:
     print("Process " + f)
@@ -47,4 +47,5 @@
         content_seg_path=content_seg_path,
         style_seg_path=style_seg_path,
         output_image_path=output_image_path,
+        cuda=args.cuda,
     )