Merge branch 'master' of https://github.com/cysmith/neural-style-tf

1 files changed, 13 insertions, 12 deletions

diff --git a/neural_style.py b/neural_style.py
index 62875f6..8d40e1b 100644
--- a/neural_style.py
+++ b/neural_style.py
@@ -479,9 +479,9 @@ def sum_total_variation_losses(sess, net, input_img):
   x = net['input']
   tv_y_size = b * (h-1) * w * d
   tv_x_size = b * h * (w-1) * d
-  loss_y = tf.nn.l2_loss(x[:,1:,:,:] - x[:,:h-1,:,:]) 
+  loss_y = tf.nn.l2_loss(x[:,1:,:,:] - x[:,:-1,:,:]) 
   loss_y /= tv_y_size
-  loss_x = tf.nn.l2_loss(x[:,:,1:,:] - x[:,:,:w-1,:]) 
+  loss_x = tf.nn.l2_loss(x[:,:,1:,:] - x[:,:,:-1,:]) 
   loss_x /= tv_x_size
   loss = 2 * (loss_y + loss_x)
   loss = tf.cast(loss, tf.float32)
@@ -535,14 +535,14 @@ def read_flow_file(path):
 
 def read_weights_file(path):
   lines = open(path).readlines()
-  header = map(int, lines[0].split(' '))
+  header = list(map(int, lines[0].split(' ')))
   w = header[0]
   h = header[1]
   vals = np.zeros((h, w), dtype=np.float32)
   for i in range(1, len(lines)):
     line = lines[i].rstrip().split(' ')
-    vals[i-1] = np.array(map(np.float32, line))
-    vals[i-1] = map(lambda x: 0. if x < 255. else 1., vals[i-1])
+    vals[i-1] = np.array(list(map(np.float32, line)))
+    vals[i-1] = list(map(lambda x: 0. if x < 255. else 1., vals[i-1]))
   # expand to 3 channels
   weights = np.dstack([vals.astype(np.float32)] * 3)
   return weights
@@ -551,7 +551,7 @@ def normalize(weights):
   denom = sum(weights)
   if denom > 0.:
     return [float(i) / denom for i in weights]
-  else: return [0. for _ in weights]
+  else: return [0.] * len(weights)
 
 def maybe_make_directory(dir_path):
   if not os.path.exists(dir_path):  
@@ -568,29 +568,30 @@ def stylize(content_img, style_imgs, init_img, frame=None):
   with tf.device(args.device), tf.Session() as sess:
     # setup network
     net = build_vgg19(content_img)
-
+    
     # style loss
     if args.style_mask:
       L_style = sum_masked_style_losses(sess, net, style_imgs)
     else:
       L_style = sum_style_losses(sess, net, style_imgs)
-
+    
     # content loss
     L_content = sum_content_losses(sess, net, content_img)
-
+    
     # denoising loss
     L_tv = sum_total_variation_losses(sess, net, init_img)
-
+    
     # loss weights
     alpha = args.content_weight
     beta  = args.style_weight
-    theta = args.tv_weight    
+    theta = args.tv_weight
     
     # total loss
     L_total  = alpha * L_content
     L_total += beta  * L_style
     L_total += theta * L_tv
-
+    
+    # video temporal loss
     if args.video and frame > 1:
       gamma      = args.temporal_weight
       L_temporal = sum_shortterm_temporal_losses(sess, net, frame, init_img)