Update app.py with potential fix

app.py

@@ -2,7 +2,13 @@ import gradio
 from huggingface_hub import Repository
 import os
 
-from utils.utils import norm_crop, estimate_norm, inverse_estimate_norm, transform_landmark_points, get_lm
+from utils.utils import (
+    norm_crop,
+    estimate_norm,
+    inverse_estimate_norm,
+    transform_landmark_points,
+    get_lm,
+)
 from networks.layers import AdaIN, AdaptiveAttention
 from tensorflow_addons.layers import InstanceNormalization
 import numpy as np
@@ -13,42 +19,77 @@ from tensorflow.keras.models import load_model
 from options.swap_options import SwapOptions
 
 
-
-token = os.environ['model_fetch']
+token = os.environ["model_fetch"]
 
 opt = SwapOptions().parse()
 
-retina_repo = Repository(local_dir="retina_model", clone_from="felixrosberg/retinaface_resnet50",
-                         private=True, use_auth_token=token, git_user="felixrosberg")
+retina_repo = Repository(
+    local_dir="retina_model",
+    clone_from="felixrosberg/retinaface_resnet50",
+    private=True,
+    use_auth_token=token,
+    git_user="felixrosberg",
+)
 
 from retina_model.models import *
 
-RetinaFace = load_model("retina_model/retinaface_res50.h5",
-                        custom_objects={"FPN": FPN,
-                                        "SSH": SSH,
-                                        "BboxHead": BboxHead,
-                                        "LandmarkHead": LandmarkHead,
-                                        "ClassHead": ClassHead})
-
-arc_repo = Repository(local_dir="arcface_model", clone_from="felixrosberg/arcface_tf",
-                      private=True, use_auth_token=token)
+RetinaFace = load_model(
+    "retina_model/retinaface_res50.h5",
+    custom_objects={
+        "FPN": FPN,
+        "SSH": SSH,
+        "BboxHead": BboxHead,
+        "LandmarkHead": LandmarkHead,
+        "ClassHead": ClassHead,
+    },
+)
+
+arc_repo = Repository(
+    local_dir="arcface_model",
+    clone_from="felixrosberg/arcface_tf",
+    private=True,
+    use_auth_token=token,
+)
 ArcFace = load_model("arcface_model/arc_res50.h5")
 ArcFaceE = load_model("arcface_model/arc_res50e.h5")
 
-g_repo = Repository(local_dir="g_model_c_hq", clone_from="felixrosberg/affa_config_c_hq",
-                    private=True, use_auth_token=token)
-G = load_model("g_model_c_hq/generator_t_28.h5", custom_objects={"AdaIN": AdaIN,
-                                                                 "AdaptiveAttention": AdaptiveAttention,
-                                                                 "InstanceNormalization": InstanceNormalization})
-
-r_repo = Repository(local_dir="reconstruction_attack", clone_from="felixrosberg/reconstruction_attack",
-                    private=True, use_auth_token=token)
-R = load_model("reconstruction_attack/reconstructor_42.h5", custom_objects={"AdaIN": AdaIN,
-                                                                            "AdaptiveAttention": AdaptiveAttention,
-                                                                            "InstanceNormalization": InstanceNormalization})
-
-permuter_repo = Repository(local_dir="identity_permuter", clone_from="felixrosberg/identitypermuter",
-                           private=True, use_auth_token=token, git_user="felixrosberg")
+g_repo = Repository(
+    local_dir="g_model_c_hq",
+    clone_from="felixrosberg/affa_config_c_hq",
+    private=True,
+    use_auth_token=token,
+)
+G = load_model(
+    "g_model_c_hq/generator_t_28.h5",
+    custom_objects={
+        "AdaIN": AdaIN,
+        "AdaptiveAttention": AdaptiveAttention,
+        "InstanceNormalization": InstanceNormalization,
+    },
+)
+
+r_repo = Repository(
+    local_dir="reconstruction_attack",
+    clone_from="felixrosberg/reconstruction_attack",
+    private=True,
+    use_auth_token=token,
+)
+R = load_model(
+    "reconstruction_attack/reconstructor_42.h5",
+    custom_objects={
+        "AdaIN": AdaIN,
+        "AdaptiveAttention": AdaptiveAttention,
+        "InstanceNormalization": InstanceNormalization,
+    },
+)
+
+permuter_repo = Repository(
+    local_dir="identity_permuter",
+    clone_from="felixrosberg/identitypermuter",
+    private=True,
+    use_auth_token=token,
+    git_user="felixrosberg",
+)
 
 from identity_permuter.id_permuter import identity_permuter
 
@@ -71,7 +112,11 @@ def run_inference(target, source, slider, adv_slider, settings):
             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))
+            source_z = ArcFace.predict(
+                np.expand_dims(
+                    tf.image.resize(source_aligned, [112, 112]) / 255.0, axis=0
+                )
+            )
         else:
             source_z = None
 
@@ -82,22 +127,31 @@ def run_inference(target, source, slider, adv_slider, settings):
 
         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()
+        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)
+            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) - 127.5) / 127.5
+            im_aligned = (
+                cv2.warpAffine(im, M, (256, 256), borderValue=0.0) - 127.5
+            ) / 127.5
 
             if "adversarial defense" in settings:
                 eps = adv_slider / 200
@@ -125,54 +179,76 @@ def run_inference(target, source, slider, adv_slider, settings):
 
                 slider_weight = slider / 100
 
