with private models

.idea/.gitignore

@@ -0,0 +1,3 @@
+# Default ignored files
+/shelf/
+/workspace.xml

.idea/AFFA-face-swap.iml

@@ -0,0 +1,15 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+  <component name="PyDocumentationSettings">
+    <option name="format" value="PLAIN" />
+    <option name="myDocStringFormat" value="Plain" />
+  </component>
+  <component name="TestRunnerService">
+    <option name="PROJECT_TEST_RUNNER" value="pytest" />
+  </component>
+</module>

.idea/inspectionProfiles/Project_Default.xml

@@ -0,0 +1,38 @@
+<component name="InspectionProjectProfileManager">
+  <profile version="1.0">
+    <option name="myName" value="Project Default" />
+    <inspection_tool class="PyPackageRequirementsInspection" enabled="true" level="WARNING" enabled_by_default="true">
+      <option name="ignoredPackages">
+        <value>
+          <list size="3">
+            <item index="0" class="java.lang.String" itemvalue="ipython" />
+            <item index="1" class="java.lang.String" itemvalue="Cython" />
+            <item index="2" class="java.lang.String" itemvalue="tensorflow-gpu" />
+          </list>
+        </value>
+      </option>
+    </inspection_tool>
+    <inspection_tool class="PyPep8Inspection" enabled="true" level="WEAK WARNING" enabled_by_default="true">
+      <option name="ignoredErrors">
+        <list>
+          <option value="E402" />
+        </list>
+      </option>
+    </inspection_tool>
+    <inspection_tool class="PyPep8NamingInspection" enabled="true" level="WEAK WARNING" enabled_by_default="true">
+      <option name="ignoredErrors">
+        <list>
+          <option value="N806" />
+          <option value="N812" />
+        </list>
+      </option>
+    </inspection_tool>
+    <inspection_tool class="PyUnresolvedReferencesInspection" enabled="true" level="WARNING" enabled_by_default="true">
+      <option name="ignoredIdentifiers">
+        <list>
+          <option value="torch.backends.cudnn" />
+        </list>
+      </option>
+    </inspection_tool>
+  </profile>
+</component>

.idea/inspectionProfiles/profiles_settings.xml

@@ -0,0 +1,6 @@
+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="USE_PROJECT_PROFILE" value="false" />
+    <version value="1.0" />
+  </settings>
+</component>

.idea/misc.xml

@@ -0,0 +1,4 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.8 (base)" project-jdk-type="Python SDK" />
+</project>

.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/AFFA-face-swap.iml" filepath="$PROJECT_DIR$/.idea/AFFA-face-swap.iml" />
+    </modules>
+  </component>
+</project>

.idea/vcs.xml

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="VcsDirectoryMappings">
+    <mapping directory="$PROJECT_DIR$" vcs="Git" />
+  </component>
+</project>

app.py

