Full Code of izikgo/AnomalyDetectionTransformations for AI

Repository: izikgo/AnomalyDetectionTransformations
Branch: master
Commit: 6327041c6377
Files: 14
Total size: 60.8 KB

Directory structure:
gitextract__9qwn17c/

├── .gitignore
├── LICENSE
├── README.md
├── experiments.py
├── intuition_experiments.py
├── models/
│   ├── __init__.py
│   ├── adgan.py
│   ├── dagmm.py
│   ├── dsebm.py
│   ├── encoders_decoders.py
│   └── wide_residual_network.py
├── multiclass_experiment.py
├── transformations.py
└── utils.py

================================================
FILE CONTENTS
================================================

================================================
FILE: .gitignore
================================================
# Custom
*.png
*.h5
*.npz
*.csv
*.zip

# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class

# C extensions
*.so

# Distribution / packaging
.Python
env/
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
*.egg-info/
.installed.cfg
*.egg

# PyInstaller
#  Usually these files are written by a python script from a template
#  before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec

# Installer logs
pip-log.txt
pip-delete-this-directory.txt

# Unit test / coverage reports
htmlcov/
.tox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*,cover
.hypothesis/

# Translations
*.mo
*.pot

# Django stuff:
*.log
local_settings.py

# Flask stuff:
instance/
.webassets-cache

# Scrapy stuff:
.scrapy

# Sphinx documentation
docs/_build/

# PyBuilder
target/

# IPython Notebook
.ipynb_checkpoints

# pyenv
.python-version

# celery beat schedule file
celerybeat-schedule

# dotenv
.env

# virtualenv
venv/
ENV/

# Spyder project settings
.spyderproject

# Rope project settings
.ropeproject


================================================
FILE: LICENSE
================================================
MIT License

Copyright (c) 2018 izikgo

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.


================================================
FILE: README.md
================================================
# Deep Anomaly Detection Using Geometric Transformations
To be presented in NIPS 2018 by Izhak Golan and Ran El-Yaniv.

## Introduction
This is the official implementation of "Deep Anomaly Detection Using Geometric Transformations".
It includes all experiments reported in the paper.

## Requirements
* Python 3.5+
* Keras 2.2.0
* Tensorflow 1.8.0
* sklearn 0.19.1

## Citation
If you use the ideas or method presented in the paper, please cite:

```
@inproceedings{NEURIPS2018_5e62d03a,
 author = {Golan, Izhak and El-Yaniv, Ran},
 booktitle = {Advances in Neural Information Processing Systems},
 editor = {S. Bengio and H. Wallach and H. Larochelle and K. Grauman and N. Cesa-Bianchi and R. Garnett},
 pages = {},
 publisher = {Curran Associates, Inc.},
 title = {Deep Anomaly Detection Using Geometric Transformations},
 url = {https://proceedings.neurips.cc/paper/2018/file/5e62d03aec0d17facfc5355dd90d441c-Paper.pdf},
 volume = {31},
 year = {2018}
}
```


================================================
FILE: experiments.py
================================================
import os
import csv
from collections import defaultdict
from glob import glob
from datetime import datetime
from multiprocessing import Manager, freeze_support, Process
import numpy as np
import scipy.stats
from scipy.special import psi, polygamma
from sklearn.metrics import roc_auc_score
from sklearn.svm import OneClassSVM
from sklearn.model_selection import ParameterGrid
from sklearn.externals.joblib import Parallel, delayed
from keras.models import Model, Input, Sequential
from keras.layers import Dense, Dropout
from keras.utils import to_categorical
from utils import load_cifar10, load_cats_vs_dogs, load_fashion_mnist, load_cifar100
from utils import save_roc_pr_curve_data, get_class_name_from_index, get_channels_axis
from transformations import Transformer
from models.wide_residual_network import create_wide_residual_network
from models.encoders_decoders import conv_encoder, conv_decoder
from models import dsebm, dagmm, adgan
import keras.backend as K

RESULTS_DIR = ''


def _transformations_experiment(dataset_load_fn, dataset_name, single_class_ind, gpu_q):
    gpu_to_use = gpu_q.get()
    os.environ["CUDA_VISIBLE_DEVICES"] = gpu_to_use

    (x_train, y_train), (x_test, y_test) = dataset_load_fn()

    if dataset_name in ['cats-vs-dogs']:
        transformer = Transformer(16, 16)
        n, k = (16, 8)
    else:
        transformer = Transformer(8, 8)
        n, k = (10, 4)
    mdl = create_wide_residual_network(x_train.shape[1:], transformer.n_transforms, n, k)
    mdl.compile('adam',
                'categorical_crossentropy',
                ['acc'])

    x_train_task = x_train[y_train.flatten() == single_class_ind]
    transformations_inds = np.tile(np.arange(transformer.n_transforms), len(x_train_task))
    x_train_task_transformed = transformer.transform_batch(np.repeat(x_train_task, transformer.n_transforms, axis=0),
                                                           transformations_inds)
    batch_size = 128

    mdl.fit(x=x_train_task_transformed, y=to_categorical(transformations_inds),
            batch_size=batch_size, epochs=int(np.ceil(200/transformer.n_transforms))
            )

    #################################################################################################
    # simplified normality score
    #################################################################################################
    # preds = np.zeros((len(x_test), transformer.n_transforms))
    # for t in range(transformer.n_transforms):
    #     preds[:, t] = mdl.predict(transformer.transform_batch(x_test, [t] * len(x_test)),
    #                               batch_size=batch_size)[:, t]
    #
    # labels = y_test.flatten() == single_class_ind
    # scores = preds.mean(axis=-1)
    #################################################################################################

    def calc_approx_alpha_sum(observations):
        N = len(observations)
        f = np.mean(observations, axis=0)

        return (N * (len(f) - 1) * (-psi(1))) / (
                N * np.sum(f * np.log(f)) - np.sum(f * np.sum(np.log(observations), axis=0)))

    def inv_psi(y, iters=5):
        # initial estimate
        cond = y >= -2.22
        x = cond * (np.exp(y) + 0.5) + (1 - cond) * -1 / (y - psi(1))

        for _ in range(iters):
            x = x - (psi(x) - y) / polygamma(1, x)
        return x

    def fixed_point_dirichlet_mle(alpha_init, log_p_hat, max_iter=1000):
        alpha_new = alpha_old = alpha_init
        for _ in range(max_iter):
            alpha_new = inv_psi(psi(np.sum(alpha_old)) + log_p_hat)
            if np.sqrt(np.sum((alpha_old - alpha_new) ** 2)) < 1e-9:
                break
            alpha_old = alpha_new
        return alpha_new

    def dirichlet_normality_score(alpha, p):
        return np.sum((alpha - 1) * np.log(p), axis=-1)

    scores = np.zeros((len(x_test),))
    observed_data = x_train_task
    for t_ind in range(transformer.n_transforms):
        observed_dirichlet = mdl.predict(transformer.transform_batch(observed_data, [t_ind] * len(observed_data)),
                                         batch_size=1024)
        log_p_hat_train = np.log(observed_dirichlet).mean(axis=0)

        alpha_sum_approx = calc_approx_alpha_sum(observed_dirichlet)
        alpha_0 = observed_dirichlet.mean(axis=0) * alpha_sum_approx

        mle_alpha_t = fixed_point_dirichlet_mle(alpha_0, log_p_hat_train)

        x_test_p = mdl.predict(transformer.transform_batch(x_test, [t_ind] * len(x_test)),
                               batch_size=1024)
        scores += dirichlet_normality_score(mle_alpha_t, x_test_p)

    scores /= transformer.n_transforms
    labels = y_test.flatten() == single_class_ind

    res_file_name = '{}_transformations_{}_{}.npz'.format(dataset_name,
                                                 get_class_name_from_index(single_class_ind, dataset_name),
                                                 datetime.now().strftime('%Y-%m-%d-%H%M'))
    res_file_path = os.path.join(RESULTS_DIR, dataset_name, res_file_name)
    save_roc_pr_curve_data(scores, labels, res_file_path)

    mdl_weights_name = '{}_transformations_{}_{}_weights.h5'.format(dataset_name,
                                                           get_class_name_from_index(single_class_ind, dataset_name),
                                                           datetime.now().strftime('%Y-%m-%d-%H%M'))
    mdl_weights_path = os.path.join(RESULTS_DIR, dataset_name, mdl_weights_name)
    mdl.save_weights(mdl_weights_path)

    gpu_q.put(gpu_to_use)


