code refactoring

Author: junyanz

constrained_opt.py

@@ -5,7 +5,7 @@
 import sys
 from lib import utils
 from PyQt4.QtCore import *
-import cv2
+
 
 class Constrained_OPT(QThread):
     def __init__(self, opt_solver, batch_size=32, n_iters=25, topK=16, morph_steps=16, interp='linear'):
@@ -35,7 +35,7 @@ def __init__(self, opt_solver, batch_size=32, n_iters=25, topK=16, morph_steps=1
         self.order = None
         self.init_constraints()  # initialize
         self.init_z()            # initialize latent vectors
-        self.just_fixed=True
+        self.just_fixed = True
         self.weights = None
 
     def is_fixed(self):
@@ -55,7 +55,7 @@ def init_z(self, frame_id=-1, image_id=-1):
             print('set z as image %d, frame %d' % (image_id, frame_id))
             self.prev_z = self.z_seq[image_id, frame_id]
 
-        if self.prev_z is None:  #random initialization
+        if self.prev_z is None:  # random initialization
             self.z_init = np_rng.uniform(-1.0, 1.0, size=(self.batch_size, nz))
             self.opt_solver.set_smoothness(0.0)
             self.z_const = self.z_init
@@ -79,25 +79,19 @@ def update(self):   # update ui
 
     def save_constraints(self):
         [im_c, mask_c, im_e, mask_e] = self.combine_constraints(self.constraints)
-        # write image
-        # im_c2 = cv2.cvtColor(im_c, cv2.COLOR_RGB2BGR)
-        # cv2.imwrite('input_color_image.png', im_c2)
-        # cv2.imwrite('input_color_mask.png', mask_c)
-        # cv2.imwrite('input_edge_map.png', im_e)
         self.prev_im_c = im_c.copy()
         self.prev_mask_c = mask_c.copy()
         self.prev_im_e = im_e.copy()
-        self.prev_mask_e =mask_e.copy()
+        self.prev_mask_e = mask_e.copy()
 
     def init_constraints(self):
         self.prev_im_c = None
         self.prev_mask_c = None
         self.prev_im_e = None
         self.prev_mask_e = None
 
-
     def combine_constraints(self, constraints):
-        if constraints is not None: #[hack]
+        if constraints is not None:  # [hack]
             # print('combine strokes')
             [im_c, mask_c, im_e, mask_e] = constraints
             if self.prev_im_c is None:
@@ -113,22 +107,21 @@ def combine_constraints(self, constraints):
             if self.prev_im_c is None:
                 im_c_f = im_c
             else:
-                im_c_f =  self.prev_im_c.copy()
-                mask_c3 = np.tile(mask_c, [1,1, im_c.shape[2]])
-                np.copyto(im_c_f, im_c, where=mask_c3.astype(np.bool))  #[hack]
+                im_c_f = self.prev_im_c.copy()
+                mask_c3 = np.tile(mask_c, [1, 1, im_c.shape[2]])
+                np.copyto(im_c_f, im_c, where=mask_c3.astype(np.bool))  # [hack]
 
             if self.prev_im_e is None:
                 im_e_f = im_e
             else:
                 im_e_f = self.prev_im_e.copy()
-                mask_e3 = np.tile(mask_e, [1,1,im_e.shape[2]])
+                mask_e3 = np.tile(mask_e, [1, 1, im_e.shape[2]])
                 np.copyto(im_e_f, im_e, where=mask_e3.astype(np.bool))
 
             return [im_c_f, mask_c_f, im_e_f, mask_e_f]
         else:
             return [self.prev_im_c, self.prev_mask_c, self.prev_im_e, self.prev_mask_e]
 
-
     def set_constraints(self, constraints):
         self.constraints = constraints
 
@@ -140,7 +133,6 @@ def get_z(self, image_id, frame_id):
         else:
             return None
 
-
     def get_image(self, image_id, frame_id, useAverage=False):
         if self.to_update:
             if self.current_ims is None or self.current_ims.size == 0:
@@ -158,7 +150,7 @@ def get_image(self, image_id, frame_id, useAverage=False):
                 frame_id = frame_id % self.img_seq.shape[1]
                 image_id = image_id % self.img_seq.shape[0]
                 if useAverage and self.weights is not None:
-                    return utils.average_image(self.img_seq[:,frame_id,...], self.weights)
+                    return utils.average_image(self.img_seq[:, frame_id, ...], self.weights)
                 else:
                     return self.img_seq[image_id, frame_id]
 
@@ -188,10 +180,10 @@ def get_current_results(self):
         return self.current_ims
 
     def run(self):  # main function
-        time_to_wait = 33 # 33 millisecond
+        time_to_wait = 33  # 33 millisecond
         while (1):
-            t1 =time()
-            if self.to_set_constraints:# update constraints
+            t1 = time()
+            if self.to_set_constraints:  # update constraints
                 self.to_set_constraints = False
 
             if self.constraints is not None and self.iter_count < self.max_iters:
@@ -204,11 +196,11 @@ def run(self):  # main function
                 self.to_update = False
                 self.iter_count += 1
 
-            t_c = int(1000*(time()-t1))
+            t_c = int(1000 * (time() - t1))
             print('update one iteration: %03d ms' % t_c, end='\r')
             sys.stdout.flush()
             if t_c < time_to_wait:
-                self.msleep(time_to_wait-t_c)
+                self.msleep(time_to_wait - t_c)
 
     def update_invert(self, constraints):
         constraints_c = self.combine_constraints(constraints)
@@ -218,7 +210,7 @@ def update_invert(self, constraints):
 
         if self.topK > 1:
             cost_sort = cost_all[order]
-            thres_top =  2 * np.mean(cost_sort[0:min(int(self.topK / 2.0), len(cost_sort))])
+            thres_top = 2 * np.mean(cost_sort[0:min(int(self.topK / 2.0), len(cost_sort))])
             ids = cost_sort - thres_top < 1e-10
             topK = np.min([self.topK, sum(ids)])
         else:
@@ -233,7 +225,7 @@ def update_invert(self, constraints):
         self.current_ims = gx_t[order]
         # compute weights
         cost_weights = cost_all[order]
-        self.weights = np.exp(-(cost_weights-np.mean(cost_weights)) / (np.std(cost_weights)+1e-10))
+        self.weights = np.exp(-(cost_weights - np.mean(cost_weights)) / (np.std(cost_weights) + 1e-10))
         self.current_zs = z_t[order]
         self.emit(SIGNAL('update_image'))
 
@@ -248,14 +240,14 @@ def gen_morphing(self, interp='linear', n_steps=8):
         z_seq = []
 
         for n in range(n_steps):
-            ratio = n / float(n_steps- 1)
+            ratio = n / float(n_steps - 1)
             z_t = utils.interp_z(z1, z2, ratio, interp=interp)
             seq = self.opt_solver.gen_samples(z0=z_t)
             img_seq.append(seq[:, np.newaxis, ...])
-            z_seq.append(z_t[:,np.newaxis,...])
+            z_seq.append(z_t[:, np.newaxis, ...])
         self.img_seq = np.concatenate(img_seq, axis=1)
         self.z_seq = np.concatenate(z_seq, axis=1)
-        print('generate morphing sequence (%.3f seconds)' % (time()-t))
+        print('generate morphing sequence (%.3f seconds)' % (time() - t))
 
     def reset(self):
         self.prev_z = self.z0
@@ -271,7 +263,4 @@ def reset(self):
         self.to_set_constraints = False
         self.iter_total = 0
         self.iter_count = 0
-        self.weights =None
-
-
-
+        self.weights = None

constrained_opt_theano.py

@@ -6,6 +6,7 @@
 from lib.theano_utils import floatX, sharedX
 import numpy as np
 
+
 class OPT_Solver():
     def __init__(self, model, batch_size=32, d_weight=0.0):
         self.model = model
@@ -16,7 +17,7 @@ def __init__(self, model, batch_size=32, d_weight=0.0):
         self.transform = model.transform
         self.transform_mask = model.transform_mask
         self.inverse_transform = model.inverse_transform
-        BS = 4 if self.nc == 1 else 8 # [hack]
+        BS = 4 if self.nc == 1 else 8  # [hack]
         self.hog = HOGNet.HOGNet(use_bin=True, NO=8, BS=BS, nc=self.nc)
         self.opt_model = self.def_invert(model, batch_size=batch_size, d_weight=d_weight, nc=self.nc)
         self.batch_size = batch_size
@@ -26,8 +27,8 @@ def get_image_size(self):
 
     def invert(self, constraints, z_i):
         [_invert, z_updates, z, beta_r, z_const] = self.opt_model
-        constraints_t  = self.preprocess_constraints(constraints)
-        [im_c_t, mask_c_t, im_e_t, mask_e_t] = constraints_t # [im_c_t, mask_c_t, im_e_t, mask_e_t]
+        constraints_t = self.preprocess_constraints(constraints)
+        [im_c_t, mask_c_t, im_e_t, mask_e_t] = constraints_t  # [im_c_t, mask_c_t, im_e_t, mask_e_t]
 
         results = _invert(im_c_t, mask_c_t, im_e_t, mask_e_t, z_i.astype(np.float32))
 
@@ -38,7 +39,6 @@ def invert(self, constraints, z_i):
         z_t = np.tanh(z.get_value()).copy()
         return gx_t, z_t, cost_all
 
-
     def preprocess_constraints(self, constraints):
         [im_c_o, mask_c_o, im_e_o, mask_e_o] = constraints
         im_c = self.transform(im_c_o[np.newaxis, :], self.nc)
@@ -65,7 +65,7 @@ def set_smoothness(self, l):
     def gen_samples(self, z0):
         samples = self.model.gen_samples(z0=z0)
         if self.nc == 1:
-            samples = np.tile(samples, [1,1,1,3])
+            samples = np.tile(samples, [1, 1, 1, 3])
         return samples
 
     def def_invert(self, model, batch_size=1, d_weight=0.5, nc=1, lr=0.1, b1=0.9, nz=100, use_bin=True):
@@ -79,8 +79,8 @@ def def_invert(self, model, batch_size=1, d_weight=0.5, nc=1, lr=0.1, b1=0.9, nz
         gx = model.model_G(z)
         # input: im_c: 255: no edge; 0: edge; transform=> 1: no edge, 0: edge
 
-        if nc == 1: # gx, range [0, 1] => edge, 1
-            gx3 = 1.0-gx #T.tile(gx, (1, 3, 1, 1))
+        if nc == 1:  # gx, range [0, 1] => edge, 1
+            gx3 = 1.0 - gx  # T.tile(gx, (1, 3, 1, 1))
         else:
             gx3 = gx
         mm_c = T.tile(m_c, (1, gx3.shape[1], 1, 1))
@@ -116,4 +116,3 @@ def def_invert(self, model, batch_size=1, d_weight=0.5, nc=1, lr=0.1, b1=0.9, nz
         _invert = theano.function(inputs=[x_c, m_c, x_e, m_e, z0], outputs=output, updates=z_updates)
         print('%.2f seconds to compile _invert function' % (time() - t))
         return [_invert, z_updates, z, d_weight_r, z_const]
-

generate_samples.py

@@ -4,6 +4,7 @@
 from lib import utils
 import cv2
 
+
 def parse_args():
     parser = argparse.ArgumentParser(description='generated random samples (dcgan_theano)')
     parser.add_argument('--model_name', dest='model_name', help='the model name', default='outdoor_64', type=str)
@@ -15,9 +16,10 @@ def parse_args():
     args = parser.parse_args()
     return args
 
+
 if __name__ == '__main__':
     args = parse_args()
-    if not args.model_file:  #if model directory is not specified
+    if not args.model_file:  # if model directory is not specified
         args.model_file = './models/%s.%s' % (args.model_name, args.model_type)
 
     if not args.output_image:
@@ -37,4 +39,4 @@ def parse_args():
     im_vis = cv2.cvtColor(im_vis, cv2.COLOR_BGR2RGB)
     cv2.imwrite(args.output_image, im_vis)
     print('samples_shape', samples.shape)
-    print('save image to %s' % args.output_image)
+    print('save image to %s' % args.output_image)

iGAN_main.py

@@ -8,6 +8,7 @@
 from pydoc import locate
 import constrained_opt
 
+
 def parse_args():
     parser = argparse.ArgumentParser(description='iGAN: Interactive Visual Synthesis Powered by GAN')
     parser.add_argument('--model_name', dest='model_name', help='the model name', default='outdoor_64', type=str)
@@ -21,14 +22,15 @@ def parse_args():
     parser.add_argument('--model_file', dest='model_file', help='the file that stores the generative model', type=str, default=None)
     parser.add_argument('--d_weight', dest='d_weight', help='captures the visual realism based on GAN discriminator', type=float, default=0.0)
     parser.add_argument('--interp', dest='interp', help='the interpolation method (linear or slerp)', type=str, default='linear')
-    parser.add_argument('--average', dest='average', help='averageExplorer mode',action="store_true", default=False)
+    parser.add_argument('--average', dest='average', help='averageExplorer mode', action="store_true", default=False)
     parser.add_argument('--shadow', dest='shadow', help='shadowDraw mode', action="store_true", default=False)
     args = parser.parse_args()
     return args
 
+
 if __name__ == '__main__':
     args = parse_args()
-    if not args.model_file:  #if the model_file is not specified
+    if not args.model_file:  # if the model_file is not specified
         args.model_file = './models/%s.%s' % (args.model_name, args.model_type)
 
     for arg in vars(args):
@@ -54,4 +56,4 @@ def parse_args():
     window.setWindowTitle('Interactive GAN')
     window.setWindowFlags(window.windowFlags() & ~Qt.WindowMaximizeButtonHint)   # fix window siz
     window.show()
-    app.exec_()
+    app.exec_()

iGAN_predict.py

@@ -16,6 +16,7 @@
 from pydoc import locate
 from lib import activations
 
+
 def def_feature(layer='conv4', up_scale=4):
     print('COMPILING...')
     t = time()
@@ -55,7 +56,8 @@ def def_bfgs(model_G, layer='conv4', npx=64, alpha=0.002):
     _invert = theano.function(inputs=[z, x, x_f], outputs=output)
 
     print('%.2f seconds to compile _bfgs function' % (time() - t))
-    return _invert,z
+    return _invert, z
+
 
 def def_predict(model_P):
     print('COMPILING...')
@@ -74,8 +76,6 @@ def def_invert_models(gen_model, layer='conv4', alpha=0.002):
     return gen_model, bfgs_model, ftr_model, predict_model
 
 
-
-
 def predict_z(gen_model, _predict, ims, batch_size=32):
     n = ims.shape[0]
     n_gen = 0
@@ -97,28 +97,26 @@ def invert_bfgs_batch(gen_model, invert_model, ftr_model, ims, z_predict=None, n
     zs = []
     recs = []
     fs = []
-    t = time()
     n_imgs = ims.shape[0]
     print('reconstruct %d images using bfgs' % n_imgs)
 
     for n in range(n_imgs):
-        im_n = ims[[n], :, :,:]
+        im_n = ims[[n], :, :, :]
         if z_predict is not None:
-            z0_n = z_predict[[n],...]
+            z0_n = z_predict[[n], ...]
         else:
             z0_n = None
-        gx, z_value, f_value = invert_bfgs(gen_model, invert_model, ftr_model,im=im_n, z_predict=z0_n, npx=npx)
+        gx, z_value, f_value = invert_bfgs(gen_model, invert_model, ftr_model, im=im_n, z_predict=z0_n, npx=npx)
         rec_im = (gx * 255).astype(np.uint8)
-        fs.append(f_value[np.newaxis,...])
-        zs.append(z_value[np.newaxis,...])
+        fs.append(f_value[np.newaxis, ...])
+        zs.append(z_value[np.newaxis, ...])
         recs.append(rec_im)
     recs = np.concatenate(recs, axis=0)
     zs = np.concatenate(zs, axis=0)
     fs = np.concatenate(fs, axis=0)
     return recs, zs, fs
 
 
-
 def invert_bfgs(gen_model, invert_model, ftr_model, im, z_predict=None, npx=64):
     _f, z = invert_model
     nz = gen_model.nz
@@ -131,14 +129,14 @@ def invert_bfgs(gen_model, invert_model, ftr_model, im, z_predict=None, npx=64):
     ftr = ftr_model(im_t)
 
     prob = optimize.minimize(f_bfgs, z_predict, args=(_f, im_t, ftr),
-                             tol=1e-6, jac=True, method='L-BFGS-B', options={'maxiter':200})
+                             tol=1e-6, jac=True, method='L-BFGS-B', options={'maxiter': 200})
     print('n_iters = %3d, f = %.3f' % (prob.nit, prob.fun))
     z_opt = prob.x
     z_opt_n = floatX(z_opt[np.newaxis, :])
     [f_opt, g, gx] = _f(z_opt_n, im_t, ftr)
     gx = gen_model.inverse_transform(gx, npx=npx)
     z_opt = np.tanh(z_opt)
-    return gx, z_opt,f_opt
+    return gx, z_opt, f_opt
 
 
 def f_bfgs(z0, _f, x, x_f):
@@ -181,11 +179,12 @@ def parse_args():
     args = parser.parse_args()
     return args
 
+
 if __name__ == "__main__":
     args = parse_args()
     if not args.model_file:  # if the model file is not specified
         args.model_file = './models/%s.%s' % (args.model_name, args.model_type)
-    if not args.output_image:# if the output image path is not specified
+    if not args.output_image:  # if the output image path is not specified
         args.output_image = args.input_image.replace('.png', '_%s.png' % args.solver)
 
     for arg in vars(args):
@@ -205,10 +204,10 @@ def parse_args():
     im = np.array(im)
     im_pre = im[np.newaxis, :, :, :]
     # run the model
-    rec, _, _  = invert_images_CNN_opt(invert_models, im_pre, solver=args.solver)
+    rec, _, _ = invert_images_CNN_opt(invert_models, im_pre, solver=args.solver)
     rec = np.squeeze(rec)
     rec_im = Image.fromarray(rec)
     # resize the image (input aspect ratio)
     rec_im = rec_im.resize((h, w))
     print('write result to %s' % args.output_image)
-    rec_im.save(args.output_image)
+    rec_im.save(args.output_image)