@@ -0,0 +1,105 @@
+import gradio
+
+import tensorflow as tf
+from huggingface_hub import Repository
+repo = Repository()
+
+from utils.utils import norm_crop, estimate_norm, inverse_estimate_norm, transform_landmark_points, get_lm, load_model_internal
+from networks.generator import get_generator
+import numpy as np
+import cv2
+from scipy.ndimage import gaussian_filter
+
+from tensorflow.keras.models import load_model
+from retinaface.models import *
+from options.swap_options import SwapOptions
+
+
+opt = SwapOptions().parse()
+
+
+#gpus = tf.config.experimental.list_physical_devices('GPU')
+#tf.config.set_visible_devices(gpus[opt.device_id], 'GPU')
+
+RetinaFace = load_model(opt.retina_path,
+                        custom_objects={"FPN": FPN,
+                                        "SSH": SSH,
+                                        "BboxHead": BboxHead,
+                                        "LandmarkHead": LandmarkHead,
+                                        "ClassHead": ClassHead})
+ArcFace = load_model(opt.arcface_path)
+
+G = load_model_internal(opt.chkp_dir + opt.log_name + "/gen/", "gen", opt.load)
+
+blend_mask_base = np.zeros(shape=(256, 256, 1))
+blend_mask_base[100:240, 32:224] = 1
+blend_mask_base = gaussian_filter(blend_mask_base, sigma=7)
+
+
+def run_inference(target, source):
+    source = np.array(source)
+    target = np.array(target)
+
+    # Prepare to load video
+    source_a = RetinaFace(np.expand_dims(source, axis=0)).numpy()[0]
+    source_h, source_w, _ = source.shape
+    source_lm = get_lm(source_a, source_w, source_h)
+    source_aligned = norm_crop(source, source_lm, image_size=256)
+    source_z = ArcFace.predict(np.expand_dims(tf.image.resize(source_aligned, [112, 112]) / 255.0, axis=0))
+
+    # read frame
+    im = target
+    im_h, im_w, _ = im.shape
+    im_shape = (im_w, im_h)
+
+    detection_scale = im_w // 640 if im_w > 640 else 1
+
+    faces = RetinaFace(np.expand_dims(cv2.resize(im,
+                                                 (im_w // detection_scale,
+                                                  im_h // detection_scale)), axis=0)).numpy()
+
+    total_img = im / 255.0
+    for annotation in faces:
+        lm_align = np.array([[annotation[4] * im_w, annotation[5] * im_h],
+                             [annotation[6] * im_w, annotation[7] * im_h],
+                             [annotation[8] * im_w, annotation[9] * im_h],
+                             [annotation[10] * im_w, annotation[11] * im_h],
+                             [annotation[12] * im_w, annotation[13] * im_h]],
+                            dtype=np.float32)
+
+        # align the detected face
+        M, pose_index = estimate_norm(lm_align, 256, "arcface", shrink_factor=1.0)
+        im_aligned = cv2.warpAffine(im, M, (256, 256), borderValue=0.0)
+
+        # face swap
+        changed_face_cage = G.predict([np.expand_dims((im_aligned - 127.5) / 127.5, axis=0),
+                                       source_z])
+        changed_face = (changed_face_cage[0] + 1) / 2
+
+        # get inverse transformation landmarks
+        transformed_lmk = transform_landmark_points(M, lm_align)
+
+        # warp image back
+        iM, _ = inverse_estimate_norm(lm_align, transformed_lmk, 256, "arcface", shrink_factor=1.0)
+        iim_aligned = cv2.warpAffine(changed_face, iM, im_shape, borderValue=0.0)
+
+        # blend swapped face with target image
+        blend_mask = cv2.warpAffine(blend_mask_base, iM, im_shape, borderValue=0.0)
+        blend_mask = np.expand_dims(blend_mask, axis=-1)
+        total_img = (iim_aligned * blend_mask + total_img * (1 - blend_mask))
+
+    if opt.compare:
+        total_img = np.concatenate((im / 255.0, total_img), axis=1)
+
+    total_img = np.clip(total_img, 0, 1)
+    total_img *= 255.0
+    total_img = total_img.astype('uint8')
+
+    return total_img
+
+
+iface = gradio.Interface(run_inference,
+                         [gradio.inputs.Image(shape=None),
+                          gradio.inputs.Image(shape=None)],
+                         gradio.outputs.Image())
+iface.launch()

networks/generator.py

@@ -0,0 +1,321 @@
+from tensorflow.keras.layers import *
+from tensorflow.keras.models import Model
+from tensorflow_addons.layers import InstanceNormalization
+from networks.layers import AdaIN, AdaptiveAttention
+
+import numpy as np
+
+
+def residual_down_block(inputs, filters, resample=True):
+    x = inputs
+
+    r = Conv2D(filters=filters, kernel_size=1, strides=1, padding='same')(x)
+    if resample:
+        r = AveragePooling2D()(r)
+
+    x = InstanceNormalization()(x)
+    x = LeakyReLU(0.2)(x)
+    x = Conv2D(filters=filters, kernel_size=3, strides=1, padding='same')(x)
+
+    if resample:
+        x = AveragePooling2D()(x)
+
+    x = Add()([x, r])
+
+    return x
+
+
+def residual_up_block(inputs, filters, resample=True, name=None):
+    x, z_id = inputs
+
+    r = Conv2D(filters=filters, kernel_size=1, strides=1, padding='same')(x)
+    if resample:
+        r = UpSampling2D(interpolation='bilinear')(r)
+
+    x = InstanceNormalization()(x)
+    x = AdaIN()([x, z_id])
+    x = LeakyReLU(0.2)(x)
+    x = Conv2D(filters=filters, kernel_size=3, strides=1, padding='same')(x)
+
+    if resample:
+        x = UpSampling2D(interpolation='bilinear')(x)
+
+    x = Add()([x, r])
+
+    return x
+
+
+def adaptive_attention(inputs, filters, name=None):
+    x_t, x_s = inputs
+
+    m = Concatenate(axis=-1)([x_t, x_s])
+    m = Conv2D(filters=filters // 4, kernel_size=3, strides=1, padding='same')(m)
+    m = LeakyReLU(0.2)(m)
+    m = InstanceNormalization()(m)
+    m = Conv2D(filters=filters, kernel_size=1, strides=1, padding='same', activation='sigmoid', name=name)(m)
+
+    x = AdaptiveAttention()([m, x_t, x_s])
+
+    return x
+
+
+def adaptive_fusion_up_block(inputs, filters, resample=True, name=None):
+    x_t, x_s, z_id = inputs
+
+    x = adaptive_attention([x_t, x_s], x_t.shape[-1], name=name)
+
+    r = Conv2D(filters=filters, kernel_size=1, strides=1, padding='same')(x)
+    if resample:
+        r = UpSampling2D(interpolation='bilinear')(r)
+
+    x = InstanceNormalization()(x)
+    x = AdaIN()([x, z_id])
+    x = LeakyReLU(0.2)(x)
+    x = Conv2D(filters=filters, kernel_size=3, strides=1, padding='same')(x)
+
+    if resample:
+        x = UpSampling2D(interpolation='bilinear')(x)
+
+    x = Add()([x, r])
+
+    return x
+
+
+def dual_adaptive_fusion_up_block(inputs, filters, resample=True, name=None):
+    x_t, x_s, z_id = inputs
+
+    x = adaptive_attention([x_t, x_s], x_t.shape[-1], name=name + '_0')
+    x = adaptive_attention([x_t, x], x_t.shape[-1], name=name + '_1')
+
+    r = Conv2D(filters=filters, kernel_size=1, strides=1, padding='same')(x)
+    if resample:
+        r = UpSampling2D(interpolation='bilinear')(r)
+
+    x = InstanceNormalization()(x)
+    x = AdaIN()([x, z_id])
+    x = LeakyReLU(0.2)(x)
+    x = Conv2D(filters=filters, kernel_size=3, strides=1, padding='same')(x)
+
+    if resample:
+        x = UpSampling2D(interpolation='bilinear')(x)
+
+    x = Add()([x, r])
+
+    return x
+
+
+def adaptive_fusion_up_block_no_add(inputs, filters, resample=True, name=None):
+    x_t, x_s, z_id = inputs
+
+    x = adaptive_attention([x_t, x_s], x_t.shape[-1], name=name)
+
+    x = InstanceNormalization()(x)
+    x = AdaIN()([x, z_id])
+    x = LeakyReLU(0.2)(x)
+    x = Conv2D(filters=filters, kernel_size=3, strides=1, padding='same')(x)
+
+    return x
+
+
+def adaptive_fusion_up_block_concat_baseline(inputs, filters, resample=True, name=None):
+    x_t, x_s, z_id = inputs
+
+    x = Concatenate(axis=-1)([x_t, x_s])
+
+    r = Conv2D(filters=filters, kernel_size=1, strides=1, padding='same')(x)
+    if resample:
+        r = UpSampling2D(interpolation='bilinear')(r)
+
+    x = InstanceNormalization()(x)
+    x = AdaIN()([x, z_id])
+    x = LeakyReLU(0.2)(x)
+    x = Conv2D(filters=filters, kernel_size=3, strides=1, padding='same')(x)
+
+    if resample:
+        x = UpSampling2D(interpolation='bilinear')(x)
+
+    x = Add(name=name if name == 'final' else None)([x, r])
+
+    return x
+
+
+def adaptive_fusion_up_block_add_baseline(inputs, filters, resample=True, name=None):
+    x_t, x_s, z_id = inputs
+
+    x = Add()([x_t, x_s])
+
+    r = Conv2D(filters=filters, kernel_size=1, strides=1, padding='same')(x)
+    if resample:
+        r = UpSampling2D(interpolation='bilinear')(r)
+
+    x = InstanceNormalization()(x)
+    x = AdaIN()([x, z_id])
+    x = LeakyReLU(0.2)(x)
+    x = Conv2D(filters=filters, kernel_size=3, strides=1, padding='same')(x)
+
+    if resample:
+        x = UpSampling2D(interpolation='bilinear')(x)
+
+    x = Add()([x, r])
+
+    return x
+
+
+def get_generator_original(mapping_depth=4, mapping_size=256):
+    x_target = Input(shape=(256, 256, 3))
+    z_source = Input(shape=(512,))
+
+    z_id = z_source
+    for m in range(np.max([mapping_depth - 1, 0])):
+        z_id = Dense(mapping_size)(z_id)
+        z_id = LeakyReLU(0.2)(z_id)
+    if mapping_depth >= 1:
+        z_id = Dense(mapping_size)(z_id)
+
+    x_0 = Conv2D(filters=64, kernel_size=3, strides=1, padding='same')(x_target)    # 256
+
+    x_1 = residual_down_block(x_0, 128)                                             # 128
+
+    x_2 = residual_down_block(x_1, 256)                                             # 64
+
+    x_3 = residual_down_block(x_2, 512)
+
+    x_4 = residual_down_block(x_3, 512)
+
+    x_5 = residual_down_block(x_4, 512)
+
+    x_6 = residual_down_block(x_5, 512, resample=False)
+
+    u_5 = residual_up_block([x_6, z_id], 512, resample=False)
+
+    u_4 = residual_up_block([u_5, z_id], 512)
+
+    u_3 = residual_up_block([u_4, z_id], 512)
+
+    u_2 = residual_up_block([u_3, z_id], 256)                   # 64
+
+    u_1 = adaptive_fusion_up_block([x_2, u_2, z_id], 128, name='aff_attention_64x64')       # 128
+
+    u_0 = adaptive_fusion_up_block([x_1, u_1, z_id], 64, name='aff_attention_128x128')        # 256
+
+    out = adaptive_fusion_up_block([x_0, u_0, z_id], 3, resample=False, name='aff_attention_256x256')
+
+    gen_model = Model([x_target, z_source], out)
+    gen_model.summary()
+
+    return gen_model
+
+
+def make_layer(l_type, inputs, filters, resample, name=None):
+    if l_type == 'affa':
+        return adaptive_fusion_up_block(inputs, filters, resample=resample, name=name)
+    if l_type == 'd_affa':
+        return dual_adaptive_fusion_up_block(inputs, filters, resample=resample, name=name)
+    elif l_type == 'concat':
+        return adaptive_fusion_up_block_concat_baseline(inputs, filters, resample=resample, name=name)
+    elif l_type == 'no_skip':
+        return residual_up_block(inputs[1:], filters, resample=resample)
+
+
+def get_generator(up_types=None, mapping_depth=4, mapping_size=256):
+
+    if up_types is None:
+        up_types = ['no_skip', 'no_skip', 'd_affa', 'd_affa', 'd_affa', 'concat']
+
+    x_target = Input(shape=(256, 256, 3))
+    z_source = Input(shape=(512,))
+
+    z_id = z_source
+    for m in range(np.max([mapping_depth - 1, 0])):
+        z_id = Dense(mapping_size)(z_id)
+        z_id = LeakyReLU(0.2)(z_id)
+    if mapping_depth >= 1:
+        z_id = Dense(mapping_size)(z_id)
+
+    x_0 = Conv2D(filters=64, kernel_size=3, strides=1, padding='same')(x_target)    # 256
+
+    x_1 = residual_down_block(x_0, 128)                                             # 128
+
+    x_2 = residual_down_block(x_1, 256)                                             # 64
+
+    x_3 = residual_down_block(x_2, 512)
+
+    x_4 = residual_down_block(x_3, 512)
+
+    x_5 = residual_down_block(x_4, 512)
+
+    x_6 = residual_down_block(x_5, 512, resample=False)
+
+    u_5 = residual_up_block([x_6, z_id], 512, resample=False)
+
+    u_4 = make_layer(up_types[0], [x_5, u_5, z_id], 512, resample=True, name='16x16')
+
+    u_3 = make_layer(up_types[1], [x_4, u_4, z_id], 512, resample=True, name='32x32')
+
+    u_2 = make_layer(up_types[2], [x_3, u_3, z_id], 256, resample=True, name='64x64')
+
+    u_1 = make_layer(up_types[3], [x_2, u_2, z_id], 128, resample=True, name='128x128')
+
+    u_0 = make_layer(up_types[4], [x_1, u_1, z_id], 64, resample=True, name='256x256')
+
+    out = make_layer(up_types[5], [x_0, u_0, z_id], 3, resample=False, name='final')
+
+    gen_model = Model([x_target, z_source], out)
+    gen_model.summary()
+
+    return gen_model
+
+
+def get_generator_large(up_types=None, mapping_depth=4, mapping_size=512):
+
+    if up_types is None:
+        up_types = ['no_skip', 'no_skip', 'affa', 'affa', 'affa', 'concat']
+
+    x_target = Input(shape=(256, 256, 3))
+    z_source = Input(shape=(512,))
+
+    z_id = z_source
+    for m in range(np.max([mapping_depth - 1, 0])):
+        z_id = Dense(mapping_size)(z_id)
+        z_id = LeakyReLU(0.2)(z_id)
+    if mapping_depth >= 1:
+        z_id = Dense(mapping_size)(z_id)
+
+    x_0 = Conv2D(filters=64, kernel_size=3, strides=1, padding='same')(x_target)    # 256
+
+    x_1 = residual_down_block(x_0, 128)                                             # 128
+
+    x_2 = residual_down_block(x_1, 256)                                             # 64
+
+    x_3 = residual_down_block(x_2, 512)
+
+    x_4 = residual_down_block(x_3, 512)
+
+    x_5 = residual_down_block(x_4, 512)
+
+    b_0 = residual_up_block([x_5, z_id], 512, resample=False)
+
+    b_1 = residual_up_block([b_0, z_id], 512, resample=False)
+
+    b_2 = residual_up_block([b_1, z_id], 512, resample=False)
+
+    u_5 = residual_up_block([b_2, z_id], 512, resample=False)
+
+    u_4 = make_layer(up_types[0], [x_5, u_5, z_id], 512, resample=True, name='16x16')
+
+    u_3 = make_layer(up_types[1], [x_4, u_4, z_id], 512, resample=True, name='32x32')
+
+    u_2 = make_layer(up_types[2], [x_3, u_3, z_id], 256, resample=True, name='64x64')
+
+    u_1 = make_layer(up_types[3], [x_2, u_2, z_id], 128, resample=True, name='128x128')
+
+    u_0 = make_layer(up_types[4], [x_1, u_1, z_id], 64, resample=True, name='256x256')
+
+    out = make_layer(up_types[5], [x_0, u_0, z_id], 3, resample=False, name='final')
+
+    gen_model = Model([x_target, z_source], out)
+    gen_model.summary()
+
+    return gen_model
+
+

options/swap_options.py

@@ -0,0 +1,43 @@
+import argparse
+
+
+class SwapOptions():
+    def __init__(self):
+        self.parser = argparse.ArgumentParser()
+        self.initialized = False
+
+    def initialize(self):
+        # paths (data, models, etc...)
+        self.parser.add_argument('--arcface_path', type=str,
+                                 default="arcface_model/arcface/arc_res50.h5",
+                                 help='path to arcface model. Used to extract identity from source.')
+
+        # Video/Image necessary models
+        self.parser.add_argument('--retina_path', type=str,
+                                 default="retinaface/retinaface_res50.h5",
+                                 help='path to retinaface model.')
+        self.parser.add_argument('--compare', type=bool,
+                                 default=True,
+                                 help='If true, concatenates the frame with the manipulated frame')
+
+        self.parser.add_argument('--load', type=int,
+                                 default=30,
+                                 help='int of number to load checkpoint weights.')
+        self.parser.add_argument('--device_id', type=int, default=0,
+                                 help='which device to use')
+
+        # logging and checkpointing
+        self.parser.add_argument('--log_dir', type=str, default='logs/runs/',
+                                 help='logging directory')
+        self.parser.add_argument('--log_name', type=str, default='affa_f',
+                                 help='name of the run, change this to track several experiments')
+
+        self.parser.add_argument('--chkp_dir', type=str, default='checkpoints/',
+                                 help='checkpoint directory (will use same name as log_name!)')
+        self.initialized = True
+
+    def parse(self):
+        if not self.initialized:
+            self.initialize()
+        self.opt = self.parser.parse_args()
+        return self.opt

requirements.txt

@@ -0,0 +1,6 @@
+tensorflow
+tensorflow-addons
+opencv-python-headless
+scipy
+pillow
+scikit-image

retinaface/anchor.py

@@ -0,0 +1,296 @@
+"""Anchor utils modified from https://github.com/biubug6/Pytorch_Retinaface"""
+import math
+import tensorflow as tf
+import numpy as np
+from itertools import product as product
+
+
+###############################################################################
+#   Tensorflow / Numpy Priors                                                 #
+###############################################################################
+def prior_box(image_sizes, min_sizes, steps, clip=False):
+    """prior box"""
+    feature_maps = [
+        [math.ceil(image_sizes[0] / step), math.ceil(image_sizes[1] / step)]
+        for step in steps]
+
+    anchors = []
+    for k, f in enumerate(feature_maps):
+        for i, j in product(range(f[0]), range(f[1])):
+            for min_size in min_sizes[k]:
+                s_kx = min_size / image_sizes[1]
+                s_ky = min_size / image_sizes[0]
+                cx = (j + 0.5) * steps[k] / image_sizes[1]
+                cy = (i + 0.5) * steps[k] / image_sizes[0]
+                anchors += [cx, cy, s_kx, s_ky]
+
+    output = np.asarray(anchors).reshape([-1, 4])
+
+    if clip:
+        output = np.clip(output, 0, 1)
+
+    return output
+
+
+def prior_box_tf(image_sizes, min_sizes, steps, clip=False):
+    """prior box"""
+    image_sizes = tf.cast(tf.convert_to_tensor(image_sizes), tf.float32)
+    feature_maps = tf.math.ceil(
+        tf.reshape(image_sizes, [1, 2]) /
+        tf.reshape(tf.cast(steps, tf.float32), [-1, 1]))
+
+    anchors = []
+    for k in range(len(min_sizes)):
+        grid_x, grid_y = _meshgrid_tf(tf.range(feature_maps[k][1]),
+                                      tf.range(feature_maps[k][0]))
+        cx = (grid_x + 0.5) * steps[k] / image_sizes[1]
+        cy = (grid_y + 0.5) * steps[k] / image_sizes[0]
+        cxcy = tf.stack([cx, cy], axis=-1)
+        cxcy = tf.reshape(cxcy, [-1, 2])
+        cxcy = tf.repeat(cxcy, repeats=tf.shape(min_sizes[k])[0], axis=0)
+
+        sx = min_sizes[k] / image_sizes[1]
+        sy = min_sizes[k] / image_sizes[0]
+        sxsy = tf.stack([sx, sy], 1)
+        sxsy = tf.repeat(sxsy[tf.newaxis],
+                         repeats=tf.shape(grid_x)[0] * tf.shape(grid_x)[1],
+                         axis=0)
+        sxsy = tf.reshape(sxsy, [-1, 2])
+
+        anchors.append(tf.concat([cxcy, sxsy], 1))
+
+    output = tf.concat(anchors, axis=0)
+
+    if clip:
+        output = tf.clip_by_value(output, 0, 1)
+
+    return output
+
+
+def _meshgrid_tf(x, y):
+    """ workaround solution of the tf.meshgrid() issue:
+        https://github.com/tensorflow/tensorflow/issues/34470"""
+    grid_shape = [tf.shape(y)[0], tf.shape(x)[0]]
+    grid_x = tf.broadcast_to(tf.reshape(x, [1, -1]), grid_shape)
+    grid_y = tf.broadcast_to(tf.reshape(y, [-1, 1]), grid_shape)
+    return grid_x, grid_y
+
+
+###############################################################################
+#   Tensorflow Encoding                                                       #
+###############################################################################
+def encode_tf(labels, priors, match_thresh, ignore_thresh,
+              variances=[0.1, 0.2]):
+    """tensorflow encoding"""
+    assert ignore_thresh <= match_thresh
+    priors = tf.cast(priors, tf.float32)
+    bbox = labels[:, :4]
+    landm = labels[:, 4:-1]
+    landm_valid = labels[:, -1]  # 1: with landm, 0: w/o landm.
+
+    # jaccard index
+    overlaps = _jaccard(bbox, _point_form(priors))
+
+    # (Bipartite Matching)
+    # [num_objects] best prior for each ground truth
+    best_prior_overlap, best_prior_idx = tf.math.top_k(overlaps, k=1)
+    best_prior_overlap = best_prior_overlap[:, 0]
+    best_prior_idx = best_prior_idx[:, 0]
+
+    # [num_priors] best ground truth for each prior
+    overlaps_t = tf.transpose(overlaps)
+    best_truth_overlap, best_truth_idx = tf.math.top_k(overlaps_t, k=1)
+    best_truth_overlap = best_truth_overlap[:, 0]
+    best_truth_idx = best_truth_idx[:, 0]
+
+    # ensure best prior
+    def _loop_body(i, bt_idx, bt_overlap):
+        bp_mask = tf.one_hot(best_prior_idx[i], tf.shape(bt_idx)[0])
+        bp_mask_int = tf.cast(bp_mask, tf.int32)
+        new_bt_idx = bt_idx * (1 - bp_mask_int) + bp_mask_int * i
+        bp_mask_float = tf.cast(bp_mask, tf.float32)
+        new_bt_overlap = bt_overlap * (1 - bp_mask_float) + bp_mask_float * 2
+        return tf.cond(best_prior_overlap[i] > match_thresh,
+                       lambda: (i + 1, new_bt_idx, new_bt_overlap),
+                       lambda: (i + 1, bt_idx, bt_overlap))
+    _, best_truth_idx, best_truth_overlap = tf.while_loop(
+        lambda i, bt_idx, bt_overlap: tf.less(i, tf.shape(best_prior_idx)[0]),
+        _loop_body, [tf.constant(0), best_truth_idx, best_truth_overlap])
+
+    matches_bbox = tf.gather(bbox, best_truth_idx)  # [num_priors, 4]
+    matches_landm = tf.gather(landm, best_truth_idx)  # [num_priors, 10]
+    matches_landm_v = tf.gather(landm_valid, best_truth_idx)  # [num_priors]
+
+    loc_t = _encode_bbox(matches_bbox, priors, variances)
+    landm_t = _encode_landm(matches_landm, priors, variances)
+    landm_valid_t = tf.cast(matches_landm_v > 0, tf.float32)
+    conf_t = tf.cast(best_truth_overlap > match_thresh, tf.float32)
+    conf_t = tf.where(
+        tf.logical_and(best_truth_overlap < match_thresh,
+                       best_truth_overlap > ignore_thresh),
+        tf.ones_like(conf_t) * -1, conf_t)    # 1: pos, 0: neg, -1: ignore
+
+    return tf.concat([loc_t, landm_t, landm_valid_t[..., tf.newaxis],
+                      conf_t[..., tf.newaxis]], axis=1)
+
+
+def _encode_bbox(matched, priors, variances):
+    """Encode the variances from the priorbox layers into the ground truth
+    boxes we have matched (based on jaccard overlap) with the prior boxes.
+    Args:
+        matched: (tensor) Coords of ground truth for each prior in point-form
+            Shape: [num_priors, 4].
+        priors: (tensor) Prior boxes in center-offset form
+            Shape: [num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        encoded boxes (tensor), Shape: [num_priors, 4]
+    """
+
+    # dist b/t match center and prior's center
+    g_cxcy = (matched[:, :2] + matched[:, 2:]) / 2 - priors[:, :2]
+    # encode variance
+    g_cxcy /= (variances[0] * priors[:, 2:])
+    # match wh / prior wh
+    g_wh = (matched[:, 2:] - matched[:, :2]) / priors[:, 2:]
+    g_wh = tf.math.log(g_wh) / variances[1]
+    # return target for smooth_l1_loss
+    return tf.concat([g_cxcy, g_wh], 1)  # [num_priors,4]
+
+
+def _encode_landm(matched, priors, variances):
+    """Encode the variances from the priorbox layers into the ground truth
+    boxes we have matched (based on jaccard overlap) with the prior boxes.
+    Args:
+        matched: (tensor) Coords of ground truth for each prior in point-form
+            Shape: [num_priors, 10].
+        priors: (tensor) Prior boxes in center-offset form
+            Shape: [num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        encoded landm (tensor), Shape: [num_priors, 10]
+    """
+
+    # dist b/t match center and prior's center
+    matched = tf.reshape(matched, [tf.shape(matched)[0], 5, 2])
+    priors = tf.broadcast_to(
+        tf.expand_dims(priors, 1), [tf.shape(matched)[0], 5, 4])
+    g_cxcy = matched[:, :, :2] - priors[:, :, :2]
+    # encode variance
+    g_cxcy /= (variances[0] * priors[:, :, 2:])
+    # g_cxcy /= priors[:, :, 2:]
+    g_cxcy = tf.reshape(g_cxcy, [tf.shape(g_cxcy)[0], -1])
+    # return target for smooth_l1_loss
+    return g_cxcy
+
+
+def _point_form(boxes):
+    """ Convert prior_boxes to (xmin, ymin, xmax, ymax)
+    representation for comparison to point form ground truth data.
+    Args:
+        boxes: (tensor) center-size default boxes from priorbox layers.
+    Return:
+        boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes.
+    """
+    return tf.concat((boxes[:, :2] - boxes[:, 2:] / 2,
+                      boxes[:, :2] + boxes[:, 2:] / 2), axis=1)
+
+
+def _intersect(box_a, box_b):
+    """ We resize both tensors to [A,B,2]:
+    [A,2] -> [A,1,2] -> [A,B,2]
+    [B,2] -> [1,B,2] -> [A,B,2]
+    Then we compute the area of intersect between box_a and box_b.
+    Args:
+      box_a: (tensor) bounding boxes, Shape: [A,4].
+      box_b: (tensor) bounding boxes, Shape: [B,4].
+    Return:
+      (tensor) intersection area, Shape: [A,B].
+    """
+    A = tf.shape(box_a)[0]
+    B = tf.shape(box_b)[0]
+    max_xy = tf.minimum(
+        tf.broadcast_to(tf.expand_dims(box_a[:, 2:], 1), [A, B, 2]),
+        tf.broadcast_to(tf.expand_dims(box_b[:, 2:], 0), [A, B, 2]))
+    min_xy = tf.maximum(
+        tf.broadcast_to(tf.expand_dims(box_a[:, :2], 1), [A, B, 2]),
+        tf.broadcast_to(tf.expand_dims(box_b[:, :2], 0), [A, B, 2]))
+    inter = tf.maximum((max_xy - min_xy), tf.zeros_like(max_xy - min_xy))
+    return inter[:, :, 0] * inter[:, :, 1]
+
+
+def _jaccard(box_a, box_b):
+    """Compute the jaccard overlap of two sets of boxes.  The jaccard overlap
+    is simply the intersection over union of two boxes.  Here we operate on
+    ground truth boxes and default boxes.
+    E.g.:
+        A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B)
+    Args:
+        box_a: (tensor) Ground truth bounding boxes, Shape: [num_objects,4]
+        box_b: (tensor) Prior boxes from priorbox layers, Shape: [num_priors,4]
+    Return:
+        jaccard overlap: (tensor) Shape: [box_a.size(0), box_b.size(0)]
+    """
+    inter = _intersect(box_a, box_b)
+    area_a = tf.broadcast_to(
+        tf.expand_dims(
+            (box_a[:, 2] - box_a[:, 0]) * (box_a[:, 3] - box_a[:, 1]), 1),
+        tf.shape(inter))  # [A,B]
+    area_b = tf.broadcast_to(
+        tf.expand_dims(
+            (box_b[:, 2] - box_b[:, 0]) * (box_b[:, 3] - box_b[:, 1]), 0),
+        tf.shape(inter))  # [A,B]
+    union = area_a + area_b - inter
+    return inter / union  # [A,B]
+
+
+###############################################################################
+#   Tensorflow Decoding                                                       #
+###############################################################################
+def decode_tf(labels, priors, variances=[0.1, 0.2]):
+    """tensorflow decoding"""
+    bbox = _decode_bbox(labels[:, :4], priors, variances)
+    landm = _decode_landm(labels[:, 4:14], priors, variances)
+    landm_valid = labels[:, 14][:, tf.newaxis]
+    conf = labels[:, 15][:, tf.newaxis]
+
+    return tf.concat([bbox, landm, landm_valid, conf], axis=1)
+
+
+def _decode_bbox(pre, priors, variances=[0.1, 0.2]):
+    """Decode locations from predictions using priors to undo
+    the encoding we did for offset regression at train time.
+    Args:
+        pre (tensor): location predictions for loc layers,
+            Shape: [num_priors,4]
+        priors (tensor): Prior boxes in center-offset form.
+            Shape: [num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        decoded bounding box predictions
+    """
+    centers = priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:]
+    sides = priors[:, 2:] * tf.math.exp(pre[:, 2:] * variances[1])
+
+    return tf.concat([centers - sides / 2, centers + sides / 2], axis=1)
+
+
+def _decode_landm(pre, priors, variances=[0.1, 0.2]):
+    """Decode landm from predictions using priors to undo
+    the encoding we did for offset regression at train time.
+    Args:
+        pre (tensor): landm predictions for loc layers,
+            Shape: [num_priors,10]
+        priors (tensor): Prior boxes in center-offset form.
+            Shape: [num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        decoded landm predictions
+    """
+    landms = tf.concat(
+        [priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:],
+         priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:],
+         priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:],
+         priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:],
+         priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:]], axis=1)
+    return landms