def _train_ocsvm_and_score(params, xtrain, test_labels, xtest):
    return roc_auc_score(test_labels, OneClassSVM(**params).fit(xtrain).decision_function(xtest))


def _raw_ocsvm_experiment(dataset_load_fn, dataset_name, single_class_ind):
    (x_train, y_train), (x_test, y_test) = dataset_load_fn()

    x_train = x_train.reshape((len(x_train), -1))
    x_test = x_test.reshape((len(x_test), -1))

    x_train_task = x_train[y_train.flatten() == single_class_ind]
    if dataset_name in ['cats-vs-dogs']:  # OC-SVM is quadratic on the number of examples, so subsample training set
        subsample_inds = np.random.choice(len(x_train_task), 5000, replace=False)
        x_train_task = x_train_task[subsample_inds]

    pg = ParameterGrid({'nu': np.linspace(0.1, 0.9, num=9),
                        'gamma': np.logspace(-7, 2, num=10, base=2)})

    results = Parallel(n_jobs=6)(
        delayed(_train_ocsvm_and_score)(d, x_train_task, y_test.flatten() == single_class_ind, x_test)
        for d in pg)

    best_params, best_auc_score = max(zip(pg, results), key=lambda t: t[-1])
    best_ocsvm = OneClassSVM(**best_params).fit(x_train_task)
    scores = best_ocsvm.decision_function(x_test)
    labels = y_test.flatten() == single_class_ind

    res_file_name = '{}_raw-oc-svm_{}_{}.npz'.format(dataset_name,
                                                     get_class_name_from_index(single_class_ind, dataset_name),
                                                     datetime.now().strftime('%Y-%m-%d-%H%M'))
    res_file_path = os.path.join(RESULTS_DIR, dataset_name, res_file_name)
    save_roc_pr_curve_data(scores, labels, res_file_path)


def _cae_ocsvm_experiment(dataset_load_fn, dataset_name, single_class_ind, gpu_q):
    gpu_to_use = gpu_q.get()
    os.environ["CUDA_VISIBLE_DEVICES"] = gpu_to_use

    (x_train, y_train), (x_test, y_test) = dataset_load_fn()

    n_channels = x_train.shape[get_channels_axis()]
    input_side = x_train.shape[2]  # channel side will always be at shape[2]
    enc = conv_encoder(input_side, n_channels)
    dec = conv_decoder(input_side, n_channels)
    x_in = Input(shape=x_train.shape[1:])
    x_rec = dec(enc(x_in))
    cae = Model(x_in, x_rec)
    cae.compile('adam', 'mse')

    x_train_task = x_train[y_train.flatten() == single_class_ind]
    x_test_task = x_test[y_test.flatten() == single_class_ind]  # This is just for visual monitoring
    cae.fit(x=x_train_task, y=x_train_task, batch_size=128, epochs=200, validation_data=(x_test_task, x_test_task))

    x_train_task_rep = enc.predict(x_train_task, batch_size=128)
    if dataset_name in ['cats-vs-dogs']:  # OC-SVM is quadratic on the number of examples, so subsample training set
        subsample_inds = np.random.choice(len(x_train_task_rep), 2500, replace=False)
        x_train_task_rep = x_train_task_rep[subsample_inds]

    x_test_rep = enc.predict(x_test, batch_size=128)
    pg = ParameterGrid({'nu': np.linspace(0.1, 0.9, num=9),
                        'gamma': np.logspace(-7, 2, num=10, base=2)})

    results = Parallel(n_jobs=6)(
        delayed(_train_ocsvm_and_score)(d, x_train_task_rep, y_test.flatten() == single_class_ind, x_test_rep)
        for d in pg)

    best_params, best_auc_score = max(zip(pg, results), key=lambda t: t[-1])
    print(best_params)
    best_ocsvm = OneClassSVM(**best_params).fit(x_train_task_rep)
    scores = best_ocsvm.decision_function(x_test_rep)
    labels = y_test.flatten() == single_class_ind

    res_file_name = '{}_cae-oc-svm_{}_{}.npz'.format(dataset_name,
                                                     get_class_name_from_index(single_class_ind, dataset_name),
                                                     datetime.now().strftime('%Y-%m-%d-%H%M'))
    res_file_path = os.path.join(RESULTS_DIR, dataset_name, res_file_name)
    save_roc_pr_curve_data(scores, labels, res_file_path)

    gpu_q.put(gpu_to_use)


def _dsebm_experiment(dataset_load_fn, dataset_name, single_class_ind, gpu_q):
    gpu_to_use = gpu_q.get()
    os.environ["CUDA_VISIBLE_DEVICES"] = gpu_to_use

    (x_train, y_train), (x_test, y_test) = dataset_load_fn()

    n_channels = x_train.shape[get_channels_axis()]
    input_side = x_train.shape[2]  # image side will always be at shape[2]
    encoder_mdl = conv_encoder(input_side, n_channels, representation_activation='relu')
    energy_mdl = dsebm.create_energy_model(encoder_mdl)
    reconstruction_mdl = dsebm.create_reconstruction_model(energy_mdl)

    # optimization parameters
    batch_size = 128
    epochs = 200
    reconstruction_mdl.compile('adam', 'mse')
    x_train_task = x_train[y_train.flatten() == single_class_ind]
    x_test_task = x_test[y_test.flatten() == single_class_ind]  # This is just for visual monitoring
    reconstruction_mdl.fit(x=x_train_task, y=x_train_task,
                           batch_size=batch_size,
                           epochs=epochs,
                           validation_data=(x_test_task, x_test_task))

    scores = -energy_mdl.predict(x_test, batch_size)
    labels = y_test.flatten() == single_class_ind
    res_file_name = '{}_dsebm_{}_{}.npz'.format(dataset_name,
                                                get_class_name_from_index(single_class_ind, dataset_name),
                                                datetime.now().strftime('%Y-%m-%d-%H%M'))
    res_file_path = os.path.join(RESULTS_DIR, dataset_name, res_file_name)
    save_roc_pr_curve_data(scores, labels, res_file_path)

    gpu_q.put(gpu_to_use)


def _dagmm_experiment(dataset_load_fn, dataset_name, single_class_ind, gpu_q):
    gpu_to_use = gpu_q.get()
    os.environ["CUDA_VISIBLE_DEVICES"] = gpu_to_use

    (x_train, y_train), (x_test, y_test) = dataset_load_fn()

    n_channels = x_train.shape[get_channels_axis()]
    input_side = x_train.shape[2]  # image side will always be at shape[2]
    enc = conv_encoder(input_side, n_channels, representation_dim=5,
                       representation_activation='linear')
    dec = conv_decoder(input_side, n_channels=n_channels, representation_dim=enc.output_shape[-1])
    n_components = 3
    estimation = Sequential([Dense(64, activation='tanh', input_dim=enc.output_shape[-1] + 2), Dropout(0.5),
                             Dense(10, activation='tanh'), Dropout(0.5),
                             Dense(n_components, activation='softmax')]
                            )

    batch_size = 256
    epochs = 200
    lambda_diag = 0.0005
    lambda_energy = 0.01
    dagmm_mdl = dagmm.create_dagmm_model(enc, dec, estimation, lambda_diag)
    dagmm_mdl.compile('adam', ['mse', lambda y_true, y_pred: lambda_energy*y_pred])

    x_train_task = x_train[y_train.flatten() == single_class_ind]
    x_test_task = x_test[y_test.flatten() == single_class_ind]  # This is just for visual monitoring
    dagmm_mdl.fit(x=x_train_task, y=[x_train_task, np.zeros((len(x_train_task), 1))],  # second y is dummy
                  batch_size=batch_size,
                  epochs=epochs,
                  validation_data=(x_test_task, [x_test_task, np.zeros((len(x_test_task), 1))]),
                  # verbose=0
                  )

    energy_mdl = Model(dagmm_mdl.input, dagmm_mdl.output[-1])

    scores = -energy_mdl.predict(x_test, batch_size)
    scores = scores.flatten()
    if not np.all(np.isfinite(scores)):
        min_finite = np.min(scores[np.isfinite(scores)])
        scores[~np.isfinite(scores)] = min_finite - 1
    labels = y_test.flatten() == single_class_ind
    res_file_name = '{}_dagmm_{}_{}.npz'.format(dataset_name,
                                                get_class_name_from_index(single_class_ind, dataset_name),
                                                datetime.now().strftime('%Y-%m-%d-%H%M'))
    res_file_path = os.path.join(RESULTS_DIR, dataset_name, res_file_name)
    save_roc_pr_curve_data(scores, labels, res_file_path)

    gpu_q.put(gpu_to_use)


