File size: 9,344 Bytes
e88171f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8acef0e
e88171f
82cef28
e88171f
82cef28
e88171f
 
 
 
 
 
 
8acef0e
82cef28
 
e88171f
8acef0e
82cef28
e88171f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5f2270e
daa2ac3
f5c17d6
 
 
 
e88171f
 
 
 
 
 
 
 
 
 
 
daa2ac3
f5c17d6
 
 
e88171f
efafb14
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
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 = "Not Working"
examples = []
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="Not Working",
                         description=description,
                         examples=examples,
                         article=article,
                         layout="vertical")
iface.launch()