added tf version to colab master notebook

colab_demo.ipynb

@@ -357,16 +357,6 @@
     "!git clone --recursive https://github.com/shaoanlu/fewshot-face-translation-GAN.git"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# There are import errors under keras == 2.2.5\n",
-    "!pip install keras==2.2.4"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": 3,
@@ -422,23 +412,9 @@
    ],
    "source": [
     "# Download pre-trined weights\n",
-    "!gdown https://drive.google.com/uc?id=1DUMmZGTGKMyEYSKy-w34IDHawVF24rIs\n",
-    "!gdown https://drive.google.com/uc?id=1xl8cg7xaRnMsyiODcXguJ83d5hwodckB"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 0,
-   "metadata": {
-    "colab": {},
-    "colab_type": "code",
-    "id": "u9b2N8INnm8P"
-   },
-   "outputs": [],
-   "source": [
     "!mkdir weights\n",
-    "!mv decoder.h5 weights/decoder.h5\n",
-    "!mv encoder.h5 weights/encoder.h5"
+    "!gdown --id 1DUMmZGTGKMyEYSKy-w34IDHawVF24rIs -O weights/encoder.h5\n",
+    "!gdown --id 1xl8cg7xaRnMsyiODcXguJ83d5hwodckB -O weights/decoder.h5"
    ]
   },
   {
@@ -897,17 +873,6 @@
     "result_img = utils.post_process_result(fn_src, fd, result_face, aligned_im, src, x0, y0, x1, y1, landmarks)\n",
     "plt.imshow(result_img)"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 0,
-   "metadata": {
-    "colab": {},
-    "colab_type": "code",
-    "id": "_Do9x9wZl0BK"
-   },
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
@@ -938,4 +903,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 1
-}
+}

models.py

@@ -1,8 +1,8 @@
 import cv2
 import os
 import numpy as np
-from keras.layers import Input
-from keras import backend as K
+from tensorflow.keras.layers import Input
+from tensorflow.keras import backend as K
 
 import networks.generator as gen
 
@@ -57,4 +57,3 @@ def inference(self, src, mask, tar, emb_tar):
                 self.preprocess_input(mask.astype(np.uint8))[None, ...],
                 emb_tar
             ])
-

networks/discriminator.py

@@ -1,45 +1,41 @@
-from keras.models import Model
-from keras.layers import *
-from keras.layers.advanced_activations import LeakyReLU
+from tensorflow.keras.models import Model
+from tensorflow.keras.layers import Input, Concatenate
 
 from .nn_blocks import *
 
-def discriminator_perceptually_aware(nc_in, input_size=IMAGE_SHAPE[0], vggface_res50=None):    
+def discriminator_perceptually_aware(nc_in, input_size, vggface_res50=None):    
     inp1 = Input(shape=(input_size, input_size, nc_in))
     inp2 = Input(shape=(input_size, input_size, nc_in))   
 
     assert not (vggface_res50 == None), "Perceptual model not found."
     perceptual_feats = vggface_res50(inp1)