def _adgan_experiment(dataset_load_fn, dataset_name, single_class_ind, gpu_q):
    gpu_to_use = gpu_q.get()
    os.environ["CUDA_VISIBLE_DEVICES"] = gpu_to_use

    (x_train, y_train), (x_test, y_test) = dataset_load_fn()
    if len(x_test) > 5000:
        # subsample x_test due to runtime complexity
        chosen_inds = np.random.choice(len(x_test), 5000, replace=False)
        x_test = x_test[chosen_inds]
        y_test = y_test[chosen_inds]

    n_channels = x_train.shape[get_channels_axis()]
    input_side = x_train.shape[2]  # image side will always be at shape[2]
    critic = conv_encoder(input_side, n_channels, representation_dim=1,
                          representation_activation='linear')
    noise_size = 256
    generator = conv_decoder(input_side, n_channels=n_channels, representation_dim=noise_size)

    def prior_gen(b_size):
        return np.random.normal(size=(b_size, noise_size))

    batch_size = 128
    epochs = 100

    x_train_task = x_train[y_train.flatten() == single_class_ind]

    def data_gen(b_size):
        chosen_inds = np.random.choice(len(x_train_task), b_size, replace=False)
        return x_train_task[chosen_inds]

    adgan.train_wgan_with_grad_penalty(prior_gen, generator, data_gen, critic, batch_size, epochs, grad_pen_coef=20)

    scores = adgan.scores_from_adgan_generator(x_test, prior_gen, generator)
    labels = y_test.flatten() == single_class_ind
    res_file_name = '{}_adgan_{}_{}.npz'.format(dataset_name,
                                                get_class_name_from_index(single_class_ind, dataset_name),
                                                datetime.now().strftime('%Y-%m-%d-%H%M'))
    res_file_path = os.path.join(RESULTS_DIR, dataset_name, res_file_name)
    save_roc_pr_curve_data(scores, labels, res_file_path)

    gpu_q.put(gpu_to_use)


def run_experiments(load_dataset_fn, dataset_name, q, n_classes):

    # CAE OC-SVM
    processes = [Process(target=_cae_ocsvm_experiment,
                         args=(load_dataset_fn, dataset_name, c, q)) for c in range(n_classes)]
    for p in processes:
        p.start()
        p.join()

    # Raw OC-SVM
    for c in range(n_classes):
        _raw_ocsvm_experiment(load_dataset_fn, dataset_name, c)

    n_runs = 5

    # Transformations
    for _ in range(n_runs):
        processes = [Process(target=_transformations_experiment,
                             args=(load_dataset_fn, dataset_name, c, q)) for c in range(n_classes)]
        if dataset_name in ['cats-vs-dogs']:  # Self-labeled set is memory consuming
            for p in processes:
                p.start()
                p.join()
        else:
            for p in processes:
                p.start()
            for p in processes:
                p.join()

    # DSEBM
    for _ in range(n_runs):
        processes = [Process(target=_dsebm_experiment,
                             args=(load_dataset_fn, dataset_name, c, q)) for c in range(n_classes)]
        for p in processes:
            p.start()
        for p in processes:
            p.join()

    # DAGMM
    for _ in range(n_runs):
        processes = [Process(target=_dagmm_experiment,
                             args=(load_dataset_fn, dataset_name, c, q)) for c in range(n_classes)]
        for p in processes:
            p.start()
        for p in processes:
            p.join()

    # ADGAN
    processes = [Process(target=_adgan_experiment,
                         args=(load_dataset_fn, dataset_name, c, q)) for c in range(n_classes)]
    for p in processes:
        p.start()
    for p in processes:
        p.join()


def create_auc_table(metric='roc_auc'):
    file_path = glob(os.path.join(RESULTS_DIR, '*', '*.npz'))
    results = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
    methods = set()
    for p in file_path:
        _, f_name = os.path.split(p)
        dataset_name, method, single_class_name = f_name.split(sep='_')[:3]
        methods.add(method)
        npz = np.load(p)
        roc_auc = npz[metric]
        results[dataset_name][single_class_name][method].append(roc_auc)

    for ds_name in results:
        for sc_name in results[ds_name]:
            for method_name in results[ds_name][sc_name]:
                roc_aucs = results[ds_name][sc_name][method_name]
                results[ds_name][sc_name][method_name] = [np.mean(roc_aucs),
                                                          0 if len(roc_aucs) == 1 else scipy.stats.sem(np.array(roc_aucs))
                                                          ]

    with open('results-{}.csv'.format(metric), 'w') as csvfile:
        fieldnames = ['dataset', 'single class name'] + sorted(list(methods))
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

        writer.writeheader()
        for ds_name in sorted(results.keys()):
            for sc_name in sorted(results[ds_name].keys()):
                row_dict = {'dataset': ds_name, 'single class name': sc_name}
                row_dict.update({method_name: '{:.3f} ({:.3f})'.format(*results[ds_name][sc_name][method_name])
                                 for method_name in results[ds_name][sc_name]})
                writer.writerow(row_dict)


if __name__ == '__main__':
    freeze_support()
    N_GPUS = 2
    man = Manager()
    q = man.Queue(N_GPUS)
    for g in range(N_GPUS):
        q.put(str(g))

    experiments_list = [
        (load_cifar10, 'cifar10', 10),
        (load_cifar100, 'cifar100', 20),
        (load_fashion_mnist, 'fashion-mnist', 10),
        (load_cats_vs_dogs, 'cats-vs-dogs', 2),
    ]

    for data_load_fn, dataset_name, n_classes in experiments_list:
        run_experiments(data_load_fn, dataset_name, q, n_classes)



================================================
FILE: intuition_experiments.py
================================================
import itertools
import numpy as np
from scipy.ndimage.filters import gaussian_filter
import matplotlib.pyplot as plt
from transformations import SimpleTransformer
from utils import load_mnist
from keras.utils import to_categorical
from keras.layers import Flatten, Conv2D, Dense, BatchNormalization, MaxPool2D, Input, Lambda, average
from keras.models import Sequential, Model
import keras.backend as K
import tensorflow as tf


(x_train, y_train), (x_test, y_test) = load_mnist()
# scale to be in [0, 1]
x_train = (x_train + 1) / 2.
x_test = (x_test + 1) / 2.

single_class_ind = 3
anomaly_class_ind = 0

x_train_single = x_train[y_train == single_class_ind]
x_test_single = x_test[y_test == single_class_ind]
x_test_anomaly = x_test[y_test == anomaly_class_ind]

transformer = SimpleTransformer()
transformations_inds = np.tile(np.arange(transformer.n_transforms), len(x_train_single))
x_train_single_transformed = transformer.transform_batch(np.repeat(x_train_single, transformer.n_transforms, axis=0),
                                                         transformations_inds)