retinaface/models.py

@@ -0,0 +1,301 @@
+import tensorflow as tf
+from tensorflow.keras import Model
+from tensorflow.keras.applications import MobileNetV2, ResNet50
+from tensorflow.keras.layers import Input, Conv2D, ReLU, LeakyReLU
+from retinaface.anchor import decode_tf, prior_box_tf
+
+
+def _regularizer(weights_decay):
+    """l2 regularizer"""
+    return tf.keras.regularizers.l2(weights_decay)
+
+
+def _kernel_init(scale=1.0, seed=None):
+    """He normal initializer"""
+    return tf.keras.initializers.he_normal()
+
+
+class BatchNormalization(tf.keras.layers.BatchNormalization):
+    """Make trainable=False freeze BN for real (the og version is sad).
+       ref: https://github.com/zzh8829/yolov3-tf2
+    """
+    def __init__(self, axis=-1, momentum=0.9, epsilon=1e-5, center=True,
+                 scale=True, name=None, **kwargs):
+        super(BatchNormalization, self).__init__(
+            axis=axis, momentum=momentum, epsilon=epsilon, center=center,
+            scale=scale, name=name, **kwargs)
+
+    def call(self, x, training=False):
+        if training is None:
+            training = tf.constant(False)
+        training = tf.logical_and(training, self.trainable)
+
+        return super().call(x, training)
+
+
+def Backbone(backbone_type='ResNet50', use_pretrain=True):
+    """Backbone Model"""
+    weights = None
+    if use_pretrain:
+        weights = 'imagenet'
+
+    def backbone(x):
+        if backbone_type == 'ResNet50':
+            extractor = ResNet50(
+                input_shape=x.shape[1:], include_top=False, weights=weights)
+            pick_layer1 = 80  # [80, 80, 512]
+            pick_layer2 = 142  # [40, 40, 1024]
+            pick_layer3 = 174  # [20, 20, 2048]
+            preprocess = tf.keras.applications.resnet.preprocess_input
+        elif backbone_type == 'MobileNetV2':
+            extractor = MobileNetV2(
+                input_shape=x.shape[1:], include_top=False, weights=weights)
+            pick_layer1 = 54  # [80, 80, 32]
+            pick_layer2 = 116  # [40, 40, 96]
+            pick_layer3 = 143  # [20, 20, 160]
+            preprocess = tf.keras.applications.mobilenet_v2.preprocess_input
+        else:
+            raise NotImplementedError(
+                'Backbone type {} is not recognized.'.format(backbone_type))
+
+        return Model(extractor.input,
+                     (extractor.layers[pick_layer1].output,
+                      extractor.layers[pick_layer2].output,
+                      extractor.layers[pick_layer3].output),
+                     name=backbone_type + '_extrator')(preprocess(x))
+
+    return backbone
+
+
+class ConvUnit(tf.keras.layers.Layer):
+    """Conv + BN + Act"""
+    def __init__(self, f, k, s, wd, act=None, **kwargs):
+        super(ConvUnit, self).__init__(**kwargs)
+        self.conv = Conv2D(filters=f, kernel_size=k, strides=s, padding='same',
+                           kernel_initializer=_kernel_init(),
+                           kernel_regularizer=_regularizer(wd),
+                           use_bias=False)
+        self.bn = BatchNormalization()
+
+        if act is None:
+            self.act_fn = tf.identity
+        elif act == 'relu':
+            self.act_fn = ReLU()
+        elif act == 'lrelu':
+            self.act_fn = LeakyReLU(0.1)
+        else:
+            raise NotImplementedError(
+                'Activation function type {} is not recognized.'.format(act))
+
+    def call(self, x):
+        return self.act_fn(self.bn(self.conv(x)))
+
+
+class FPN(tf.keras.layers.Layer):
+    """Feature Pyramid Network"""
+    def __init__(self, out_ch, wd, **kwargs):
+        super(FPN, self).__init__(**kwargs)
+        act = 'relu'
+        self.out_ch = out_ch
+        self.wd = wd
+        if (out_ch <= 64):
+            act = 'lrelu'
+
+        self.output1 = ConvUnit(f=out_ch, k=1, s=1, wd=wd, act=act)
+        self.output2 = ConvUnit(f=out_ch, k=1, s=1, wd=wd, act=act)
+        self.output3 = ConvUnit(f=out_ch, k=1, s=1, wd=wd, act=act)
+        self.merge1 = ConvUnit(f=out_ch, k=3, s=1, wd=wd, act=act)
+        self.merge2 = ConvUnit(f=out_ch, k=3, s=1, wd=wd, act=act)
+
+    def call(self, x):
+        output1 = self.output1(x[0])  # [80, 80, out_ch]
+        output2 = self.output2(x[1])  # [40, 40, out_ch]
+        output3 = self.output3(x[2])  # [20, 20, out_ch]
+
+        up_h, up_w = tf.shape(output2)[1], tf.shape(output2)[2]
+        up3 = tf.image.resize(output3, [up_h, up_w], method='nearest')
+        output2 = output2 + up3
+        output2 = self.merge2(output2)
+
+        up_h, up_w = tf.shape(output1)[1], tf.shape(output1)[2]
+        up2 = tf.image.resize(output2, [up_h, up_w], method='nearest')
+        output1 = output1 + up2
+        output1 = self.merge1(output1)
+
+        return output1, output2, output3
+    
+    def get_config(self):
+        config = {
+            'out_ch': self.out_ch,
+            'wd': self.wd,
+        }
+        base_config = super(FPN, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+
+class SSH(tf.keras.layers.Layer):
+    """Single Stage Headless Layer"""
+    def __init__(self, out_ch, wd, **kwargs):
+        super(SSH, self).__init__(**kwargs)
+        assert out_ch % 4 == 0
+        self.out_ch = out_ch
+        self.wd = wd
+        act = 'relu'
+        if (out_ch <= 64):
+            act = 'lrelu'
+
+        self.conv_3x3 = ConvUnit(f=out_ch // 2, k=3, s=1, wd=wd, act=None)
+
+        self.conv_5x5_1 = ConvUnit(f=out_ch // 4, k=3, s=1, wd=wd, act=act)
+        self.conv_5x5_2 = ConvUnit(f=out_ch // 4, k=3, s=1, wd=wd, act=None)
+
+        self.conv_7x7_2 = ConvUnit(f=out_ch // 4, k=3, s=1, wd=wd, act=act)
+        self.conv_7x7_3 = ConvUnit(f=out_ch // 4, k=3, s=1, wd=wd, act=None)
+
+        self.relu = ReLU()
+
+    def call(self, x):
+        conv_3x3 = self.conv_3x3(x)
+
+        conv_5x5_1 = self.conv_5x5_1(x)
+        conv_5x5 = self.conv_5x5_2(conv_5x5_1)
+
+        conv_7x7_2 = self.conv_7x7_2(conv_5x5_1)
+        conv_7x7 = self.conv_7x7_3(conv_7x7_2)
+
+        output = tf.concat([conv_3x3, conv_5x5, conv_7x7], axis=3)
+        output = self.relu(output)
+
+        return output
+    
+    def get_config(self):
+        config = {
+            'out_ch': self.out_ch,
+            'wd': self.wd,
+        }
+        base_config = super(SSH, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+
+class BboxHead(tf.keras.layers.Layer):
+    """Bbox Head Layer"""
+    def __init__(self, num_anchor, wd, **kwargs):
+        super(BboxHead, self).__init__(**kwargs)
+        self.num_anchor = num_anchor
+        self.wd = wd
+        self.conv = Conv2D(filters=num_anchor * 4, kernel_size=1, strides=1)
+
+    def call(self, x):
+        h, w = tf.shape(x)[1], tf.shape(x)[2]
+        x = self.conv(x)
+
+        return tf.reshape(x, [-1, h * w * self.num_anchor, 4])
+    
+    def get_config(self):
+        config = {
+            'num_anchor': self.num_anchor,
+            'wd': self.wd,
+        }
+        base_config = super(BboxHead, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+
+class LandmarkHead(tf.keras.layers.Layer):
+    """Landmark Head Layer"""
+    def __init__(self, num_anchor, wd, name='LandmarkHead', **kwargs):
+        super(LandmarkHead, self).__init__(name=name, **kwargs)
+        self.num_anchor = num_anchor
+        self.wd = wd
+        self.conv = Conv2D(filters=num_anchor * 10, kernel_size=1, strides=1)
+
+    def call(self, x):
+        h, w = tf.shape(x)[1], tf.shape(x)[2]
+        x = self.conv(x)
+
+        return tf.reshape(x, [-1, h * w * self.num_anchor, 10])
+
+    def get_config(self):
+        config = {
+            'num_anchor': self.num_anchor,
+            'wd': self.wd,
+        }
+        base_config = super(LandmarkHead, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+
+class ClassHead(tf.keras.layers.Layer):
+    """Class Head Layer"""
+    def __init__(self, num_anchor, wd, name='ClassHead', **kwargs):
+        super(ClassHead, self).__init__(name=name, **kwargs)
+        self.num_anchor = num_anchor
+        self.wd = wd
+        self.conv = Conv2D(filters=num_anchor * 2, kernel_size=1, strides=1)
+
+    def call(self, x):
+        h, w = tf.shape(x)[1], tf.shape(x)[2]
+        x = self.conv(x)
+
+        return tf.reshape(x, [-1, h * w * self.num_anchor, 2])
+
+    def get_config(self):
+        config = {
+            'num_anchor': self.num_anchor,
+            'wd': self.wd,
+        }
+        base_config = super(ClassHead, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+
+def RetinaFaceModel(cfg, training=False, iou_th=0.4, score_th=0.02,
+                    name='RetinaFaceModel'):
+    """Retina Face Model"""
+    input_size = cfg['input_size'] if training else None
+    wd = cfg['weights_decay']
+    out_ch = cfg['out_channel']
+    num_anchor = len(cfg['min_sizes'][0])
+    backbone_type = cfg['backbone_type']
+
+    # define model
+    x = inputs = Input([input_size, input_size, 3], name='input_image')
+
+    x = Backbone(backbone_type=backbone_type)(x)
+
+    fpn = FPN(out_ch=out_ch, wd=wd)(x)
+
+    features = [SSH(out_ch=out_ch, wd=wd)(f)
+                for i, f in enumerate(fpn)]
+
+    bbox_regressions = tf.concat(
+        [BboxHead(num_anchor, wd=wd)(f)
+         for i, f in enumerate(features)], axis=1)
+    landm_regressions = tf.concat(
+        [LandmarkHead(num_anchor, wd=wd, name=f'LandmarkHead_{i}')(f)
+         for i, f in enumerate(features)], axis=1)
+    classifications = tf.concat(
+        [ClassHead(num_anchor, wd=wd, name=f'ClassHead_{i}')(f)
+         for i, f in enumerate(features)], axis=1)
+
+    classifications = tf.keras.layers.Softmax(axis=-1)(classifications)
+
+    if training:
+        out = (bbox_regressions, landm_regressions, classifications)
+    else:
+        # only for batch size 1
+        preds = tf.concat(  # [bboxes, landms, landms_valid, conf]
+            [bbox_regressions[0],
+             landm_regressions[0],
+             tf.ones_like(classifications[0, :, 0][..., tf.newaxis]),
+             classifications[0, :, 1][..., tf.newaxis]], 1)
+        priors = prior_box_tf((tf.shape(inputs)[1], tf.shape(inputs)[2]), cfg['min_sizes'], cfg['steps'], cfg['clip'])
+        decode_preds = decode_tf(preds, priors, cfg['variances'])
+
+        selected_indices = tf.image.non_max_suppression(
+            boxes=decode_preds[:, :4],
+            scores=decode_preds[:, -1],
+            max_output_size=tf.shape(decode_preds)[0],
+            iou_threshold=iou_th,
+            score_threshold=score_th)
+
+        out = tf.gather(decode_preds, selected_indices)
+
+    return Model(inputs, out, name=name), Model(inputs, [bbox_regressions, landm_regressions, classifications], name=name + '_bb_only')

retinaface/ops.py

@@ -0,0 +1,27 @@
+from retinaface.anchor import decode_tf, prior_box_tf
+import tensorflow as tf
+
+
+def extract_detections(bbox_regressions, landm_regressions, classifications, image_sizes, iou_th=0.4, score_th=0.02):
+    min_sizes = [[16, 32], [64, 128], [256, 512]]
+    steps = [8, 16, 32]
+    variances = [0.1, 0.2]
+    preds = tf.concat(  # [bboxes, landms, landms_valid, conf]
+        [bbox_regressions,
+         landm_regressions,
+         tf.ones_like(classifications[:, 0][..., tf.newaxis]),
+         classifications[:, 1][..., tf.newaxis]], 1)
+    priors = prior_box_tf(image_sizes, min_sizes, steps, False)
+    decode_preds = decode_tf(preds, priors, variances)
+
+    selected_indices = tf.image.non_max_suppression(
+        boxes=decode_preds[:, :4],
+        scores=decode_preds[:, -1],
+        max_output_size=tf.shape(decode_preds)[0],
+        iou_threshold=iou_th,
+        score_threshold=score_th)
+
+    out = tf.gather(decode_preds, selected_indices)
+
+    return out
+

utils/utils.py

@@ -0,0 +1,377 @@
+import json
+from tensorflow.keras.models import model_from_json
+from networks.layers import AdaIN, AdaptiveAttention
+import tensorflow as tf
+
+import numpy as np
+import cv2
+import math
+from skimage import transform as trans
+from scipy.signal import convolve2d
+from skimage.color import rgb2yuv, yuv2rgb
+
+from PIL import Image
+
+
+def save_model_internal(model, path, name, num):
+    json_model = model.to_json()
+    with open(path + name + '.json', "w") as json_file:
+        json_file.write(json_model)
+
+    model.save_weights(path + name + '_' + str(num) + '.h5')
+
+
+def load_model_internal(path, name, num):
+    with open(path + name + '.json', 'r') as json_file:
+        model_dict = json_file.read()
+
+    mod = model_from_json(model_dict, custom_objects={'AdaIN': AdaIN, 'AdaptiveAttention': AdaptiveAttention})
+    mod.load_weights(path + name + '_' + str(num) + '.h5')
+
+    return mod
+
+
+def save_training_meta(state_dict, path, num):
+    with open(path + str(num) + '.json', 'w') as json_file:
+        json.dump(state_dict, json_file, indent=2)
+
+
+def load_training_meta(path, num):
+    with open(path + str(num) + '.json', 'r') as json_file:
+        state_dict = json.load(json_file)
+    return state_dict
+
+
+def log_info(sw, results_dict, iteration):
+    with sw.as_default():
+        for key in results_dict.keys():
+            tf.summary.scalar(key, results_dict[key], step=iteration)
+
+
+src1 = np.array([[51.642, 50.115], [57.617, 49.990], [35.740, 69.007],
+                 [51.157, 89.050], [57.025, 89.702]],
+                dtype=np.float32)
+# <--left
+src2 = np.array([[45.031, 50.118], [65.568, 50.872], [39.677, 68.111],
+                 [45.177, 86.190], [64.246, 86.758]],
+                dtype=np.float32)
+
+# ---frontal
+src3 = np.array([[39.730, 51.138], [72.270, 51.138], [56.000, 68.493],
+                 [42.463, 87.010], [69.537, 87.010]],
+                dtype=np.float32)
+
+# -->right
+src4 = np.array([[46.845, 50.872], [67.382, 50.118], [72.737, 68.111],
+                 [48.167, 86.758], [67.236, 86.190]],
+                dtype=np.float32)
+
+# -->right profile
+src5 = np.array([[54.796, 49.990], [60.771, 50.115], [76.673, 69.007],
+                 [55.388, 89.702], [61.257, 89.050]],
+                dtype=np.float32)
+
+src = np.array([src1, src2, src3, src4, src5])
+src_map = {112: src, 224: src * 2}
+
+# Left eye, right eye, nose, left mouth, right mouth
+arcface_src = np.array(
+    [[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366],
+     [41.5493, 92.3655], [70.7299, 92.2041]],
+    dtype=np.float32)
+
+arcface_src = np.expand_dims(arcface_src, axis=0)
+
+
+def extract_face(img, bb, absolute_center, mode='arcface', extention_rate=0.05, debug=False):
+    """Extract face from image given a bounding box"""
+    # bbox
+    x1, y1, x2, y2 = bb + 60
+    adjusted_absolute_center = (absolute_center[0] + 60, absolute_center[1] + 60)
+    if debug:
+        print(bb + 60)
+        x1, y1, x2, y2 = bb
+        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 3)
+        cv2.circle(img, absolute_center, 1, (255, 0, 255), 2)
+        Image.fromarray(img).show()
+        x1, y1, x2, y2 = bb + 60
+    # Pad image in case face is out of frame
+    padded_img = np.zeros(shape=(248, 248, 3), dtype=np.uint8)
+    padded_img[60:-60, 60:-60, :] = img
+
+    if debug:
+        cv2.rectangle(padded_img, (x1, y1), (x2, y2), (0, 255, 255), 3)
+        cv2.circle(padded_img, adjusted_absolute_center, 1, (255, 255, 255), 2)
+        Image.fromarray(padded_img).show()
+
+    y_len = abs(y1 - y2)
+    x_len = abs(x1 - x2)
+
+    new_len = (y_len + x_len) // 2
+
+    extension = int(new_len * extention_rate)
+
+    x_adjust = (x_len - new_len) // 2
+    y_adjust = (y_len - new_len) // 2
+
+    x_1_adjusted = x1 + x_adjust - extension
+    x_2_adjusted = x2 - x_adjust + extension
+
+    if mode == 'arcface':
+        y_1_adjusted = y1 - extension
+        y_2_adjusted = y2 - 2 * y_adjust + extension
+    else:
+        y_1_adjusted = y1 + 2 * y_adjust - extension
+        y_2_adjusted = y2 + extension
+
+    move_x = adjusted_absolute_center[0] - (x_1_adjusted + x_2_adjusted) // 2
+    move_y = adjusted_absolute_center[1] - (y_1_adjusted + y_2_adjusted) // 2
+
+    x_1_adjusted = x_1_adjusted + move_x
+    x_2_adjusted = x_2_adjusted + move_x
+    y_1_adjusted = y_1_adjusted + move_y
+    y_2_adjusted = y_2_adjusted + move_y
+
+    # print(y_1_adjusted, y_2_adjusted, x_1_adjusted, x_2_adjusted)
+
+    return padded_img[y_1_adjusted:y_2_adjusted, x_1_adjusted:x_2_adjusted]
+
+
+def distance(a, b):
+    return np.sqrt((a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2)
+
+
+def euclidean_distance(a, b):
+    x1 = a[0]; y1 = a[1]
+    x2 = b[0]; y2 = b[1]
+    return np.sqrt(((x2 - x1) * (x2 - x1)) + ((y2 - y1) * (y2 - y1)))
+
+
+def align_face(img, landmarks, debug=False):
+    nose, right_eye, left_eye = landmarks
+
+    left_eye_x = left_eye[0]
+    left_eye_y = left_eye[1]
+
+    right_eye_x = right_eye[0]
+    right_eye_y = right_eye[1]
+
+    center_eye = ((left_eye[0] + right_eye[0]) // 2, (left_eye[1] + right_eye[1]) // 2)
+
+    if left_eye_y < right_eye_y:
+        point_3rd = (right_eye_x, left_eye_y)
+        direction = -1
+    else:
+        point_3rd = (left_eye_x, right_eye_y)
+        direction = 1
+
+    if debug:
+        cv2.circle(img, point_3rd, 1, (255, 0, 0), 1)
+        cv2.circle(img, center_eye, 1, (255, 0, 0), 1)
+
+        cv2.line(img, right_eye, left_eye, (0, 0, 0), 1)
+        cv2.line(img, left_eye, point_3rd, (0, 0, 0), 1)
+        cv2.line(img, right_eye, point_3rd, (0, 0, 0), 1)
+
+    a = euclidean_distance(left_eye, point_3rd)
+    b = euclidean_distance(right_eye, left_eye)
+    c = euclidean_distance(right_eye, point_3rd)
+
+    cos_a = (b * b + c * c - a * a) / (2 * b * c)
+
+    angle = np.arccos(cos_a)
+
+    angle = (angle * 180) / np.pi
+
+    if direction == -1:
+        angle = 90 - angle
+        ang = math.radians(direction * angle)
+    else:
+        ang = math.radians(direction * angle)
+        angle = 0 - angle
+
+    M = cv2.getRotationMatrix2D((64, 64), angle, 1)
+    new_img = cv2.warpAffine(img, M, (128, 128),
+                            flags=cv2.INTER_CUBIC)
+
+    rotated_nose = (int((nose[0] - 64) * np.cos(ang) - (nose[1] - 64) * np.sin(ang) + 64),
+                    int((nose[0] - 64) * np.sin(ang) + (nose[1] - 64) * np.cos(ang) + 64))
+
+    rotated_center_eye = (int((center_eye[0] - 64) * np.cos(ang) - (center_eye[1] - 64) * np.sin(ang) + 64),
+                          int((center_eye[0] - 64) * np.sin(ang) + (center_eye[1] - 64) * np.cos(ang) + 64))
+
+    abolute_center = (rotated_center_eye[0], (rotated_nose[1] + rotated_center_eye[1]) // 2)
+
+    if debug:
+        cv2.circle(new_img, rotated_nose, 1, (0, 0, 255), 1)
+        cv2.circle(new_img, rotated_center_eye, 1, (0, 0, 255), 1)
+        cv2.circle(new_img, abolute_center, 1, (0, 0, 255), 1)
+
+    return new_img, abolute_center
+
+
+def estimate_norm(lmk, image_size=112, mode='arcface', shrink_factor=1.0):
+    assert lmk.shape == (5, 2)
+    tform = trans.SimilarityTransform()
+    lmk_tran = np.insert(lmk, 2, values=np.ones(5), axis=1)
+    min_M = []
+    min_index = []
+    min_error = float('inf')
+    src_factor = image_size / 112
+    if mode == 'arcface':
+        src = arcface_src * shrink_factor + (1 - shrink_factor) * 56
+        src = src * src_factor
+    else:
+        src = src_map[image_size] * src_factor
+    for i in np.arange(src.shape[0]):
+        tform.estimate(lmk, src[i])
+        M = tform.params[0:2, :]
+        results = np.dot(M, lmk_tran.T)
+        results = results.T
+        error = np.sum(np.sqrt(np.sum((results - src[i])**2, axis=1)))
+        #         print(error)
+        if error < min_error:
+            min_error = error
+            min_M = M
+            min_index = i
+    return min_M, min_index
+
+
+def inverse_estimate_norm(lmk, t_lmk, image_size=112, mode='arcface', shrink_factor=1.0):
+    assert lmk.shape == (5, 2)
+    tform = trans.SimilarityTransform()
+    lmk_tran = np.insert(lmk, 2, values=np.ones(5), axis=1)
+    min_M = []
+    min_index = []
+    min_error = float('inf')
+    src_factor = image_size / 112
+    if mode == 'arcface':
+        src = arcface_src * shrink_factor + (1 - shrink_factor) * 56
+        src = src * src_factor
+    else:
+        src = src_map[image_size] * src_factor
+    for i in np.arange(src.shape[0]):
+        tform.estimate(t_lmk, lmk)
+        M = tform.params[0:2, :]
+        results = np.dot(M, lmk_tran.T)
+        results = results.T
+        error = np.sum(np.sqrt(np.sum((results - src[i])**2, axis=1)))
+        #         print(error)
+        if error < min_error:
+            min_error = error
+            min_M = M
+            min_index = i
+    return min_M, min_index
+
+
+def norm_crop(img, landmark, image_size=112, mode='arcface', shrink_factor=1.0):
+    """
+    Align and crop the image based of the facial landmarks in the image. The alignment is done with
+    a similarity transformation based of source coordinates.
+    :param img: Image to transform.
+    :param landmark: Five landmark coordinates in the image.
+    :param image_size: Desired output size after transformation.
+    :param mode: 'arcface' aligns the face for the use of Arcface facial recognition model. Useful for
+    both facial recognition tasks and face swapping tasks.
+    :param shrink_factor: Shrink factor that shrinks the source landmark coordinates. This will include more border
+    information around the face. Useful when you want to include more background information when performing face swaps.
+    The lower the shrink factor the more of the face is included. Default value 1.0 will align the image to be ready
+    for the Arcface recognition model, but usually omits part of the chin. Value of 0.0 would transform all source points
+    to the middle of the image, probably rendering the alignment procedure useless.
+
+    If you process the image with a shrink factor of 0.85 and then want to extract the identity embedding with arcface,
+    you simply do a central crop of factor 0.85 to yield same cropped result as using shrink factor 1.0. This will
+    reduce the resolution, the recommendation is to processed images to output resolutions higher than 112 is using
+    Arcface. This will make sure no information is lost by resampling the image after central crop.
+    :return: Returns the transformed image.
+    """
+    M, pose_index = estimate_norm(landmark, image_size, mode, shrink_factor=shrink_factor)
+    warped = cv2.warpAffine(img, M, (image_size, image_size), borderValue=0.0)
+    return warped
+
+
+def transform_landmark_points(M, points):
+    lmk_tran = np.insert(points, 2, values=np.ones(5), axis=1)
+    transformed_lmk = np.dot(M, lmk_tran.T)
+    transformed_lmk = transformed_lmk.T
+
+    return transformed_lmk
+
+
+def multi_convolver(image, kernel, iterations):
+    if kernel == "Sharpen":
+        kernel = np.array([[0, -1, 0],
+                           [-1, 5, -1],
+                           [0, -1, 0]])
+    elif kernel == "Unsharp_mask":
+        kernel = np.array([[1, 4, 6, 4, 1],
+                           [4, 16, 24, 16, 1],
+                           [6, 24, -476, 24, 1],
+                           [4, 16, 24, 16, 1],
+                           [1, 4, 6, 4, 1]]) * (-1 / 256)
+    elif kernel == "Blur":
+        kernel = (1 / 16.0) * np.array([[1., 2., 1.],
+                                        [2., 4., 2.],
+                                        [1., 2., 1.]])
+    for i in range(iterations):
+        image = convolve2d(image, kernel, 'same', boundary='fill', fillvalue = 0)
+    return image
+
+
+def convolve_rgb(image, kernel, iterations=1):
+    img_yuv = rgb2yuv(image)
+    img_yuv[:, :, 0] = multi_convolver(img_yuv[:, :, 0], kernel,
+                                       iterations)
+    final_image = yuv2rgb(img_yuv)
+
+    return final_image.astype('float32')
+
+
+def generate_mask_from_landmarks(lms, im_size):
+    blend_mask_lm = np.zeros(shape=(im_size, im_size, 3), dtype='float32')
+
+    # EYES
+    blend_mask_lm = cv2.circle(blend_mask_lm,
+                               (int(lms[0][0]), int(lms[0][1])), 12, (255, 255, 255), 30)
+    blend_mask_lm = cv2.circle(blend_mask_lm,
+                               (int(lms[1][0]), int(lms[1][1])), 12, (255, 255, 255), 30)
+    blend_mask_lm = cv2.circle(blend_mask_lm,
+                               (int((lms[0][0] + lms[1][0]) / 2), int((lms[0][1] + lms[1][1]) / 2)),
+                               16, (255, 255, 255), 65)
+
+    # NOSE
+    blend_mask_lm = cv2.circle(blend_mask_lm,
+                               (int(lms[2][0]), int(lms[2][1])), 5, (255, 255, 255), 5)
+    blend_mask_lm = cv2.circle(blend_mask_lm,
+                               (int((lms[0][0] + lms[1][0]) / 2), int(lms[2][1])), 16, (255, 255, 255), 100)
+
+    # MOUTH
+    blend_mask_lm = cv2.circle(blend_mask_lm,
+                               (int(lms[3][0]), int(lms[3][1])), 6, (255, 255, 255), 30)
+    blend_mask_lm = cv2.circle(blend_mask_lm,
+                               (int(lms[4][0]), int(lms[4][1])), 6, (255, 255, 255), 30)
+
+    blend_mask_lm = cv2.circle(blend_mask_lm,
+                               (int((lms[3][0] + lms[4][0]) / 2), int((lms[3][1] + lms[4][1]) / 2)),
+                               16, (255, 255, 255), 40)
+    return blend_mask_lm
+
+
+def display_distance_text(im, distance, lms, im_w, im_h, scale=2):
+    blended_insert = cv2.putText(im, str(distance)[:4],
+                                 (int(lms[4] * im_w * 0.5), int(lms[5] * im_h * 0.8)),
+                                 cv2.FONT_HERSHEY_SIMPLEX, scale * 0.5, (0.08, 0.16, 0.08), int(scale * 2))
+    blended_insert = cv2.putText(blended_insert, str(distance)[:4],
+                                 (int(lms[4] * im_w * 0.5), int(lms[5] * im_h * 0.8)),
+                                 cv2.FONT_HERSHEY_SIMPLEX, scale*  0.5, (0.3, 0.7, 0.32), int(scale * 1))
+    return blended_insert
+
+
+def get_lm(annotation, im_w, im_h):
+    lm_align = np.array([[annotation[4] * im_w, annotation[5] * im_h],
+                         [annotation[6] * im_w, annotation[7] * im_h],
+                         [annotation[8] * im_w, annotation[9] * im_h],
+                         [annotation[10] * im_w, annotation[11] * im_h],
+                         [annotation[12] * im_w, annotation[13] * im_h]],
+                        dtype=np.float32)
+    return lm_align