File size: 10,431 Bytes
c873df0
 
 
 
 
 
 
 
 
 
 
 
 
 
8547084
 
c873df0
 
8547084
c873df0
dea2840
c873df0
d9edcea
c873df0
8547084
c873df0
 
 
 
8547084
 
3a17234
dea2840
 
 
c873df0
dea2840
 
c873df0
 
 
8547084
 
dea2840
 
 
c873df0
8547084
 
c873df0
dea2840
c873df0
dea2840
 
c873df0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8547084
c873df0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
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 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()