mdl = Sequential([Conv2D(32, (3, 3), padding='same', input_shape=(32, 32, 1), activation='relu'),
                  BatchNormalization(axis=-1),
                  MaxPool2D(),
                  Flatten(),
                  Dense(10, activation='relu'),
                  BatchNormalization(axis=-1),
                  Dense(transformer.n_transforms, activation='softmax')])

mdl.compile('adam',
            'categorical_crossentropy',
            ['acc'])

batch_size = 64
mdl.fit(x=x_train_single_transformed,
        y=to_categorical(transformations_inds),
        batch_size=batch_size,
        validation_split=0.1,
        epochs=10)

single_class_preds = np.zeros((len(x_test_single), transformer.n_transforms))
for t in range(transformer.n_transforms):
    single_class_preds[:, t] = mdl.predict(transformer.transform_batch(x_test_single, [t] * len(x_test_single)),
                                           batch_size=batch_size)[:, t]
single_class_scores = single_class_preds.mean(axis=-1)

anomaly_class_preds = np.zeros((len(x_test_anomaly), transformer.n_transforms))
for t in range(transformer.n_transforms):
    anomaly_class_preds[:, t] = mdl.predict(transformer.transform_batch(x_test_anomaly, [t] * len(x_test_anomaly)),
                                            batch_size=batch_size)[:, t]
anomaly_class_scores = anomaly_class_preds.mean(axis=-1)


def affine(x, is_flip, k_rotate):
    return tf.image.rot90(tf.image.flip_left_right(x) if is_flip else x,
                          k=k_rotate)


x_in = Input(batch_shape=mdl.input_shape)
transformations_sm_responses = [mdl(Lambda(affine, arguments={'is_flip': is_flip, 'k_rotate': k_rotate})(x_in))
                                for is_flip, k_rotate in itertools.product((False, True), range(4))]
out = average([Lambda(lambda sm_res: sm_res[:, j:j+1])(tens) for j, tens in enumerate(transformations_sm_responses)])


inference_mdl = Model(x_in, out)

grads_tensor = K.gradients([inference_mdl.output], [inference_mdl.input])[0]
grads_fn = K.function([inference_mdl.input], [grads_tensor])


def optimize_anomaly_images():
    for im_ind in range(len(x_test_anomaly)):
        im = x_test_anomaly[im_ind:im_ind+1].copy()

        eta = 5
        for _ in range(200):
            grads = grads_fn([im])[0]
            grads[np.abs(grads * im) < np.percentile(np.abs(grads * im), 80)] = 0
            im_diff = grads * eta
            im_diff *= 0.99
            im += im_diff
            im = gaussian_filter(im, 0.28)

        im = np.clip(im, 0, 1)
        im[im < np.percentile(np.abs(im), 80)] = 0

        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(4, 2))
        ax1.imshow(x_test_anomaly[im_ind].squeeze(), cmap='Greys_r')
        ax1.grid(False)
        ax1.tick_params(which='both', bottom=False, top=False, left=False, right=False,
                        labelbottom=False, labeltop=False, labelleft=False, labelright=False)
        ax2.imshow(im.squeeze(), cmap='Greys_r')
        ax2.grid(False)
        ax2.tick_params(which='both', bottom=False, top=False, left=False, right=False,
                        labelbottom=False, labeltop=False, labelleft=False, labelright=False)

        fig.savefig('0_{}.png'.format(im_ind))
        plt.close()
        print('0_3_{} done'.format(im_ind))


def optimize_normal_images():
    for im_ind in range(len(x_train_single)):
        im = x_train_single[im_ind:im_ind+1].copy()

        eta = 5
        for _ in range(200):
            grads = grads_fn([im])[0]
            grads[np.abs(grads * im) < np.percentile(np.abs(grads * im), 80)] = 0
            im_diff = grads * eta
            im_diff *= 0.99
            im += im_diff
            im = gaussian_filter(im, 0.28)

        im = np.clip(im, 0, 1)
        im[im < np.percentile(np.abs(im), 80)] = 0

        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(4, 2))
        ax1.imshow(x_train_single[im_ind].squeeze(), cmap='Greys_r')
        ax1.grid(False)
        ax1.tick_params(which='both', bottom=False, top=False, left=False, right=False,
                        labelbottom=False, labeltop=False, labelleft=False, labelright=False)
        ax2.imshow(im.squeeze(), cmap='Greys_r')
        ax2.grid(False)
        ax2.tick_params(which='both', bottom=False, top=False, left=False, right=False,
                        labelbottom=False, labeltop=False, labelleft=False, labelright=False)

        fig.savefig('3_{}.png'.format(im_ind))
        plt.close()
        print('3_3_{} done'.format(im_ind))


optimize_normal_images()
optimize_anomaly_images()


================================================
FILE: models/__init__.py
================================================


================================================
FILE: models/adgan.py
================================================
import numpy as np
from keras.optimizers import Adam
from keras.models import Model
from keras.layers.merge import subtract
from keras.utils.generic_utils import Progbar
from keras.engine.topology import Input, Layer
from keras.callbacks import CallbackList
import keras.backend as K


class GradPenLayer(Layer):
    def call(self, inputs, **kwargs):
        interp_in, critic_interp_score_in = inputs[0], inputs[1]
        interp_critic_grad = K.batch_flatten(K.gradients(critic_interp_score_in, [interp_in])[0])
        interp_critic_grad_norm = K.sqrt(K.sum(K.square(interp_critic_grad), axis=-1, keepdims=True))
        return K.square(interp_critic_grad_norm - 1.)  # two sided regularisation
        # return K.square(K.relu(interp_critic_grad_norm - 1.))  # one sided regularisation

    def compute_output_shape(self, input_shape):
        return input_shape[0], 1


def train_wgan_with_grad_penalty(prior_gen, generator, data_gen, critic, batch_size, epochs,
                                 batches_per_epoch=100, optimizer=Adam(lr=1e-4, beta_1=0, beta_2=0.9),
                                 grad_pen_coef=10., critic_gen_train_ratio=2, callbacks=None):
    # build model to train the critic
    data_shape = critic.input_shape[1:]
    real_critic_input = Input(shape=data_shape, name='real_in')
    fake_critic_input = Input(shape=data_shape, name='fake_in')
    interp_critic_input = Input(shape=data_shape, name='interp_in')

    real_critic_score = critic(real_critic_input)
    fake_critic_score = critic(fake_critic_input)
    interp_critic_score = critic(interp_critic_input)

    critic_loss = subtract([fake_critic_score, real_critic_score])
    gradient_penalty = GradPenLayer()([interp_critic_input, interp_critic_score])

    critic_train_mdl = Model([real_critic_input, fake_critic_input, interp_critic_input],
                             [critic_loss, gradient_penalty])

    critic_train_mdl.compile(optimizer=optimizer,
                             loss=lambda y_true, y_pred: y_pred,
                             loss_weights=[1., grad_pen_coef])

    # build model to train generator
    prior_input = Input(shape=generator.input_shape[1:], name='prior_in')
    critic.trainable = False
    critic_on_generator_score = critic(generator(prior_input))
    generator_train_mdl = Model(prior_input, critic_on_generator_score)
    generator_train_mdl.compile(optimizer=optimizer, loss=lambda y_true, y_pred: -y_pred)

    # init callbacks
    callbacks = callbacks or []
    callbacks = CallbackList(callbacks)
    callbacks.set_model({'generator': generator, 'critic': critic})
    callbacks.set_params({
        'batch_size': batch_size,
        'epochs': epochs,
        'steps': batches_per_epoch,
        'samples': batches_per_epoch * batch_size,
        'prior_gen': prior_gen,
        'data_gen': data_gen,
    })

    # train
    print('Training on {} samples for {} epochs'.format(batches_per_epoch * batch_size, epochs))
    callbacks.on_train_begin()
    for e in range(epochs):
        print('Epoch {}/{}'.format(e + 1, epochs))
        callbacks.on_epoch_begin(e)
        progbar = Progbar(target=batches_per_epoch*batch_size)
        dummy_y = np.array([None]*batch_size)
        for b in range(batches_per_epoch):
            callbacks.on_batch_begin(b)
            batch_losses = np.zeros(shape=3)
            for critic_upd in range(critic_gen_train_ratio):
                real_batch = data_gen(batch_size)
                fake_batch = generator.predict(prior_gen(batch_size))
                weights = np.random.uniform(size=batch_size)
                weights = weights.reshape((-1,) + (1,)*(len(real_batch.shape)-1))
                interp_batch = weights * real_batch + (1. - weights) * fake_batch

                x_batch = {'real_in': real_batch, 'fake_in': fake_batch, 'interp_in': interp_batch}
                cur_losses = np.array(critic_train_mdl.train_on_batch(x=x_batch, y=[dummy_y, dummy_y]))
                batch_losses += cur_losses

            generator_train_mdl.train_on_batch(x=prior_gen(batch_size), y=dummy_y)

            losses_names = ('total_loss', 'critic_loss', 'gradient_pen')
            progbar.add(batch_size, zip(losses_names, batch_losses))
            callbacks.on_batch_end(b)

        progbar.update(batches_per_epoch*batch_size)
        callbacks.on_epoch_end(e)

    callbacks.on_train_end()


