fawkes/fawkes_dev/eval_cloak.py

import argparse
import os
import sys

import numpy as np

sys.path.append("/home/shansixioing/fawkes/fawkes")
from utils import extract_faces, get_dataset_path, init_gpu, load_extractor, load_victim_model

import random
import glob
from keras.preprocessing import image
from keras.utils import to_categorical
from keras.applications.vgg16 import preprocess_input


def select_samples(data_dir):
    all_data_path = []
    for cls in os.listdir(data_dir):
        cls_dir = os.path.join(data_dir, cls)
        for data_path in os.listdir(cls_dir):
            all_data_path.append(os.path.join(cls_dir, data_path))
    return all_data_path


def generator_wrap(protect_images, test=False, validation_split=0.1):
    train_data_dir, test_data_dir, num_classes, num_images = get_dataset_path(args.dataset)

    idx = 0
    path2class = {}
    path2imgs_list = {}

    for target_path in sorted(glob.glob(train_data_dir + "/*")):
        path2class[target_path] = idx
        path2imgs_list[target_path] = glob.glob(os.path.join(target_path, "*"))
        idx += 1
        if idx >= args.num_classes:
            break

    path2class["protected"] = idx

    np.random.seed(12345)

    while True:
        batch_X = []
        batch_Y = []
        cur_batch_path = np.random.choice(list(path2class.keys()), args.batch_size)
        for p in cur_batch_path:
            cur_y = path2class[p]
            if test and p == 'protected':
                continue
            # protect class images in train dataset
            elif p == 'protected':
                cur_x = random.choice(protect_images)
            else:
                cur_path = random.choice(path2imgs_list[p])
                im = image.load_img(cur_path, target_size=(224, 224))
                cur_x = image.img_to_array(im)

            cur_x = preprocess_input(cur_x)
            batch_X.append(cur_x)
            batch_Y.append(cur_y)

        batch_X = np.array(batch_X)
        batch_Y = to_categorical(np.array(batch_Y), num_classes=args.num_classes + 1)

        yield batch_X, batch_Y


def eval_uncloaked_test_data(cloak_data, n_classes):
    original_label = cloak_data.path2idx[list(cloak_data.protect_class_path)[0]]
    protect_test_X = cloak_data.protect_test_X
    original_Y = [original_label] * len(protect_test_X)
    original_Y = to_categorical(original_Y, n_classes)
    return protect_test_X, original_Y


def eval_cloaked_test_data(cloak_data, n_classes, validation_split=0.1):
    split = int(len(cloak_data.cloaked_protect_train_X) * (1 - validation_split))
    cloaked_test_X = cloak_data.cloaked_protect_train_X[split:]
    original_label = cloak_data.path2idx[list(cloak_data.protect_class_path)[0]]
    original_Y = [original_label] * len(cloaked_test_X)
    original_Y = to_categorical(original_Y, n_classes)
    return cloaked_test_X, original_Y


def main():
    init_gpu(args.gpu)
    #
    # if args.dataset == 'pubfig':
    #     N_CLASSES = 65
    #     CLOAK_DIR = args.cloak_data
    # elif args.dataset == 'scrub':
    #     N_CLASSES = 530
    #     CLOAK_DIR = args.cloak_data
    # else:
    #     raise ValueError

    print("Build attacker's model")

    image_paths = glob.glob(os.path.join(args.directory, "*"))
    original_image_paths = sorted([path for path in image_paths if "_cloaked" not in path.split("/")[-1]])

    protect_image_paths = sorted([path for path in image_paths if "_cloaked" in path.split("/")[-1]])

    original_imgs = np.array([extract_faces(image.img_to_array(image.load_img(cur_path))) for cur_path in
                     original_image_paths[:150]])
    original_y = to_categorical([args.num_classes] * len(original_imgs), num_classes=args.num_classes + 1)

    protect_imgs = [extract_faces(image.img_to_array(image.load_img(cur_path))) for cur_path in
                    protect_image_paths]

    train_generator = generator_wrap(protect_imgs,
                                     validation_split=args.validation_split)
    test_generator = generator_wrap(protect_imgs, test=True,
                                    validation_split=args.validation_split)

    base_model = load_extractor(args.transfer_model)
    model = load_victim_model(teacher_model=base_model, number_classes=args.num_classes + 1)

    # cloaked_test_X, cloaked_test_Y = eval_cloaked_test_data(cloak_data, args.num_classes,
    #                                                         validation_split=args.validation_split)

    # try:
    train_data_dir, test_data_dir, num_classes, num_images = get_dataset_path(args.dataset)
    model.fit_generator(train_generator, steps_per_epoch=num_images // 32,
                        validation_data=(original_imgs, original_y),
                        epochs=args.n_epochs,
                        verbose=1,
                        use_multiprocessing=True, workers=5)
    # except KeyboardInterrupt:
    #     pass

    _, acc_original = model.evaluate(original_imgs, original_y, verbose=0)
    print("Accuracy on uncloaked/original images TEST: {:.4f}".format(acc_original))
    # EVAL_RES['acc_original'] = acc_original

    _, other_acc = model.evaluate_generator(test_generator, verbose=0, steps=50)
    print("Accuracy on other classes {:.4f}".format(other_acc))
    # EVAL_RES['other_acc'] = other_acc
    # dump_dictionary_as_json(EVAL_RES, os.path.join(CLOAK_DIR, "eval_seed{}.json".format(args.seed_idx)))


def parse_arguments(argv):
    parser = argparse.ArgumentParser()

    parser.add_argument('--gpu', type=str,
                        help='GPU id', default='0')

    parser.add_argument('--dataset', type=str,
                        help='name of dataset', default='scrub')
    parser.add_argument('--num_classes', type=int,
                        help='name of dataset', default=520)

    parser.add_argument('--directory', '-d', type=str,
                        help='name of the cloak result directory',
                        default='img/')

    parser.add_argument('--transfer_model', type=str,
                        help='the feature extractor used for tracker model training. ', default='low_extract')
    parser.add_argument('--batch_size', type=int, default=32)
    parser.add_argument('--validation_split', type=float, default=0.1)
    parser.add_argument('--n_epochs', type=int, default=3)
    return parser.parse_args(argv)


if __name__ == '__main__':
    args = parse_arguments(sys.argv[1:])
    main()