-                target_z = ArcFace.predict(np.expand_dims(tf.image.resize(im_aligned, [112, 112]) * 0.5 + 0.5, axis=0))
+                target_z = ArcFace.predict(
+                    np.expand_dims(
+                        tf.image.resize(im_aligned, [112, 112]) * 0.5 + 0.5, axis=0
+                    )
+                )
                 source_z = IDP.predict(target_z)
 
                 source_z = slider_weight * source_z + (1 - slider_weight) * target_z
 
             if "reconstruction attack" in settings:
-                source_z = ArcFaceE.predict(np.expand_dims(tf.image.resize(im_aligned, [112, 112]) * 0.5 + 0.5, axis=0))
+                source_z = ArcFaceE.predict(
+                    np.expand_dims(
+                        tf.image.resize(im_aligned, [112, 112]) * 0.5 + 0.5, axis=0
+                    )
+                )
 
             # face swap
             if "reconstruction attack" not in settings:
-                changed_face_cage = G.predict([np.expand_dims(im_aligned, axis=0),
-                                               source_z])
+                changed_face_cage = G.predict(
+                    [np.expand_dims(im_aligned, axis=0), source_z]
+                )
                 changed_face = changed_face_cage[0] * 0.5 + 0.5
 
                 # 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)
+                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 = 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))
+                total_img = iim_aligned * blend_mask + total_img * (1 - blend_mask)
             else:
-                changed_face_cage = R.predict([np.expand_dims(im_aligned, axis=0),
-                                               source_z])
+                changed_face_cage = R.predict(
+                    [np.expand_dims(im_aligned, axis=0), source_z]
+                )
                 changed_face = changed_face_cage[0] * 0.5 + 0.5
 
                 # 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)
+                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 = 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))
+                total_img = iim_aligned * blend_mask + total_img * (1 - blend_mask)
 
         if "compare" in settings:
             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')
+        total_img = total_img.astype("uint8")
 
         return total_img
     except Exception as e:
@@ -180,38 +256,47 @@ def run_inference(target, source, slider, adv_slider, settings):
         return None
 
 
-description = "Performs subject agnostic identity transfer from a source face to all target faces. \n\n" \
-              "Implementation and demo of FaceDancer, accepted to WACV 2023. \n\n" \
-              "Pre-print: https://arxiv.org/abs/2210.10473 \n\n" \
-              "Code: https://github.com/felixrosberg/FaceDancer \n\n" \
-               "\n\n" \
-              "Options:\n\n" \
-              "-Compare returns the target image concatenated with the results.\n\n" \
-              "-Anonymize will ignore the source image and perform an identity permutation of target faces.\n\n" \
-              "-Reconstruction attack will attempt to invert the face swap or the anonymization.\n\n" \
-              "-Adversarial defense will add a permutation noise that disrupts the reconstruction attack.\n\n" \
-              "NOTE: There is no guarantees with the anonymization process currently.\n\n" \
-              "NOTE: source image with too high resolution may not work properly!"
-examples = [["assets/rick.jpg", "assets/musk.jpg", 100, 10, ["compare"]],
-            ["assets/musk.jpg", "assets/musk.jpg", 100, 10, ["anonymize"]]]
+description = (
+    "Performs subject agnostic identity transfer from a source face to all target faces. \n\n"
+    "Implementation and demo of FaceDancer, accepted to WACV 2023. \n\n"
+    "Pre-print: https://arxiv.org/abs/2210.10473 \n\n"
+    "Code: https://github.com/felixrosberg/FaceDancer \n\n"
+    "\n\n"
+    "Options:\n\n"
+    "-Compare returns the target image concatenated with the results.\n\n"
+    "-Anonymize will ignore the source image and perform an identity permutation of target faces.\n\n"
+    "-Reconstruction attack will attempt to invert the face swap or the anonymization.\n\n"
+    "-Adversarial defense will add a permutation noise that disrupts the reconstruction attack.\n\n"
+    "NOTE: There is no guarantees with the anonymization process currently.\n\n"
+    "NOTE: source image with too high resolution may not work properly!"
+)
+examples = [
+    ["assets/rick.jpg", "assets/musk.jpg", 100, 10, ["compare"]],
+    ["assets/musk.jpg", "assets/musk.jpg", 100, 10, ["anonymize"]],
+]
 article = """
 Demo is based of recent research from my Ph.D work. Results expects to be published in the coming months.
 """
 
-iface = gradio.Interface(run_inference,
-                         [gradio.inputs.Image(shape=None, label='Target'),
-                          gradio.inputs.Image(shape=None, label='Source'),
-                          gradio.inputs.Slider(0, 100, default=100, label="Anonymization ratio (%)"),
-                          gradio.inputs.Slider(0, 100, default=100, label="Adversarial defense ratio (%)"),
-                          gradio.inputs.CheckboxGroup(["compare",
-                                                       "anonymize",
-                                                       "reconstruction attack",
-                                                       "adversarial defense"],
-                                                      label='Options')],
-                         gradio.outputs.Image(),
-                         title="Face Swap",
-                         description=description,
-                         examples=examples,
-                         article=article,
-                         layout="vertical")
+iface = gradio.Interface(
+    run_inference,
+    [
+        gradio.inputs.Image(shape=None, label="Target"),
+        gradio.inputs.Image(shape=None, label="Source"),
+        gradio.inputs.Slider(0, 100, default=100, label="Anonymization ratio (%)"),
+        gradio.inputs.Slider(
+            0, 100, default=100, label="Adversarial defense ratio (%)"
+        ),
+        gradio.inputs.CheckboxGroup(
+            ["compare", "anonymize", "reconstruction attack", "adversarial defense"],
+            label="Options",
+        ),
+    ],
+    gradio.outputs.Image(),
+    title="Face Swap",
+    description=description,
+    examples=examples,
+    article=article,
+    layout="vertical",
+)
 iface.launch()