def scores_from_adgan_generator(x_test, prior_gen, generator, n_seeds=8, k=5, z_lr=0.25, gen_lr=5e-5):
    generator.trainable = True
    initial_weights = generator.get_weights()

    gen_opt = Adam(lr=gen_lr, beta_1=0.5)
    z_opt = Adam(lr=z_lr, beta_1=0.5)

    x_ph = K.placeholder((1,)+x_test.shape[1:])
    z = K.variable(prior_gen(1))
    rec_loss = K.mean(K.square(x_ph - generator(z)))
    z_train_fn = K.function([x_ph], [rec_loss], updates=z_opt.get_updates(rec_loss, [z]))
    g_train_fn = K.function([x_ph, K.learning_phase()], [rec_loss],
                            updates=gen_opt.get_updates(rec_loss, generator.trainable_weights))

    gen_opt_initial_params = gen_opt.get_weights()
    z_opt_initial_params = z_opt.get_weights()

    scores = []
    for x in x_test:
        x = np.expand_dims(x, axis=0)
        losses = []
        for j in range(n_seeds):
            K.set_value(z, prior_gen(1))
            generator.set_weights(initial_weights)
            gen_opt.set_weights(gen_opt_initial_params)
            z_opt.set_weights(z_opt_initial_params)
            for _ in range(k):
                z_train_fn([x])
                g_train_fn([x, 1])
            loss = z_train_fn([x])[0]
            losses.append(loss)

        score = -np.mean(losses)
        scores.append(score)

    return np.array(scores)


================================================
FILE: models/dagmm.py
================================================
import numpy as np
from keras.models import Input, Model
from keras.layers import concatenate, Lambda, Layer, add
import keras.backend as K
import tensorflow as tf


class GaussianMixtureComponent(Layer):
    def __init__(self, lambd_diag=0.005, **kwargs):
        self.lambd_diag = lambd_diag
        super().__init__(**kwargs)

    def build(self, input_shapes):
        z_shape, _ = input_shapes
        self.phi = self.add_weight(name='phi',
                                   shape=(1,),
                                   initializer='ones',
                                   trainable=False)
        self.mu = self.add_weight(name='mu',
                                  shape=(1, z_shape[1]),
                                  initializer='uniform',
                                  trainable=False)
        self.sigma = self.add_weight(name='sig',
                                     shape=(z_shape[1], z_shape[1]),
                                     initializer='identity',
                                     trainable=False)
        super().build(input_shapes)

    def call(self, inputs, training=None):
        z, gamma_k = inputs

        gamma_k_sum = K.sum(gamma_k)
        est_phi = K.mean(gamma_k, axis=0)
        est_mu = K.dot(K.transpose(gamma_k), z) / gamma_k_sum
        est_sigma = K.dot(K.transpose(z - est_mu),
                          gamma_k * (z - est_mu)) / gamma_k_sum

        est_sigma = est_sigma + (K.random_normal(shape=(K.int_shape(z)[1], 1), mean=1e-3, stddev=1e-4) * K.eye(K.int_shape(z)[1]))

        self.add_update(K.update(self.phi, est_phi), inputs)
        self.add_update(K.update(self.mu, est_mu), inputs)
        self.add_update(K.update(self.sigma, est_sigma), inputs)

        est_sigma_diag_inv = K.eye(K.int_shape(self.sigma)[0]) / est_sigma
        self.add_loss(self.lambd_diag * K.sum(est_sigma_diag_inv), inputs)

        phi = K.in_train_phase(est_phi, self.phi, training)
        mu = K.in_train_phase(est_mu, self.mu, training)
        sigma = K.in_train_phase(est_sigma, self.sigma, training)
        return GaussianMixtureComponent._calc_component_density(z, phi, mu, sigma)

    @staticmethod
    def _calc_component_density(z, phi, mu, sigma):
        sig_inv = tf.matrix_inverse(sigma)
        sig_sqrt_det = K.sqrt(tf.matrix_determinant(2 * np.pi * sigma) + K.epsilon())
        density = phi * (K.exp(-0.5 * K.sum(K.dot(z - mu, sig_inv) * (z - mu),
                                            axis=-1,
                                            keepdims=True)) / sig_sqrt_det) + K.epsilon()

        return density

    def compute_output_shape(self, input_shapes):
        z_shape, gamma_k_shape = input_shapes
        return z_shape[0], 1


def create_dagmm_model(encoder, decoder, estimation_encoder, lambd_diag=0.005):
    x_in = Input(batch_shape=encoder.input_shape)
    zc = encoder(x_in)

    decoder.name = 'reconstruction'
    x_rec = decoder(zc)
    euclid_dist = Lambda(lambda args: K.sqrt(K.sum(K.batch_flatten(K.square(args[0] - args[1])),
                                                   axis=-1, keepdims=True) /
                                             K.sum(K.batch_flatten(K.square(args[0])),
                                                   axis=-1, keepdims=True)),
                         output_shape=(1,))([x_in, x_rec])
    cos_sim = Lambda(lambda args: K.batch_dot(K.l2_normalize(K.batch_flatten(args[0]), axis=-1),
                                              K.l2_normalize(K.batch_flatten(args[1]), axis=-1),
                                              axes=-1),
                     output_shape=(1,))([x_in, x_rec])

    zr = concatenate([euclid_dist, cos_sim])
    z = concatenate([zc, zr])

    gamma = estimation_encoder(z)

    gamma_ks = [Lambda(lambda g: g[:, k:k + 1], output_shape=(1,))(gamma)
                for k in range(estimation_encoder.output_shape[-1])]

    components = [GaussianMixtureComponent(lambd_diag)([z, gamma_k])
                  for gamma_k in gamma_ks]
    density = add(components) if len(components) > 1 else components[0]
    energy = Lambda(lambda dens: -K.log(dens), name='energy')(density)

    dagmm = Model(x_in, [x_rec, energy])

    return dagmm


================================================
FILE: models/dsebm.py
================================================
from keras.layers import *
from keras.models import Model
import keras.backend as K


class Prior(Layer):
    def __init__(self,
                 bias_initializer='zeros',
                 bias_regularizer=None,
                 bias_constraint=None,
                 **kwargs):
        if 'input_shape' not in kwargs and 'input_dim' in kwargs:
            kwargs['input_shape'] = (kwargs.pop('input_dim'),)
        super(Prior, self).__init__(**kwargs)
        self.bias_initializer = initializers.get(bias_initializer)
        self.bias_regularizer = regularizers.get(bias_regularizer)
        self.bias_constraint = constraints.get(bias_constraint)
        self.input_spec = InputSpec(min_ndim=2)
        self.bias = None

    def build(self, input_shape):
        self.bias = self.add_weight(shape=input_shape[1:],
                                    initializer=self.bias_initializer,
                                    name='bias',
                                    regularizer=self.bias_regularizer,
                                    constraint=self.bias_constraint)
        super(Prior, self).build(input_shape)

    def compute_output_shape(self, input_shape):
        return input_shape[0], 1

    def call(self, x, **kwargs):
        return K.sum(K.batch_flatten(K.square(K.bias_add(x, -self.bias))), axis=-1, keepdims=True)


