Spaces:
Running
Running
File size: 7,977 Bytes
3eebbc1 |
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 |
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Date : 2020-05-17
# @Author : Shawn Shan (shansixiong@cs.uchicago.edu)
# @Link : https://www.shawnshan.com/
import argparse
import glob
import logging
import os
import sys
logging.getLogger('tensorflow').setLevel(logging.ERROR)
os.environ["KMP_AFFINITY"] = "noverbose"
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
import tensorflow as tf
tf.get_logger().setLevel('ERROR')
tf.autograph.set_verbosity(3)
import numpy as np
from fawkes.differentiator import FawkesMaskGeneration
from fawkes.utils import init_gpu, reverse_process_cloaked, \
Faces, load_extractor
from fawkes.align_face import aligner
def generate_cloak_images(protector, image_X, target_emb=None):
cloaked_image_X = protector.compute(image_X, target_emb)
return cloaked_image_X
IMG_SIZE = 112
PREPROCESS = 'raw'
class Fawkes(object):
def __init__(self, gpu, batch_size, mode="low"):
self.gpu = gpu
self.batch_size = batch_size
self.mode = mode
th, max_step, lr, extractors = self.mode2param(self.mode)
self.th = th
self.lr = lr
self.max_step = max_step
if gpu is not None:
init_gpu(gpu)
self.aligner = aligner()
self.protector = None
self.protector_param = None
self.feature_extractors_ls = [load_extractor(name) for name in extractors]
def mode2param(self, mode):
if mode == 'low':
th = 0.004
max_step = 40
lr = 25
extractors = ["extractor_2"]
elif mode == 'mid':
th = 0.012
max_step = 75
lr = 20
extractors = ["extractor_0", "extractor_2"]
elif mode == 'high':
th = 0.017
max_step = 150
lr = 15
extractors = ["extractor_0", "extractor_2"]
else:
raise Exception("mode must be one of 'min', 'low', 'mid', 'high'")
return th, max_step, lr, extractors
def run_protection(self, image, sd=1e7, batch_size=1, format='png', separate_target=True, debug=False,
no_align=False, maximize=True, save_last_on_failed=True):
current_param = "-".join([str(x) for x in [self.th, sd, self.lr, self.max_step, batch_size, format,
separate_target, debug]])
faces = Faces(image, self.aligner, verbose=1, no_align=no_align)
original_images = faces.cropped_faces
if len(original_images) == 0:
print("No face detected. ")
return 2
original_images = np.array(original_images)
if current_param != self.protector_param:
self.protector_param = current_param
if self.protector is not None:
del self.protector
if batch_size == -1:
batch_size = len(original_images)
self.protector = FawkesMaskGeneration(self.feature_extractors_ls,
batch_size=batch_size,
mimic_img=True,
intensity_range=PREPROCESS,
initial_const=sd,
learning_rate=self.lr,
max_iterations=self.max_step,
l_threshold=self.th,
verbose=debug,
maximize=maximize,
keep_final=False,
image_shape=(IMG_SIZE, IMG_SIZE, 3),
loss_method='features',
tanh_process=True,
save_last_on_failed=save_last_on_failed,
)
protected_images = generate_cloak_images(self.protector, original_images)
faces.cloaked_cropped_faces = protected_images
final_images, images_without_face = faces.merge_faces(
reverse_process_cloaked(protected_images, preprocess=PREPROCESS),
reverse_process_cloaked(original_images, preprocess=PREPROCESS))
if images_without_face:
return None
else:
return [img for img in final_images][0]
# for i in range(len(final_images)):
# if i in images_without_face:
# continue
# p_img = final_images[i]
# path = image_paths[i]
# file_name = "{}_cloaked.{}".format(".".join(path.split(".")[:-1]), format)
# dump_image(p_img, file_name, format=format)
# print("Done!")
# return 1
def main(*argv):
if not argv:
argv = list(sys.argv)
try:
import signal
signal.signal(signal.SIGPIPE, signal.SIG_DFL)
except Exception as e:
pass
parser = argparse.ArgumentParser()
parser.add_argument('--directory', '-d', type=str,
help='the directory that contains images to run protection', default='imgs/')
parser.add_argument('--gpu', '-g', type=str,
help='the GPU id when using GPU for optimization', default='0')
parser.add_argument('--mode', '-m', type=str,
help='cloak generation mode, select from min, low, mid, high. The higher the mode is, '
'the more perturbation added and stronger protection',
default='low')
parser.add_argument('--feature-extractor', type=str,
help="name of the feature extractor used for optimization",
default="arcface_extractor_0")
parser.add_argument('--th', help='only relevant with mode=custom, DSSIM threshold for perturbation', type=float,
default=0.01)
parser.add_argument('--max-step', help='only relevant with mode=custom, number of steps for optimization', type=int,
default=1000)
parser.add_argument('--sd', type=int, help='only relevant with mode=custom, penalty number, read more in the paper',
default=1e6)
parser.add_argument('--lr', type=float, help='only relevant with mode=custom, learning rate', default=2)
parser.add_argument('--batch-size', help="number of images to run optimization together", type=int, default=1)
parser.add_argument('--separate_target', help="whether select separate targets for each faces in the directory",
action='store_true')
parser.add_argument('--no-align', help="whether to detect and crop faces",
action='store_true')
parser.add_argument('--debug', help="turn on debug and copy/paste the stdout when reporting an issue on github",
action='store_true')
parser.add_argument('--format', type=str,
help="format of the output image",
default="png")
args = parser.parse_args(argv[1:])
assert args.format in ['png', 'jpg', 'jpeg']
if args.format == 'jpg':
args.format = 'jpeg'
image_paths = glob.glob(os.path.join(args.directory, "*"))
image_paths = [path for path in image_paths if "_cloaked" not in path.split("/")[-1]]
protector = Fawkes(args.gpu, args.batch_size, mode=args.mode)
protector.run_protection(image_paths, th=args.th, sd=args.sd, lr=args.lr,
max_step=args.max_step,
batch_size=args.batch_size, format=args.format,
separate_target=args.separate_target, debug=args.debug, no_align=args.no_align)
if __name__ == '__main__':
main(*sys.argv)
|