Update app.py

app.py

@@ -1,69 +1,215 @@
-from flask import Flask, render_template, request, redirect, url_for, send_from_directory
+import gradio
+from huggingface_hub import Repository
 import os
-import subprocess
-import uuid  # Import the UUID module
-from werkzeug.utils import secure_filename
-
-app = Flask(__name__)
-
-# Configuration for file uploads and output
-UPLOAD_FOLDER = 'uploads'
-OUTPUT_FOLDER = 'output_files'
-ALLOWED_EXTENSIONS = {'jpg', 'jpeg', 'png', 'gif', 'mp4', 'avi', 'mov'}
-
-app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
-app.config['OUTPUT_FOLDER'] = OUTPUT_FOLDER
-
-def allowed_file(filename):
-    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
-
-@app.route('/', methods=['GET', 'POST'])
-def index():
-    if request.method == 'POST':
-        # Handle file uploads and processing options
-        source_file = request.files['source']
-        target_file = request.files['target']
-        frame_processors = request.form.getlist('frame_processor')
-        
-        if source_file and allowed_file(source_file.filename) and target_file and allowed_file(target_file.filename):
-            # Generate unique filenames with UUIDs for uploaded files
-            source_filename = str(uuid.uuid4()) + '_' + secure_filename(source_file.filename)
-            target_filename = str(uuid.uuid4()) + '_' + secure_filename(target_file.filename)
-            
-            # Save uploaded files
-            source_path = os.path.join(app.config['UPLOAD_FOLDER'], source_filename)
-            target_path = os.path.join(app.config['UPLOAD_FOLDER'], target_filename)
-            source_file.save(source_path)
-            target_file.save(target_path)
-            
-            # Determine output file name with UUID
-            output_filename = str(uuid.uuid4()) + '.jpg'  # Default output format is JPEG
-            
-            # Build and execute the processing command here
-            processing_command = ['python', 'run.py', '-s', source_path, '-t', target_path, '-o', os.path.join(app.config['OUTPUT_FOLDER'], output_filename)]
-            processing_command.extend(['--frame-processor', *frame_processors])
-            
-            try:
-                # Execute the processing command using subprocess
-                subprocess.run(processing_command, check=True)
-                
-                # Redirect to the output page
-                return redirect(url_for('output', filename=output_filename))
-            except subprocess.CalledProcessError:
-                return render_template('error.html')
-    
-    return render_template('index.html')
-
-@app.route('/output/<filename>')
-def output(filename):
-    return render_template('output.html', filename=filename)
-
-@app.route('/download/<filename>')
-def download(filename):
-    return send_from_directory(app.config['OUTPUT_FOLDER'], filename, as_attachment=True)
-
-if __name__ == '__main__':
-    os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
-    os.makedirs(app.config['OUTPUT_FOLDER'], exist_ok=True)
-    app.run(host="0.0.0.0", port=7860, debug=True)
-    
+
+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
+import cv2
+from scipy.ndimage import gaussian_filter
+
+from tensorflow.keras.models import load_model
+from options.swap_options import SwapOptions
+
+# .
+# token = os.environ['model_fetch']
+
+opt = SwapOptions().parse()
+token = os.environ['token']
+
+retina_repo = Repository(local_dir="retina_models", clone_from="felixrosberg/RetinaFace")
+
+from retinaface.models import *
+
+RetinaFace = load_model("retina_models/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")
+ArcFace = load_model("arcface_model/ArcFace-Res50.h5")
+ArcFaceE = load_model("arcface_model/ArcFacePerceptual-Res50.h5")
+
+g_repo = Repository(local_dir="g_model_c_hq", clone_from="felixrosberg/FaceDancer",use_auth_token=token)
+G = load_model("g_model_c_hq/FaceDancer_config_c_HQ.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
+
+# IDP = identity_permuter(emb_size=32, min_arg=False)
+# IDP.load_weights("identity_permuter/id_permuter.h5")
+
+blend_mask_base = np.zeros(shape=(256, 256, 1))
+blend_mask_base[80:244, 32:224] = 1
+blend_mask_base = gaussian_filter(blend_mask_base, sigma=7)
+
+
+def run_inference(target, source, slider, adv_slider, settings):
+    try:
+        source = np.array(source)
+        target = np.array(target)
+
+        # Prepare to load video
+        if "anonymize" not in settings:
+            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))
+        else:
+            source_z = None
+
+        # 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) - 127.5) / 127.5
+
+            if "adversarial defense" in settings:
+                eps = adv_slider / 200
+                X = tf.convert_to_tensor(np.expand_dims(im_aligned, axis=0))
+                with tf.GradientTape() as tape:
+                    tape.watch(X)
+
+                    X_z = ArcFaceE(tf.image.resize(X * 0.5 + 0.5, [112, 112]))
+                    output = R([X, X_z])
+
+                    loss = tf.reduce_mean(tf.abs(0 - output))
+
+                gradient = tf.sign(tape.gradient(loss, X))
+
+                adv_x = X + eps * gradient
+                im_aligned = tf.clip_by_value(adv_x, -1, 1)[0]
+
+            if "anonymize" in settings and "reconstruction attack" not in settings:
+                """source_z = ArcFace.predict(np.expand_dims(tf.image.resize(im_aligned, [112, 112]) / 255.0, axis=0))
+                anon_ratio = int(512 * (slider / 100))
+                anon_vector = np.ones(shape=(1, 512))
+                anon_vector[:, :anon_ratio] = -1
+                np.random.shuffle(anon_vector)
+                source_z *= anon_vector"""
+
+                slider_weight = slider / 100
+
+                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))
+
+            # face swap
+            if "reconstruction attack" not in settings:
+                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)
+
+                # 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))
+            else:
+                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)
+
+                # 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 "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')
+
+        return total_img
+    except Exception as e:
+        print(e)
+        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"]]]
+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.Image(shape=None, type="pil", label='Target'),
+                          gradio.Image(shape=None, type="pil", label='Source'),
+                          gradio.Slider(0, 100, default=100, label="Anonymization ratio (%)"),
+                          gradio.Slider(0, 100, default=100, label="Adversarial defense ratio (%)"),
+                          gradio.CheckboxGroup(["compare",
+                                                       "anonymize",
+                                                       "reconstruction attack",
+                                                       "adversarial defense"],
+                                                      label='Options')],
+                         "image",
+                         title="Face Swap",
+                         description=description,
+                         examples=examples,
+                         article=article,
+                         layout="vertical")
+iface.launch()