-    pf1_0 = resize_tensor(perceptual_feats[0], [input_size//4, input_size//4])
     pf1_1 = resize_tensor(perceptual_feats[1], [input_size//8, input_size//8])
     pf1_2 = resize_tensor(perceptual_feats[2], [input_size//16, input_size//16])
     pf1_3 = resize_tensor(perceptual_feats[3], [input_size//32, input_size//32])
     perceptual_feats = vggface_res50(inp2)
-    pf2_0 = resize_tensor(perceptual_feats[0], [input_size//4, input_size//4])
     pf2_1 = resize_tensor(perceptual_feats[1], [input_size//8, input_size//8])
     pf2_2 = resize_tensor(perceptual_feats[2], [input_size//16, input_size//16])
     pf2_3 = resize_tensor(perceptual_feats[3], [input_size//32, input_size//32])
 
-    x = concatenate([inp1, inp2])
+    x = Concatenate()([inp1, inp2])
     nc_base = 32
     for i in range(3):
         x = conv_block_d(x, int(nc_base*(2**(i))), False)
     x = res_block(x, 128)
-    x = concatenate([x, conv_block_d(pf1_1, f=128, k=1, strides=1), conv_block_d(pf2_1, f=128, k=1, strides=1)])
+    x = Concatenate()([x, conv_block_d(pf1_1, f=128, k=1, strides=1), conv_block_d(pf2_1, f=128, k=1, strides=1)])
     x = conv_block_d(x, 256, False)
     x = res_block(x, 256)
-    x = concatenate([x, conv_block_d(pf1_2, f=128, k=1, strides=1), conv_block_d(pf2_2, f=256, k=1, strides=1)])
+    x = Concatenate()([x, conv_block_d(pf1_2, f=128, k=1, strides=1), conv_block_d(pf2_2, f=256, k=1, strides=1)])
     x = conv_block_d(x, 512, False)
-    x = concatenate([x, conv_block_d(pf1_3, f=128, k=1, strides=1), conv_block_d(pf2_3, f=256, k=1, strides=1)])
+    x = Concatenate()([x, conv_block_d(pf1_3, f=128, k=1, strides=1), conv_block_d(pf2_3, f=256, k=1, strides=1)])
     out = Conv2D(1, kernel_size=4, use_bias=False, padding="same")(x)   
     return Model(inputs=[inp1, inp2], outputs=out) 
 
 def discriminator(nc_in, input_size=224):    
     inp = Input(shape=(input_size, input_size, nc_in))
     inp_segm = Input(shape=(input_size, input_size, nc_in))
-    x = concatenate([inp, inp_segm])
+    x = Concatenate()([inp, inp_segm])
 
-    nc_base = 64
     x = conv_block_d(x, 32, False)
     x = conv_block_d(x, 32, False)
     x = conv_block_d(x, 64, True)

networks/generator.py

@@ -1,14 +1,13 @@
-from keras.models import Model
-from keras.layers import *
-from keras.applications import *
+from tensorflow.keras.models import Model
+from tensorflow.keras.layers import Input, Concatenate, Reshape
 
 from .nn_blocks import *
 
 def encoder(nc_in=3, input_size=224):
     inp = Input(shape=(input_size, input_size, nc_in))
     inp_ref = Input(shape=(input_size, input_size, nc_in))
     inp_segm_mask = Input(shape=(input_size, input_size, nc_in))
-    x = concatenate([inp, inp_ref, inp_segm_mask])
+    x = Concatenate()([inp, inp_ref, inp_segm_mask])
     conv1 = conv_block(x, 64, use_norm=True, strides=2)
     conv2 = conv_block(conv1, 128, use_norm=True, strides=2)
     conv3 = conv_block(conv2, 256, use_norm=True, strides=2)
@@ -25,28 +24,27 @@ def encoder(nc_in=3, input_size=224):
 
 def decoder(nc_conv_in=512, input_size=14, nc_in=3, latent_dim=512):
     conv4 = Input(shape=(input_size, input_size, nc_conv_in))
-    inp_segm_mask = Input(shape=(input_size, input_size, nc_in))
+    inp_segm_mask = Input(shape=(input_size*16, input_size*16, nc_in))
     inp_ds2 = Input(shape=(input_size*8, input_size*8, nc_in))
     inp_ds4 = Input(shape=(input_size*4, input_size*4, nc_in))
     inp_emb = Input(shape=(latent_dim,))
-    
+
     emb_mean_var = embddding_fc_block(inp_emb)
     emb_mean_var = Reshape((1, 1, 256))(emb_mean_var)
     emb = Reshape((1, 1, latent_dim))(inp_emb)
     emb = UpSampling2D((input_size,input_size))(emb)
-    
+
     x = adain_resblock(conv4, emb_mean_var, 512) 
     x = adain_resblock(x, emb_mean_var, 512)   
-    x = concatenate([x, emb])
+    x = Concatenate()([x, emb])
     x = upscale_nn(x, 512)
     x = SPADE_res_block(x, resize_tensor(inp_segm_mask, [input_size*2, input_size*2]), 512, block_id='3')
     x = upscale_nn(x, 256)
     x = SPADE_res_block(x, resize_tensor(inp_segm_mask, [input_size*4, input_size*4]), 256, block_id='4')
-    x = concatenate([x, inp_ds4])
+    x = Concatenate()([x, inp_ds4])
     x = upscale_nn(x, 128)
-    x = concatenate([x, inp_ds2])
+    x = Concatenate()([x, inp_ds2])
     x = upscale_nn(x, 64)
-
-    out = Conv2D(3, kernel_size=4, padding='same', activation="tanh")(x)
-    return Model([conv4, inp_ds2, inp_ds4, inp_segm_mask, inp_emb], out)  
 
+    out = Conv2D(3, kernel_size=4, padding='same', activation="tanh")(x)
+    return Model([conv4, inp_ds2, inp_ds4, inp_segm_mask, inp_emb], out)

networks/instance_normalization.py

@@ -1,9 +1,7 @@
-from keras.engine import Layer, InputSpec
-from keras import initializers, regularizers, constraints
-from keras import backend as K
-from keras.utils.generic_utils import get_custom_objects
-
-import numpy as np
+from tensorflow.keras.layers import Layer, InputSpec
+from tensorflow.keras import initializers, regularizers, constraints
+from tensorflow.keras import backend as K
+from tensorflow.keras.utils import get_custom_objects
 
 
 class InstanceNormalization(Layer):

networks/nn_blocks.py

@@ -1,5 +1,6 @@
-from keras.layers import *
-from keras.layers.advanced_activations import LeakyReLU
+from tensorflow.keras.layers import Lambda, BatchNormalization, Conv2D, Dense, LeakyReLU, ReLU, Add, UpSampling2D
+from tensorflow.keras import regularizers
+from tensorflow.keras.layers.experimental.preprocessing import Resizing
 import tensorflow as tf
 
 from networks.instance_normalization import InstanceNormalization
@@ -22,25 +23,25 @@ def normalization(inp, norm='none', group='16'):
     #    def spade_norm(x):
     #        meanC, varC = tf.nn.moments(x, [1, 2], keep_dims=True)
     #        sigmaC = tf.sqrt(tf.add(varC, 1e-5))
-    #        return (content - meanC) / sigmaC
+    #        return (x - meanC) / sigmaC
     #    x = Lambda(lambda xx: spade_norm(xx))(x)
     else:
         x = x
     return x
 
 def conv_block(input_tensor, f, use_norm=False, k=3, strides=2):
     x = input_tensor
-    if not k == 1:
+    if k != 1:
         x = ReflectPadding2D(x)
     x = Conv2D(f, kernel_size=k, strides=strides, kernel_regularizer=regularizers.l2(w_l2),  
                kernel_initializer=conv_init, use_bias=(not use_norm))(x)
     x = normalization(x, norm, f) if use_norm else x
-    x = Activation("relu")(x)
+    x = ReLU()(x)
     return x
 
 def conv_block_d(input_tensor, f, use_norm=False, k=3, strides=2):
     x = input_tensor
-    if not k == 1:
+    if k != 1:
         x = ReflectPadding2D(x)
     x = Conv2D(f, kernel_size=k, strides=strides, kernel_regularizer=regularizers.l2(w_l2), 
                kernel_initializer=conv_init, use_bias=(not use_norm))(x)
@@ -55,32 +56,32 @@ def res_block(input_tensor, f, use_norm=True):
     x = Conv2D(f, kernel_size=3, kernel_regularizer=regularizers.l2(w_l2), 
                kernel_initializer=conv_init, use_bias=not use_norm)(x)
     x = normalization(x, norm, f) if use_norm else x
-    x = Activation('relu')(x)
+    x = ReLU()(x)
     x = ReflectPadding2D(x)
     x = Conv2D(f, kernel_size=3, kernel_regularizer=regularizers.l2(w_l2), 
                kernel_initializer=conv_init)(x)
-    x = add([x, input_tensor])
-    x = Activation('relu')(x)
+    x = Add()([x, input_tensor])
+    x = ReLU()(x)
     return x
 
 def embddding_fc_block(input_tensor):
     x = input_tensor
 
     x = Dense(256, kernel_regularizer=regularizers.l2(w_l2))(x)
     x = normalization(x, norm, 256)
-    x = Activation('relu')(x)
+    x = ReLU()(x)
     x = Dense(256, kernel_regularizer=regularizers.l2(w_l2))(x)
     x = normalization(x, norm, 256)
-    x = Activation('relu')(x)
+    x = ReLU()(x)
     x = Dense(256, kernel_regularizer=regularizers.l2(w_l2))(x)
     x = normalization(x, norm, 256)
-    x = Activation('relu')(x)
+    x = ReLU()(x)
     return x
 
 def adain_resblock(input_tensor, embeddings, f):
     # conv -> norm -> activ -> conv -> norm -> add
     def AdaIN(content, style_var, style_mean, epsilon=1e-5):
-        meanC, varC = tf.nn.moments(content, [1, 2], keep_dims=True)
+        meanC, varC = tf.nn.moments(content, [1, 2], keepdims=True)
         sigmaC = tf.sqrt(tf.add(varC, epsilon))
         return (content - meanC) * style_var / sigmaC + style_mean
 
@@ -92,37 +93,37 @@ def AdaIN(content, style_var, style_mean, epsilon=1e-5):
     x = Conv2D(f, kernel_size=3, kernel_regularizer=regularizers.l2(w_l2), 
                kernel_initializer=conv_init, use_bias=False)(x)
     x = Lambda(lambda x: AdaIN(x[0], x[1], x[2]))([x, style_var, style_mean])
-    x = Activation('relu')(x)
+    x = ReLU()(x)
     x = ReflectPadding2D(x)
     x = Conv2D(f, kernel_size=3, kernel_regularizer=regularizers.l2(w_l2), 
                kernel_initializer=conv_init, use_bias=False)(x)
     x = Lambda(lambda x: AdaIN(x[0], x[1], x[2]))([x, style_var, style_mean])
-    x = add([x, input_tensor])
-    x = Activation('relu')(x)
+    x = Add()([x, input_tensor])
+    x = ReLU()(x)
     return x
 
 def SPADE(input_tensor, cond_input_tensor, f, block_id=0):
     x = input_tensor
     x = InstanceNormalization(name=f"SPADE_norm{block_id}")(x)# x = normalization(x, "SPADE_norm", f)
     y = cond_input_tensor
     y = Conv2D(128, kernel_size=3, padding='same')(y)
-    y = Activation('relu')(y)           
+    y = ReLU()(y)           
     gamma = Conv2D(f, kernel_size=3, padding='same')(y)
     beta = Conv2D(f, kernel_size=3, padding='same')(y)    
-    x = add([x, multiply([x, gamma])])
-    x = add([x, beta])
+    x = Add()([x, tf.multiply(x, gamma)])
+    x = Add()([x, beta])
     return x
 
 def SPADE_res_block(input_tensor, cond_input_tensor, f, block_id=0):
     x = input_tensor
     x = SPADE(x, cond_input_tensor, f, block_id=block_id+"_0")
-    x = Activation('relu')(x)
+    x = ReLU()(x)
     x = Conv2D(f, kernel_size=3, kernel_initializer=conv_init, padding='same')(x)
     x = SPADE(x, cond_input_tensor, f, block_id=block_id+"_1")
-    x = Activation('relu')(x)
+    x = ReLU()(x)
     x = Conv2D(f, kernel_size=3, kernel_initializer=conv_init, padding='same')(x)
-    x = add([x, input_tensor])
-    x = Activation('relu')(x)
+    x = Add()([x, input_tensor])
+    x = ReLU()(x)
     return x
 
 def upscale_nn(input_tensor, f, use_norm=True, w_l2=w_l2, norm=norm):
@@ -135,7 +136,4 @@ def upscale_nn(input_tensor, f, use_norm=True, w_l2=w_l2, norm=norm):
     return x
 
 def resize_tensor(inp, shape):
-    if isinstance(shape, int):
-        return Lambda(lambda x: tf.image.resize_images(x, [shape, shape]))(inp)
-    elif isinstance(shape, list):
-        return Lambda(lambda x: tf.image.resize_images(x, [shape[0], shape[1]]))(inp)
+    return Resizing(*shape)(inp)

utils/utils.py

@@ -27,7 +27,7 @@ def get_tar_landmarks(img, landmarks_type=68):
              [0.81675393, 0.52105263], [0.81675393, 0.46315789], [0.81675393, 0.38421053]]
         )
     else:
-        raise NotImplementedError(f"Only 68 points landmarks model is provided. Received {landmarks_pnts}.")
+        raise NotImplementedError(f"Only 68 points landmarks model is provided. Received {landmarks_type}.")
     tar_landmarks = [(int(xy[0]*img_sz[0]), int(xy[1]*img_sz[1])) for xy in avg_landmarks]
     return tar_landmarks
 
@@ -47,7 +47,7 @@ def landmarks_match(src_im, src_landmarks, tar_landmarks, border_mode=cv2.BORDER
     result = cv2.warpAffine(src_im, M, (src_size[1], src_size[0]), borderMode=border_mode, borderValue=border_value)
     return result, M     
 
-def get_68_landmarks_edge_image(img, face, lms, eyes_binary_mask=False, apply_dilate_or_erode=True):
+def get_68_landmarks_edge_image(img, face, lms):
     result = np.zeros_like(img, np.uint8)
     stroke_size = np.maximum(np.min(face.shape[:2])//50, 2)    
 
@@ -125,8 +125,7 @@ def get_segm_mask(im, face_im, x0, y0, x1, y1, landmarks):
     seg_mask = get_68_landmarks_edge_image(
         im, 
         face_im, 
-        landmarks[0],
-        apply_dilate_or_erode=False)
+        landmarks[0])
     seg_mask = seg_mask[int(x0):int(x1), int(y0):int(y1), :] 
     return seg_mask
 
@@ -203,8 +202,8 @@ def get_tar_inputs(fns, fd, fv):
         aligned_face: A RGB image.
         emb_tar: A numpy array of shape (512,). Latent embeddings of aligned_face.
     """
-    if not type(fns) == list:
-        if type(fns) == str:
+    if not isinstance(fns, list):
+        if isinstance(fns, str):
             fns = [fns]
         else:
             raise ValueError("Received and unknown filename type. fns shoulbe be a list or a string.")