def create_energy_model(encoder_mdl):
    x_in = Input(batch_shape=encoder_mdl.input_shape)

    encoded = encoder_mdl(x_in)
    prior = Prior()(x_in)
    energy = Lambda(lambda args: 0.5*args[0] - K.sum(K.batch_flatten(args[1]), axis=-1, keepdims=True),
                    output_shape=(1,))([prior, encoded])
    return Model(x_in, energy)


def create_reconstruction_model(energy_mdl):
    x_in = Input(batch_shape=energy_mdl.input_shape)
    x = GaussianNoise(stddev=.5)(x_in)  # only active in training
    energy = energy_mdl(x)
    rec = Lambda(lambda args: args[1] - K.gradients(args[0], args[1]),
                 output_shape=energy_mdl.input_shape[1:])([energy, x])

    return Model(x_in, rec)


================================================
FILE: models/encoders_decoders.py
================================================
from keras.models import Model
from keras.layers import Input, Conv2D, Activation, BatchNormalization, Flatten, Dense, Conv2DTranspose, Reshape
from utils import get_channels_axis


def conv_encoder(input_side=32, n_channels=3, representation_dim=256, representation_activation='tanh',
                 intermediate_activation='relu'):
    nf = 64
    input_shape = (n_channels, input_side, input_side) if get_channels_axis() == 1 else (input_side, input_side,
                                                                                         n_channels)

    x_in = Input(shape=input_shape)
    enc = x_in

    # downsample x0.5
    enc = Conv2D(nf, kernel_size=(3, 3), strides=(2, 2), padding='same')(enc)
    enc = BatchNormalization(axis=get_channels_axis())(enc)
    enc = Activation(intermediate_activation)(enc)

    # downsample x0.5
    enc = Conv2D(nf * 2, kernel_size=(3, 3), strides=(2, 2), padding='same')(enc)
    enc = BatchNormalization(axis=get_channels_axis())(enc)
    enc = Activation(intermediate_activation)(enc)

    # downsample x0.5
    enc = Conv2D(nf * 4, kernel_size=(3, 3), strides=(2, 2), padding='same')(enc)
    enc = BatchNormalization(axis=get_channels_axis())(enc)
    enc = Activation(intermediate_activation)(enc)

    if input_side == 64:
        # downsample x0.5
        enc = Conv2D(nf * 8, kernel_size=(3, 3), strides=(2, 2), padding='same')(enc)
        enc = BatchNormalization(axis=get_channels_axis())(enc)
        enc = Activation(intermediate_activation)(enc)

    enc = Flatten()(enc)
    rep = Dense(representation_dim, activation=representation_activation)(enc)

    return Model(x_in, rep)


def conv_decoder(output_side=32, n_channels=3, representation_dim=256, activation='relu'):
    nf = 64

    rep_in = Input(shape=(representation_dim,))

    g = Dense(nf * 4 * 4 * 4)(rep_in)
    g = BatchNormalization(axis=-1)(g)
    g = Activation(activation)(g)

    conv_shape = (nf * 4, 4, 4) if get_channels_axis() == 1 else (4, 4, nf * 4)
    g = Reshape(conv_shape)(g)

    # upsample x2
    g = Conv2DTranspose(nf * 2, kernel_size=(3, 3), strides=(2, 2), padding='same')(g)
    g = BatchNormalization(axis=get_channels_axis())(g)
    g = Activation(activation)(g)

    # upsample x2
    g = Conv2DTranspose(nf, kernel_size=(3, 3), strides=(2, 2), padding='same')(g)
    g = BatchNormalization(axis=get_channels_axis())(g)
    g = Activation(activation)(g)

    if output_side == 64:
        # upsample x2
        g = Conv2DTranspose(nf, kernel_size=(3, 3), strides=(2, 2), padding='same')(g)
        g = BatchNormalization(axis=get_channels_axis())(g)
        g = Activation(activation)(g)

    # upsample x2
    g = Conv2DTranspose(n_channels, kernel_size=(3, 3), strides=(2, 2), padding='same')(g)
    g = Activation('tanh')(g)

    return Model(rep_in, g)




================================================
FILE: models/wide_residual_network.py
================================================
from keras.models import Model
from keras.layers import *
from keras.regularizers import l2
from keras.initializers import _compute_fans
from keras.optimizers import SGD
from keras import backend as K


WEIGHT_DECAY = 0.5 * 0.0005


class SGDTorch(SGD):
    @interfaces.legacy_get_updates_support
    def get_updates(self, loss, params):
        grads = self.get_gradients(loss, params)
        self.updates = [K.update_add(self.iterations, 1)]

        lr = self.lr
        if self.initial_decay > 0:
            lr *= (1. / (1. + self.decay * K.cast(self.iterations,
                                                  K.dtype(self.decay))))
        # momentum
        shapes = [K.int_shape(p) for p in params]
        moments = [K.zeros(shape) for shape in shapes]
        self.weights = [self.iterations] + moments
        for p, g, m in zip(params, grads, moments):
            v = self.momentum * m + g  # velocity
            self.updates.append(K.update(m, v))

            if self.nesterov:
                new_p = p - lr * (self.momentum * v + g)
            else:
                new_p = p - lr * v

            # Apply constraints.
            if getattr(p, 'constraint', None) is not None:
                new_p = p.constraint(new_p)

            self.updates.append(K.update(p, new_p))
        return self.updates


def _get_channels_axis():
    return -1 if K.image_data_format() == 'channels_last' else 1


def _conv_kernel_initializer(shape, dtype=None):
    fan_in, fan_out = _compute_fans(shape)
    stddev = np.sqrt(2. / fan_in)
    return K.random_normal(shape, 0., stddev, dtype)


def _dense_kernel_initializer(shape, dtype=None):
    fan_in, fan_out = _compute_fans(shape)
    stddev = 1. / np.sqrt(fan_in)
    return K.random_uniform(shape, -stddev, stddev, dtype)


def batch_norm():
    return BatchNormalization(axis=_get_channels_axis(), momentum=0.9, epsilon=1e-5,
                              beta_regularizer=l2(WEIGHT_DECAY), gamma_regularizer=l2(WEIGHT_DECAY))


def conv2d(output_channels, kernel_size, strides=1):
    return Convolution2D(output_channels, kernel_size, strides=strides, padding='same', use_bias=False,
                         kernel_initializer=_conv_kernel_initializer, kernel_regularizer=l2(WEIGHT_DECAY))


def dense(output_units):
    return Dense(output_units, kernel_initializer=_dense_kernel_initializer, kernel_regularizer=l2(WEIGHT_DECAY),
                 bias_regularizer=l2(WEIGHT_DECAY))


def _add_basic_block(x_in, out_channels, strides, dropout_rate=0.0):
    is_channels_equal = K.int_shape(x_in)[_get_channels_axis()] == out_channels

    bn1 = batch_norm()(x_in)
    bn1 = Activation('relu')(bn1)
    out = conv2d(out_channels, 3, strides)(bn1)
    out = batch_norm()(out)
    out = Activation('relu')(out)
    out = Dropout(dropout_rate)(out)
    out = conv2d(out_channels, 3, 1)(out)
    shortcut = x_in if is_channels_equal else conv2d(out_channels, 1, strides)(bn1)
    out = add([out, shortcut])
    return out


def _add_conv_group(x_in, out_channels, n, strides, dropout_rate=0.0):
    out = _add_basic_block(x_in, out_channels, strides, dropout_rate)
    for _ in range(1, n):
        out = _add_basic_block(out, out_channels, 1, dropout_rate)
    return out


def create_wide_residual_network(input_shape, num_classes, depth, widen_factor=1, dropout_rate=0.0,
                                 final_activation='softmax'):
    n_channels = [16, 16*widen_factor, 32*widen_factor, 64*widen_factor]
    assert ((depth - 4) % 6 == 0), 'depth should be 6n+4'
    n = (depth - 4) // 6

    inp = Input(shape=input_shape)
    conv1 = conv2d(n_channels[0], 3)(inp)  # one conv at the beginning (spatial size: 32x32)
    conv2 = _add_conv_group(conv1, n_channels[1], n, 1, dropout_rate)  # Stage 1 (spatial size: 32x32)
    conv3 = _add_conv_group(conv2, n_channels[2], n, 2, dropout_rate)  # Stage 2 (spatial size: 16x16)
    conv4 = _add_conv_group(conv3, n_channels[3], n, 2, dropout_rate)  # Stage 3 (spatial size: 8x8)

    out = batch_norm()(conv4)
    out = Activation('relu')(out)
    out = AveragePooling2D(8)(out)
    out = Flatten()(out)

    out = dense(num_classes)(out)
    out = Activation(final_activation)(out)

    return Model(inp, out)


================================================
FILE: multiclass_experiment.py
================================================
from glob import glob
import os
import numpy as np
import matplotlib.pyplot as plt
from keras.utils import to_categorical
from keras.preprocessing.image import ImageDataGenerator
from keras.callbacks import LearningRateScheduler
from keras.layers import Activation
from keras import backend as K
from keras.models import Model
from models.wide_residual_network import create_wide_residual_network, SGDTorch, dense
from utils import load_cifar10, normalize_minus1_1, save_roc_pr_curve_data
from transformations import Transformer


def load_tinyimagenet(tinyimagenet_path='./'):
    images = [plt.imread(fp) for fp in glob(os.path.join(tinyimagenet_path, '*.jpg'))]

    for i in range(len(images)):
        if len(images[i].shape) != 3:
            images[i] = np.stack([images[i], images[i], images[i]], axis=-1)

    images = np.stack(images)
    images = normalize_minus1_1(K.cast_to_floatx(images))
    return images


def train_cifar10():
    (x_train, y_train), (x_test, y_test) = load_cifar10()

    idg = ImageDataGenerator(
        horizontal_flip=True,
        height_shift_range=4,
        width_shift_range=4,
        fill_mode='reflect'
    )

    idg.fit(x_train)

    n = 16
    k = 8
    mdl = create_wide_residual_network(x_train.shape[1:], 10, n, k)
    mdl.compile(SGDTorch(lr=.1, momentum=0.9, nesterov=True), 'categorical_crossentropy', ['acc'])

    lr_cb = LearningRateScheduler(lambda e: 0.1 * (0.2 ** (e >= 160 and 3 or e >= 120 and 2 or e >= 60 and 1 or 0)))

    batch_size = 128
    mdl.fit_generator(
        generator=idg.flow(x_train, to_categorical(y_train), batch_size=batch_size),
        epochs=200,
        validation_data=(idg.standardize(x_test), to_categorical(y_test)),
        callbacks=[lr_cb]
    )
    mdl.save_weights('cifar10_WRN_{}-{}.h5'.format(n, k))


def train_cifar10_transformations():
    (x_train, y_train), _ = load_cifar10()

    transformer = Transformer(8, 8)

    def data_gen(x, y, batch_size):
        while True:
            ind_permutation = np.random.permutation(len(x))
            for b_start_ind in range(0, len(x), batch_size):
                batch_inds = ind_permutation[b_start_ind:b_start_ind + batch_size]
                x_batch = x[batch_inds]
                y_batch = y[batch_inds].flatten()

                if K.image_data_format() == 'channels_first':
                    x_batch = np.transpose(x_batch, (0, 2, 3, 1))

                y_t_batch = np.random.randint(0, transformer.n_transforms, size=len(x_batch))

                x_batch = transformer.transform_batch(x_batch, y_t_batch)

                if K.image_data_format() == 'channels_first':
                    x_batch = np.transpose(x_batch, (0, 3, 1, 2))

                yield (x_batch, [to_categorical(y_batch, num_classes=10), to_categorical(y_t_batch, num_classes=transformer.n_transforms)])

    n = 16
    k = 8
    base_mdl = create_wide_residual_network(x_train.shape[1:], 10, n, k)

    transformations_cls_out = Activation('softmax')(dense(transformer.n_transforms)(base_mdl.get_layer(index=-3).output))

    mdl = Model(base_mdl.input, [base_mdl.output, transformations_cls_out])

    mdl.compile(SGDTorch(lr=.1, momentum=0.9, nesterov=True), 'categorical_crossentropy', ['acc'])

    lr_cb = LearningRateScheduler(lambda e: 0.1 * (0.2 ** (e >= 160 and 3 or e >= 120 and 2 or e >= 60 and 1 or 0)))

    batch_size = 128
    mdl.fit_generator(
        generator=data_gen(x_train, y_train, batch_size=batch_size),
        steps_per_epoch=len(x_train) // batch_size,
        epochs=200,
        callbacks=[lr_cb]
    )
    mdl.save_weights('cifar10_WRN_doublehead-transformations_{}-{}.h5'.format(n, k))


def transformation_cifar10_vs_tinyimagenet():
    _, (x_test, y_test) = load_cifar10()
    x_test_out = load_tinyimagenet('/home/izikgo/Imagenet_resize/Imagenet_resize/')

    transformer = Transformer(8, 8)
    n = 16
    k = 8
    base_mdl = create_wide_residual_network(x_test.shape[1:], 10, n, k)

    transformations_cls_out = Activation('softmax')(dense(transformer.n_transforms)(base_mdl.get_layer(index=-3).output))

    mdl = Model(base_mdl.input, [base_mdl.output, transformations_cls_out])
    mdl.load_weights('cifar10_WRN_doublehead-transformations_{}-{}.h5'.format(n, k))

    scores_mdl = Model(mdl.input, mdl.output[1])
    x_test_all = np.concatenate((x_test, x_test_out))
    preds = np.zeros((len(x_test_all), transformer.n_transforms))
    for t in range(transformer.n_transforms):
        preds[:, t] = scores_mdl.predict(transformer.transform_batch(x_test_all, [t] * len(x_test_all)),
                                  batch_size=128)[:, t]

    labels = np.concatenate((np.ones(len(x_test)), np.zeros(len(x_test_out))))
    scores = preds.mean(axis=-1)

    save_roc_pr_curve_data(scores, labels, 'cifar10-vs-tinyimagenet_transformations.npz')


if __name__ == '__main__':
    train_cifar10_transformations()
    transformation_cifar10_vs_tinyimagenet()


================================================
FILE: transformations.py
================================================
import abc
import itertools
import numpy as np

from keras.preprocessing.image import apply_affine_transform


class AffineTransformation(object):
    def __init__(self, flip, tx, ty, k_90_rotate):
        self.flip = flip
        self.tx = tx
        self.ty = ty
        self.k_90_rotate = k_90_rotate

    def __call__(self, x):
        res_x = x
        if self.flip:
            res_x = np.fliplr(res_x)
        if self.tx != 0 or self.ty != 0:
            res_x = apply_affine_transform(res_x, tx=self.tx, ty=self.ty, channel_axis=2, fill_mode='reflect')
        if self.k_90_rotate != 0:
            res_x = np.rot90(res_x, self.k_90_rotate)

        return res_x


class AbstractTransformer(abc.ABC):
    def __init__(self):
        self._transformation_list = None
        self._create_transformation_list()

    @property
    def n_transforms(self):
        return len(self._transformation_list)

    @abc.abstractmethod
    def _create_transformation_list(self):
        return

    def transform_batch(self, x_batch, t_inds):
        assert len(x_batch) == len(t_inds)

        transformed_batch = x_batch.copy()
        for i, t_ind in enumerate(t_inds):
            transformed_batch[i] = self._transformation_list[t_ind](transformed_batch[i])
        return transformed_batch


class Transformer(AbstractTransformer):
    def __init__(self, translation_x=8, translation_y=8):
        self.max_tx = translation_x
        self.max_ty = translation_y
        super().__init__()

    def _create_transformation_list(self):
        transformation_list = []
        for is_flip, tx, ty, k_rotate in itertools.product((False, True),
                                                           (0, -self.max_tx, self.max_tx),
                                                           (0, -self.max_ty, self.max_ty),
                                                           range(4)):
            transformation = AffineTransformation(is_flip, tx, ty, k_rotate)
            transformation_list.append(transformation)

        self._transformation_list = transformation_list


class SimpleTransformer(AbstractTransformer):
    def _create_transformation_list(self):
        transformation_list = []
        for is_flip, k_rotate in itertools.product((False, True),
                                                    range(4)):
            transformation = AffineTransformation(is_flip, 0, 0, k_rotate)
            transformation_list.append(transformation)

        self._transformation_list = transformation_list



================================================
FILE: utils.py
================================================
from glob import glob
import os
import numpy as np
import cv2
from sklearn.metrics import roc_curve, precision_recall_curve, auc
from keras.datasets import mnist, fashion_mnist, cifar100, cifar10
from keras.backend import cast_to_floatx


def resize_and_crop_image(input_file, output_side_length, greyscale=False):
    img = cv2.imread(input_file)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB if not greyscale else cv2.COLOR_BGR2GRAY)
    height, width = img.shape[:2]
    new_height = output_side_length
    new_width = output_side_length
    if height > width:
        new_height = int(output_side_length * height / width)
    else:
        new_width = int(output_side_length * width / height)
    resized_img = cv2.resize(img, (new_width, new_height), interpolation=cv2.INTER_AREA)
    height_offset = (new_height - output_side_length) // 2
    width_offset = (new_width - output_side_length) // 2
    cropped_img = resized_img[height_offset:height_offset + output_side_length,
                              width_offset:width_offset + output_side_length]
    assert cropped_img.shape[:2] == (output_side_length, output_side_length)
    return cropped_img


def normalize_minus1_1(data):
    return 2*(data/255.) - 1


def get_channels_axis():
    import keras
    idf = keras.backend.image_data_format()
    if idf == 'channels_first':
        return 1
    assert idf == 'channels_last'
    return 3


def load_fashion_mnist():
    (X_train, y_train), (X_test, y_test) = fashion_mnist.load_data()
    X_train = normalize_minus1_1(cast_to_floatx(np.pad(X_train, ((0, 0), (2, 2), (2, 2)), 'constant')))
    X_train = np.expand_dims(X_train, axis=get_channels_axis())
    X_test = normalize_minus1_1(cast_to_floatx(np.pad(X_test, ((0, 0), (2, 2), (2, 2)), 'constant')))
    X_test = np.expand_dims(X_test, axis=get_channels_axis())
    return (X_train, y_train), (X_test, y_test)


def load_mnist():
    (X_train, y_train), (X_test, y_test) = mnist.load_data()
    X_train = normalize_minus1_1(cast_to_floatx(np.pad(X_train, ((0, 0), (2, 2), (2, 2)), 'constant')))
    X_train = np.expand_dims(X_train, axis=get_channels_axis())
    X_test = normalize_minus1_1(cast_to_floatx(np.pad(X_test, ((0, 0), (2, 2), (2, 2)), 'constant')))
    X_test = np.expand_dims(X_test, axis=get_channels_axis())
    return (X_train, y_train), (X_test, y_test)


def load_cifar10():
    (X_train, y_train), (X_test, y_test) = cifar10.load_data()
    X_train = normalize_minus1_1(cast_to_floatx(X_train))
    X_test = normalize_minus1_1(cast_to_floatx(X_test))
    return (X_train, y_train), (X_test, y_test)


def load_cifar100(label_mode='coarse'):
    (X_train, y_train), (X_test, y_test) = cifar100.load_data(label_mode=label_mode)
    X_train = normalize_minus1_1(cast_to_floatx(X_train))
    X_test = normalize_minus1_1(cast_to_floatx(X_test))
    return (X_train, y_train), (X_test, y_test)


def save_roc_pr_curve_data(scores, labels, file_path):
    scores = scores.flatten()
    labels = labels.flatten()

    scores_pos = scores[labels == 1]
    scores_neg = scores[labels != 1]

    truth = np.concatenate((np.zeros_like(scores_neg), np.ones_like(scores_pos)))
    preds = np.concatenate((scores_neg, scores_pos))
    fpr, tpr, roc_thresholds = roc_curve(truth, preds)
    roc_auc = auc(fpr, tpr)

    # pr curve where "normal" is the positive class
    precision_norm, recall_norm, pr_thresholds_norm = precision_recall_curve(truth, preds)
    pr_auc_norm = auc(recall_norm, precision_norm)

    # pr curve where "anomaly" is the positive class
    precision_anom, recall_anom, pr_thresholds_anom = precision_recall_curve(truth, -preds, pos_label=0)
    pr_auc_anom = auc(recall_anom, precision_anom)

    np.savez_compressed(file_path,
                        preds=preds, truth=truth,
                        fpr=fpr, tpr=tpr, roc_thresholds=roc_thresholds, roc_auc=roc_auc,
                        precision_norm=precision_norm, recall_norm=recall_norm,
                        pr_thresholds_norm=pr_thresholds_norm, pr_auc_norm=pr_auc_norm,
                        precision_anom=precision_anom, recall_anom=recall_anom,
                        pr_thresholds_anom=pr_thresholds_anom, pr_auc_anom=pr_auc_anom)


def create_cats_vs_dogs_npz(cats_vs_dogs_path='./'):
    labels = ['cat', 'dog']
    label_to_y_dict = {l: i for i, l in enumerate(labels)}

    def _load_from_dir(dir_name):
        glob_path = os.path.join(cats_vs_dogs_path, dir_name, '*.*.jpg')
        imgs_paths = glob(glob_path)
        images = [resize_and_crop_image(p, 64) for p in imgs_paths]
        x = np.stack(images)
        y = [label_to_y_dict[os.path.split(p)[-1][:3]] for p in imgs_paths]
        y = np.array(y)
        return x, y

    x_train, y_train = _load_from_dir('train')
    x_test, y_test = _load_from_dir('test')

    np.savez_compressed(os.path.join(cats_vs_dogs_path, 'cats_vs_dogs.npz'),
                        x_train=x_train, y_train=y_train,
                        x_test=x_test, y_test=y_test)


def load_cats_vs_dogs(cats_vs_dogs_path='./'):
    npz_file = np.load(os.path.join(cats_vs_dogs_path, 'cats_vs_dogs.npz'))
    x_train = normalize_minus1_1(cast_to_floatx(npz_file['x_train']))
    y_train = npz_file['y_train']
    x_test = normalize_minus1_1(cast_to_floatx(npz_file['x_test']))
    y_test = npz_file['y_test']

    return (x_train, y_train), (x_test, y_test)


def get_class_name_from_index(index, dataset_name):
    ind_to_name = {
        'cifar10': ('airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck'),
        'cifar100': ('aquatic mammals', 'fish', 'flowers', 'food containers', 'fruit and vegetables',
                     'household electrical devices', 'household furniture', 'insects', 'large carnivores',
                     'large man-made outdoor things', 'large natural outdoor scenes', 'large omnivores and herbivores',
                     'medium-sized mammals', 'non-insect invertebrates', 'people', 'reptiles', 'small mammals', 'trees',
                     'vehicles 1', 'vehicles 2'),
        'fashion-mnist': ('t-shirt', 'trouser', 'pullover', 'dress', 'coat', 'sandal', 'shirt', 'sneaker', 'bag',
                          'ankle-boot'),
        'cats-vs-dogs': ('cat', 'dog'),
    }

    return ind_to_name[dataset_name][index]