r * F.randn(shape=x.shape)
    xa = F.concatenate(*[xu, xn], axis=0).apply(need_grad=True)
    ya = model(xa)
    grads = nn.grad([ya], [xa])
    norms = [F.sum(g ** 2.0, axis=[1]) ** 0.5 for g in grads]
    gp = sum([F.mean((n - 1.0) ** 2.0) for n in norms])
    return gp.apply(persistent=True)


def feed_data(inputs, di, with_normals):
    x, n, r = inputs
    pts, nrm, rad = di.next()
    x.d = pts
    r.d = rad
    if with_normals:
        n.d = nrm


def evaluate(model, pts_true, grid_size, volume_factor, monitor_distances,
             i, save_interval_epoch=1):
    if i % save_interval_epoch != 0:
        return
    pts, vol = utils.compute_pts_vol(model, grid_size, volume_factor)
    mesh = utils.create_mesh_from_volume(vol)
    pcd = mesh.sample_points_poisson_disk(len(pts_true), seed=412)
    pts_pred = np.asarray(pcd.points)
    pts_pred = utils.normalize(pts_pred)
    # Pair-wise distance
    cd0, cd1, cd, hd0, hd1, hd = utils.chamfer_hausdorff_dists(
        pts_pred, pts_true)
    for m, d in zip(monitor_distances, [cd0, cd1, cd, hd0, hd1, hd]):
        m.add(i, d)


def save_params(monitor_path, i, save_interval_epoch=1):
    if i % save_interval_epoch != 0:
        return
    nn.save_parameters(os.path.join(monitor_path, "param_{:05d}.h5".format(i)))


def main(args):
    # Context
    ctx = get_extension_context("cudnn", device_id=args.device_id)
    nn.set_default_context(ctx)

    # Network
    model = MLP(args.dims, args.ldims)
    # points loss
    x = nn.Variable([args.batch_size, 3]).apply(need_grad=True)
    y = model(x)
    loss_points = points_loss(y)
    # normals loss
    n = nn.Variable([args.batch_size, 3])
    loss_normals = normals_loss(y, x, n, args.with_normals)
    # eikonal regularization
    r = nn.Variable([args.batch_size, 1])
    reg_eikonal = eikonal_reg(x, r, model, args.volume_factor)
    # total loss
    loss = loss_points + args.tau * loss_normals + args.lam * reg_eikonal

    # Dataset (input is normalized in [-1, 1])
    ds = point_cloud_data_source(
        args.fpath, args.knn, args.test_rate, test=False)
    di = point_cloud_data_iterator(ds, args.batch_size)
    ds = point_cloud_data_source(args.fpath, -1, test=True)
    pts_true = ds.points

    # Solver
    solver = S.Adam(args.learning_rate, args.beta0, args.beta1)
    solver.set_parameters(nn.get_parameters())

    # Monitor
    monitor = M.Monitor(args.monitor_path)
    monitor_time = M.MonitorTimeElapsed(
        "Training time", monitor, interval=100, verbose=False)
    monitor_points_loss = M.MonitorSeries(
        "Points loss", monitor, interval=100, verbose=False)
    monitor_normals_loss = M.MonitorSeries(
        "Normals loss", monitor, interval=100, verbose=False)
    monitor_eikonal_reg = M.MonitorSeries(
        "Eikonal reg", monitor, interval=100, verbose=False)
    monitor_total_loss = M.MonitorSeries(
        "Total loss", monitor, interval=100, verbose=False)

    monitor_distances = [
                         M.MonitorSeries("Chamfer Pred True",
                                         monitor, interval=1, verbose=False),
                         M.MonitorSeries("Chamfer True Pred",
                                         monitor, interval=1, verbose=False),
                         M.MonitorSeries("Chamfer", monitor,
                                         interval=1, verbose=False),
                         M.MonitorSeries("Hausdorff Pred True",
                                         monitor, interval=1, verbose=False),
                         M.MonitorSeries("Hausdorff True Pred",
                                         monitor, interval=1, verbose=False),
                         M.MonitorSeries("Hausdorff", monitor, interval=1, verbose=False)]

    # Train
    iter_per_epoch = di.size // args.batch_size
    for i in tqdm.tqdm(range(args.train_epoch), desc="train-loop"):
        # evaluate
        evaluate(model, pts_true, args.grid_size, args.volume_factor, monitor_distances,
                 i, args.save_interval_epoch)
        for j in range(iter_per_epoch):
            # feed data
            feed_data([x, n, r], di, args.with_normals)
            # zero_grad, forward, backward, update
            solver.zero_grad()
            loss.forward()
            loss.backward(clear_buffer=True)
            solver.update()
            # monitor
            monitor_points_loss.add(i * iter_per_epoch + j, loss_points.d)
            monitor_normals_loss.add(
                i * iter_per_epoch + j, loss_normals.d) if args.with_normals else None
            monitor_eikonal_reg.add(i * iter_per_epoch + j, reg_eikonal.d)
            monitor_total_loss.add(i * iter_per_epoch + j, loss.d)
            monitor_time.add(i)
        # save
        save_params(args.monitor_path, i, args.save_interval_epoch)
    save_params(args.monitor_path, args.train_epoch)
    evaluate(model, pts_true, args.grid_size, args.volume_factor, monitor_distances,
             args.train_epoch)


if __name__ == '__main__':
    args = get_args()
    utils.save_args(args)
    main(args)


# shape-reconstruction/implicit-geometric-regularization/visualize.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import shutil
import os
import numpy as np
import open3d as o3d
from args import get_args

from functools import partial


def save_as_image(fpath, geo, cpath):
    path, ext = os.path.splitext(args.fpath)
    fpath = "{}.jpeg".format(path)
    cpath = "{}.json".format(path)

    def capture_image(vis):
        ctr = vis.get_view_control()
        vis.capture_screen_image(fpath)
        print("{} was saved.".format(fpath))
        cp = ctr.convert_to_pinhole_camera_parameters()
        o3d.io.write_pinhole_camera_parameters(cpath, cp)
        print("{} was saved.".format(cpath))
        return False

    def apply_camera_parameters(vis):
        if cpath != "":
            ctr = vis.get_view_control()
            cp = o3d.io.read_pinhole_camera_parameters(cpath)
            ctr.convert_from_pinhole_camera_parameters(cp)

    key_to_callback = {}
    key_to_callback[ord("S")] = capture_image
    key_to_callback[ord("A")] = apply_camera_parameters
    o3d.visualization.draw_geometries_with_key_callbacks(
        [geo], key_to_callback)


def save_as_images(dpath, mesh):
    np.random.seed(412)
    save_as_images.vis = o3d.visualization.Visualizer()
    if not os.path.exists(dpath):
        os.makedirs(dpath)
    else:
        shutil.rmtree(dpath)
        os.makedirs(dpath)
    save_as_images.n = 0
    save_as_images.i = 0
    save_as_images.a = np.random.randint(-15, 15)
    save_as_images.b = np.random.randint(-15, 15)

    def move_forward(vis):
        fpath = os.path.join(dpath, "{:05d}.jpeg".format(save_as_images.n))
        ctr = vis.get_view_control()
        ctr.rotate(save_as_images.a, save_as_images.b)
        vis.capture_screen_image(fpath)
        print("{} is saved.".format(fpath))
        save_as_images.i += 1
        save_as_images.n += 1
        if save_as_images.i % 200 == 0:
            save_as_images.a = np.random.randint(-15, 15)
            save_as_images.b = np.random.randint(-15, 15)
            save_as_images.i = 0
        return False

    vis = save_as_images.vis
    vis.create_window()
    vis.add_geometry(mesh)
    vis.register_animation_callback(move_forward)
    vis.run()
    vis.destroy_window()


def draw_multiple_visualizations(mesh, rate):
    path, ext = os.path.splitext(args.fpath)

    def capture_image(vis, type):
        ctr = vis.get_view_control()
        fpath = "{}_{}.jpeg".format(path, type)
        vis.capture_screen_image(fpath)
        print("{} was saved.".format(fpath))
        return False

    vis = o3d.visualization.VisualizerWithKeyCallback()
    vis.register_key_callback(ord("S"), partial(capture_image, type="pcd"))
    vis.create_window(window_name='TopLeft', width=960,
                      height=540, left=0, top=0)
    n = len(mesh.vertices)
    pcd = mesh.sample_points_poisson_disk(int(n * rate))
    print("{} / {} points were sampled.".format(len(pcd.points), n))
    vis.add_geometry(pcd)

    vis2 = o3d.visualization.VisualizerWithKeyCallback()
    vis2.register_key_callback(ord("S"), partial(capture_image, type="mesh"))
    vis2.create_window(window_name='TopRight', width=960,
                       height=540, left=960, top=0)
    vis2.add_geometry(mesh)

    ctr = vis.get_view_control()
    ctr2 = vis2.get_view_control()

    while True:
        cp = ctr.convert_to_pinhole_camera_parameters()
        vis.update_geometry(pcd)
        if not vis.poll_events():
            break
        vis.update_renderer()

        ctr2.convert_from_pinhole_camera_parameters(cp)
        vis2.update_geometry(mesh)
        if not vis2.poll_events():
            break
        vis2.update_renderer()

    vis.destroy_window()
    vis2.destroy_window()


def main(args):
    if args.command == "mesh":
        mesh = o3d.io.read_triangle_mesh(args.fpath)
        mesh.compute_vertex_normals()
        save_as_image(args.fpath, mesh, args.cpath)
    elif args.command == "pcd":
        pcd = o3d.io.read_point_cloud(args.fpath)
        save_as_image(args.fpath, pcd, args.cpath)
    elif args.command == "images":
        mesh = o3d.io.read_triangle_mesh(args.fpath)
        mesh.compute_vertex_normals()
        dpath, ext = os.path.splitext(args.fpath)
        save_as_images(dpath, mesh)
    elif args.command == "pcdmesh":
        mesh = o3d.io.read_triangle_mesh(args.fpath)
        mesh.compute_vertex_normals()
        draw_multiple_visualizations(mesh, args.rate)
    else:
        raise ValueError("Command={} is not supported".format(args.command))


if __name__ == '__main__':
    args = get_args()
    main(args)


# speech-synthesis/HiFiGAN/dataset.py

# Copyright (c) 2017 Sony Corporation. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pathlib import Path

import numpy as np
from nnabla.utils.data_source import DataSource


class LJSpeechDataSource(DataSource):
    r"""Data source for LJ Speech dataset.
    Args:
        metadata (str): The metadata containing text and audio indices.
        hp (HParams): A container containing all hyper parameters.
        shuffle (bool, optional): If `shuffle==True`, data will be loaded in
            a shuffled order. Defaults to False.
        rng (:obj:`numpy.random.RandomState`, optional): Numpy random number
            generator. Defaults to None.
    """

    def __init__(self, metadata, hp, shuffle=False, rng=None):

        if rng is None:
            rng = np.random.RandomState(hp.seed)
        super().__init__(shuffle=shuffle, rng=rng)

        with open(Path(hp.precomputed_path) / 'meta_train.csv') as f:
            files = [line.split('|')[0] for line in f.readlines()]

        n = len(files)
        index = self._rng.permutation(n) if shuffle else np.arange(n)
        if hasattr(hp, 'comm'):  # distributed learning
            num = n // hp.comm.n_procs
            index = index[num * hp.comm.rank:num * (hp.comm.rank + 1)]

        self._files = [files[i] for i in index]
        self._path = Path(hp.precomputed_path) / "data"
        self._size = len(self._files)
        self._variables = ["wave"]
        self.hp = hp

        self.reset()

    def reset(self):
        if self._shuffle:
            self._indexes = self._rng.permutation(self._size)
        else:
            self._indexes = np.arange(self._size)
        super().reset()

    def _get_data(self, position):
        r"""Return a tuple of data."""
        hp = self.hp
        index = self._indexes[position]
        name = self._files[index]
        w = np.load(self._path / (name + '.npz'))['wave']
        if len(w) > hp.segment_length:
            idx = self._rng.randint(0, len(w) - hp.segment_length)
            w = w[idx:idx + hp.segment_length]
        else:
            w = np.pad(w, (0, hp.segment_length - len(w)), mode='constant')
        return w[None, None, ...]


# speech-synthesis/HiFiGAN/hparams.py

# Copyright (c) 2017 Sony Corporation. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys
common_utils_path = os.path.abspath(os.path.join(
    os.path.dirname(__file__), '..', '..', 'utils'))
sys.path.append(common_utils_path)

from neu.tts.hparams import HParams

hparams = HParams(

    # dataset parameters
    corpus_path='./datasets/LJSpeech-1.1/',
    precomputed_path='./datasets/hifi/',

    output_path="./log/hifigan/",      # directory to save results

    segment_length=8192,
    resblock_kernel_sizes=[3, 7, 11],
    resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
    upsample_initial_channel=512,
    resblock="1",
    upsample_rates=[8, 8, 2, 2],
    periods=[2, 3, 5, 7, 11],
    n_layers_D=5,                       # number of layers in discriminators

    # spectrogram parameters
    sr=22050,                           # sampling rate used to read audios
    # length of the windowed signal after padding with zeros.
    n_fft=1024,
    n_mels=80,                          # number of mel filters
    mel_fmin=0.0,                       # minimum mel bank
    mel_fmax=8000.0,                    # maximum mel bank
    hop_length=256,
    win_length=1024,                    # window length

    seed=123456,                        # random seed

    # optimization parameters
    batch_size=16,                      # batch size
    epoch=3501,                         # number of epochs
    print_frequency=25,                 # number of iterations before printing
                                        # to log file
    epochs_per_checkpoint=100,          # number of epochs for each checkpoint

    alpha=2e-4,                         # learning rate
    beta1=0.8,
    beta2=0.99,
    lr_decay=0.999,
)


# speech-synthesis/HiFiGAN/main.py

# Copyright (c) 2017 Sony Corporation. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os
from pathlib import Path

import sys
sys.path.append(str(Path().cwd().parents[1] / 'utils'))

import nnabla as nn
import numpy as np
from neu.comm import CommunicatorWrapper
from neu.tts.optimizer import Optimizer
from nnabla.ext_utils import get_extension_context
from nnabla.logger import logger
from nnabla.utils.data_iterator import data_iterator
from nnabla.utils.learning_rate_scheduler import ExponentialScheduler

from dataset import LJSpeechDataSource
from hparams import hparams as hp
from model.model import Discriminator, Generator
from train import HiFiGANTrainer


def run(args):
    """Runs the algorithm."""
    Path(hp.output_path).mkdir(parents=True, exist_ok=True)

    # setup nnabla context
    ctx = get_extension_context(args.context)
    nn.set_default_context(ctx)
    hp.comm = CommunicatorWrapper(ctx)

    if hp.comm.n_procs > 1 and hp.comm.rank == 0:
        n_procs = hp.comm.n_procs
        logger.info(f'Distributed training with {n_procs} processes.')

    rng = np.random.RandomState(hp.seed)

    # train data
    train_loader = data_iterator(
        LJSpeechDataSource('meta_train.csv', hp, shuffle=True, rng=rng),
        batch_size=hp.batch_size, with_memory_cache=False
    )
    # valid data
    valid_loader = data_iterator(
        LJSpeechDataSource('meta_test.csv', hp, shuffle=False, rng=rng),
        batch_size=hp.batch_size, with_memory_cache=False
    )
    dataloader = dict(train=train_loader, valid=valid_loader)

    # build model
    gen = Generator(hp)
    dis = Discriminator(hp)

    # setup optimizer
    iter_interval = train_loader.size//hp.batch_size
    g_optim = Optimizer(
        lr_scheduler=ExponentialScheduler(
            hp.alpha, gamma=hp.lr_decay, iter_interval=iter_interval),
        name='AdamW', alpha=hp.alpha, beta1=hp.beta1, beta2=hp.beta2
    )
    d_optim = Optimizer(
        lr_scheduler=ExponentialScheduler(
            hp.alpha, gamma=hp.lr_decay, iter_interval=iter_interval),
        name='AdamW', alpha=hp.alpha, beta1=hp.beta1, beta2=hp.beta2
    )

    HiFiGANTrainer(gen, dis, g_optim, d_optim, dataloader, hp).run()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--context', '-c', type=str, default='cudnn',
                        help="'cudnn' is highly recommended.")
    parser.add_argument("--device-id", "-d", type=str, default='-1',
                        help='A list of device ids to use.\
                        This is only valid if you specify `-c cudnn`. \
                        Defaults to use all available GPUs.')
    for key, value in hp.__dict__.items():
        name = "--" + key
        if type(value) == list:
            nargs, t = '+', type(value[0])
        else:
            nargs, t = None, type(value)
        parser.add_argument(name, type=t, nargs=nargs, default=value)

    args = parser.parse_args()
    for k, v in vars(args).items():
        hp.__dict__[k] = v

    # setup context for nnabla
    if args.device_id != '-1':
        os.environ["CUDA_VISIBLE_DEVICES"] = args.device_id

    run(args)


# speech-synthesis/HiFiGAN/preprocess.py

# Copyright (c) 2017 Sony Corporation. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import multiprocessing
import random
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path

import librosa as lr
import numpy as np
from librosa.filters import mel as librosa_mel_fn
from tqdm import tqdm

from hparams import hparams as hp

mel_basis = librosa_mel_fn(
    hp.sr, hp.n_fft, n_mels=hp.n_mels,
    fmin=hp.mel_fmin, fmax=hp.mel_fmax
)


def process(line):
    r"""Read audio waveform and preprocess it.

    Args:
        line (str): A line from metadata.
    """
    path = Path(hp.corpus_path) / 'wavs'
    meta = line.strip().split('|')
    wave = lr.load(path / f'{meta[0]}.wav', sr=hp.sr)[0]
    np.savez(
        Path(hp.precomputed_path) / 'data' / (meta[0] + '.npz'),
        wave=wave
    )


def run():
    path = Path(hp.corpus_path)
    save_path = Path(hp.precomputed_path)
    save_path.joinpath('data').mkdir(parents=True, exist_ok=True)

    with open(path / 'metadata.csv', encoding='utf-8') as f:
        files = f.readlines()

    num_cores = multiprocessing.cpu_count()
    with ProcessPoolExecutor(max_workers=num_cores) as executor:
        futures = [executor.submit(process, f) for f in files]
        for f in tqdm(as_completed(futures), total=len(futures)):
            pass

    # split data into train/test sets
    file_train = Path(hp.precomputed_path) / 'meta_train.csv'
    file_valid = Path(hp.precomputed_path) / 'meta_test.csv'
    if not (file_train.exists() and file_valid.exists()):
        random.Random(hp.seed).shuffle(files)
        with open(file_train, 'w') as f:
            f.writelines(files[:int(0.98 * len(files))])
        with open(file_valid, 'w') as f:
            f.writelines(files[int(0.98 * len(files)):])


if __name__ == "__main__":
    run()


# speech-synthesis/HiFiGAN/synthesize.py

# Copyright (c) 2017 Sony Corporation. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import argparse
import os
from pathlib import Path

import sys
sys.path.append(str(Path().cwd().parents[1] / 'utils'))

import nnabla as nn
from nnabla.ext_utils import get_extension_context
import numpy as np
from scipy.io import wavfile

from hparams import hparams as hp
from model.model import Generator


def synthesize(args):
    hp.batch_size = 1
    model = Generator(hp)
    model.training = False
    model.load_parameters(args.f_model, raise_if_missing=True)
    wave = np.load(args.f_mel)
    x_mel = nn.Variable.from_numpy_array(wave[None, ...])
    o_aud = model(x_mel)

    o_aud.forward(clear_buffer=True)
    wavfile.write(args.f_output, rate=hp.sr, data=o_aud.d[0, 0].copy())


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--context', '-c', type=str, default='cudnn',
                        help="Extension module. 'cudnn' is highly recommended.")
    parser.add_argument("--device-id", "-d", type=str, default='-1',
                        help='A list of device ids to use.\
                        This is only valid if you specify `-c cudnn`. \
                        Defaults to use all available GPUs.')
    parser.add_argument("--f-model", "-m", type=str,
                        help='File path to the trained model.')
    parser.add_argument("--f-mel", "-f", type=str,
                        help='File path to the mel file.')
    parser.add_argument("--f-output", "-o", type=str, default='sample.wav',
                        help='File path to the synthetic output waveform.')
    args = parser.parse_args()

    # setup context for nnabla
    if args.device_id != '-1':
        os.environ["CUDA_VISIBLE_DEVICES"] = args.device_id

    # setup nnabla context
    ctx = get_extension_context(args.context, device_id='0')
    nn.set_default_context(ctx)

    synthesize(args)


# speech-synthesis/HiFiGAN/train.py

# Copyright (c) 2017 Sony Corporation. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
from pathlib import Path

import nnabla as nn
import nnabla.functions as F
import soundfile as sf
from librosa.filters import mel as librosa_mel_fn
from neu.tts.logger import ProgressMeter
from neu.variable_utils import set_persistent_all

from utils import compute_mel


def discriminator_loss(fea, scale):
    loss = sum([F.mean((out[-1] - scale)**2) for out in fea])
    return loss


def feature_loss(fea_real, fea_fake):
    loss = list()
    for o1, o2 in zip(fea_real, fea_fake):
        for f1, f2 in zip(o1[:-1], o2[:-1]):
            loss.append(F.mean(F.absolute_error(f1, f2)))
    return sum(loss)


class HiFiGANTrainer:
    r"""Trainer for HiFi-GAN."""

    def __init__(self, gen, dis, gen_optim, dis_optim, dataloader, hp):
        self.gen = gen
        self.dis = dis
        self.gen_optim = gen_optim
        self.dis_optim = dis_optim
        self.dataloader = dataloader
        self.hp = hp

        # compute mel basis
        mel_basis = librosa_mel_fn(hp.sr, hp.n_fft, n_mels=hp.n_mels,
                                   fmin=hp.mel_fmin, fmax=hp.mel_fmax)
        self.mel_basis = nn.Variable.from_numpy_array(mel_basis[None, ...])

        self.one_epoch_train = dataloader['train'].size // hp.batch_size
        self.one_epoch_valid = dataloader['valid'].size // hp.batch_size
        self.placeholder = dict()

        self.monitor = ProgressMeter(
            self.one_epoch_train, hp.output_path, quiet=hp.comm.rank > 0)
        hp.save(Path(hp.output_path) / 'settings.json')

    def callback_on_start(self):
        r"""Calls this on starting the training."""
        self.cur_epoch = 0
        checkpoint = Path(self.hp.output_path) / 'checkpoint.json'
        if checkpoint.is_file():
            self.load_checkpoint_model(str(checkpoint))

        self.update_graph('train')
        params = self.gen.get_parameters(grad_only=True)
        self.gen_optim.set_parameters(params)

        dis_params = self.dis.get_parameters(grad_only=True)
        self.dis_optim.set_parameters(dis_params)

        if checkpoint.is_file():
            self.load_checkpoint_optim(str(checkpoint))

        self._grads = [x.grad for x in params.values()]
        self._discs = [x.grad for x in dis_params.values()]

        self.update_graph('valid')

        self.log_variables = [
            'g_loss_avd', 'g_loss_mel',
            'g_loss', 'd_loss', 'g_loss_fea'
        ]

    def run(self):
        r"""Run the training process."""
        self.callback_on_start()

        for cur_epoch in range(self.cur_epoch + 1, self.hp.epoch + 1):
            self.monitor.reset()
            lr = self.gen_optim.get_learning_rate()
            self.monitor.info(f'Running epoch={cur_epoch}\tlr={lr:.5f}\n')
            self.cur_epoch = cur_epoch

            for i in range(self.one_epoch_train):
                self.train_on_batch(i)
                if i % (self.hp.print_frequency) == 0:
                    self.monitor.display(i, self.log_variables)

            for i in range(self.one_epoch_valid):
                self.valid_on_batch(i)

            self.callback_on_epoch_end()

        self.callback_on_finish()
        self.monitor.close()

    def _zero_grads(self):
        self.gen_optim.zero_grad()
        self.dis_optim.zero_grad()

    def update_graph(self, key='train'):
        r"""Builds the graph and update the placeholder.
        Args:
            key (str, optional): Type of computational graph.
                Defaults to 'train'.
        """
        assert key in ('train', 'valid')
        hp = self.hp

        self.gen.training = key == 'train'
        self.dis.training = key == 'train'

        # define input variables
        x_real = nn.Variable((hp.batch_size, 1, hp.segment_length))
        x_real_mel = compute_mel(x_real, self.mel_basis, hp)

        x_fake = self.gen(x_real_mel)
        x_fake_mel = compute_mel(x_fake, self.mel_basis, hp)

        dis_real_x = self.dis(x_real)
        dis_fake_x = self.dis(x_fake)

        # ------------------------------ Discriminator -----------------------
        d_loss = (
            discriminator_loss(dis_real_x, 1.0)
            + discriminator_loss(dis_fake_x, 0.0)
        )
        # -------------------------------- Generator -------------------------
        g_loss_avd = discriminator_loss(dis_fake_x, 1.0)
        g_loss_mel = F.mean(F.absolute_error(x_real_mel, x_fake_mel))
        g_loss_fea = feature_loss(dis_real_x, dis_fake_x)
        g_loss = g_loss_avd + 45*g_loss_mel + 2*g_loss_fea

        set_persistent_all(
            g_loss_mel, g_loss_avd, g_loss_fea,
            d_loss, x_fake, g_loss,
        )

        self.placeholder[key] = dict(
            x_real=x_real, x_fake=x_fake, d_loss=d_loss,
            g_loss_avd=g_loss_avd, g_loss_mel=g_loss_mel,
            g_loss_fea=g_loss_fea, g_loss=g_loss,
        )

    def train_on_batch(self, i):
        r"""Updates the model parameters."""
        hp = self.hp
        bs, p = hp.batch_size, self.placeholder['train']
        dl = self.dataloader['train']
        p['x_real'].d = dl.next()[0]

        # ----------------------------- train discriminator ----------------
        self._zero_grads()
        p['x_fake'].need_grad = False
        p['d_loss'].forward()
        p['d_loss'].backward(clear_buffer=True)
        self.monitor.update('d_loss', p['d_loss'].d.copy(), bs)
        hp.comm.all_reduce(self._discs, division=True, inplace=False)
        self.dis_optim.update()
        p['x_fake'].need_grad = True

        # ------------------------------ train generator -------------------
        self._zero_grads()
        p['g_loss'].forward()
        p['g_loss'].backward(clear_buffer=True)
        self.monitor.update('g_loss', p['g_loss'].d.copy(), bs)
        self.monitor.update('g_loss_avd', p['g_loss_avd'].d.copy(), bs)
        self.monitor.update('g_loss_mel', p['g_loss_mel'].d.copy(), bs)
        self.monitor.update('g_loss_fea', p['g_loss_fea'].d.copy(), bs)

        hp.comm.all_reduce(self._grads, division=True, inplace=False)
        self.gen_optim.update()

    def valid_on_batch(self, i):
        r"""Performs validation."""
        p, dl = self.placeholder['valid'], self.dataloader['valid']
        p['x_real'].d = dl.next()[0]
        p['x_fake'].forward(clear_buffer=True)

    def callback_on_epoch_end(self):
        hp = self.hp
        if (hp.comm.rank == 0):
            path = Path(hp.output_path) / 'artifacts'
            path.joinpath('states').mkdir(parents=True, exist_ok=True)
            path.joinpath('samples').mkdir(parents=True, exist_ok=True)
            self.save_checkpoint(path / 'states')
            self.write_samples(path / 'samples')
            if self.cur_epoch % hp.epochs_per_checkpoint == 0:
                path = path / f"epoch_{self.cur_epoch}"
                path.mkdir(parents=True, exist_ok=True)
                self.gen.save_parameters(str(path / 'model.h5'))
                self.dis.save_parameters(str(path / 'cls.h5'))

    def callback_on_finish(self):
        if self.hp.comm.rank == 0:
            path = Path(self.hp.output_path)
            self.gen.save_parameters(str(path / 'model.h5'))
            self.dis.save_parameters(str(path / 'cls.h5'))

    def save_checkpoint(self, path):
        r"""Save the current states of the trainer."""
        if self.hp.comm.rank == 0:
            path = Path(path)
            self.gen.save_parameters(str(path / 'model.h5'))
            self.dis.save_parameters(str(path / 'cls.h5'))
            self.gen_optim._solver.save_states(str(path / 'gen_optim.h5'))
            self.dis_optim._solver.save_states(str(path / 'dis_optim.h5'))
            with open(Path(self.hp.output_path) / 'checkpoint.json', 'w') as f:
                json.dump(
                    dict(cur_epoch=self.cur_epoch,
                         params_path=str(path),
                         gen_optim_n_iters=self.gen_optim._iter,
                         dis_optim_n_iters=self.dis_optim._iter,),
                    f
                )
            self.monitor.info(f"Checkpoint saved: {str(path)}\n")

    def load_checkpoint_model(self, checkpoint):
        r"""Load the last states of the trainer."""
        with open(checkpoint, 'r') as file:
            info = json.load(file)
            path = Path(info['params_path'])
        self.gen.load_parameters(str(path / 'model.h5'), raise_if_missing=True)
        self.dis.load_parameters(str(path / 'cls.h5'), raise_if_missing=True)

    def load_checkpoint_optim(self, checkpoint):
        r"""Load the last states of the trainer."""
        with open(checkpoint, 'r') as file:
            info = json.load(file)
            path = Path(info['params_path'])
            self.gen_optim._iter = info['gen_optim_n_iters']
            self.dis_optim._iter = info['dis_optim_n_iters']
            self.cur_epoch = info['cur_epoch']
        self.gen_optim._solver.load_states(str(path / 'gen_optim.h5'))
        self.dis_optim._solver.load_states(str(path / 'dis_optim.h5'))

    def write_samples(self, path):
        r"""write a few samples."""
        hp = self.hp
        p = self.placeholder['valid']
        p['x_fake'].forward(clear_buffer=True)
        X, Z = p['x_real'].d.copy(), p['x_fake'].d.copy()
        for i, (x, z) in enumerate(zip(X, Z)):
            sf.write(str(path / f'real_{i}.wav'), x[0], hp.sr)
            sf.write(str(path / f'output_{i}.wav'), z[0], hp.sr)


# speech-synthesis/HiFiGAN/utils.py

# Copyright (c) 2017 Sony Corporation. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import nnabla as nn
import nnabla.functions as F
import numpy as np


def stft(x, window_size, stride, fft_size,
         window_type='hanning', center=True, pad_mode='reflect'):
    if window_type == 'hanning':
        window_func = np.hanning(window_size + 1)[:-1]
    elif window_type == 'hamming':
        window_func = np.hamming(window_size + 1)[:-1]
    elif window_type == 'rectangular' or window_type is None:
        window_func = np.ones(window_size)
    else:
        raise ValueError("Unknown window type {}.".format(window_type))

    # pad window if `fft_size > window_size`
    if fft_size > window_size:
        diff = fft_size - window_size
        window_func = np.pad(
            window_func, (diff // 2, diff - diff // 2), mode='constant')
    elif fft_size < window_size:
        raise ValueError(
            "FFT size has to be as least as large as window size.")

    # compute STFT filter coefficients
    mat_r = np.zeros((fft_size // 2 + 1, 1, fft_size))
    mat_i = np.zeros((fft_size // 2 + 1, 1, fft_size))

    for w in range(fft_size // 2 + 1):
        for t in range(fft_size):
            mat_r[w, 0, t] = np.cos(2. * np.pi * w * t / fft_size)
            mat_i[w, 0, t] = -np.sin(2. * np.pi * w * t / fft_size)

    conv_r = nn.Variable.from_numpy_array(mat_r * window_func)
    conv_i = nn.Variable.from_numpy_array(mat_i * window_func)

    if center:
        # pad at begin/end (per default this is a reflection padding)
        p = (fft_size - stride) // 2
        x = F.pad(x, (p, p), mode=pad_mode)

    # compute STFT

    y_r = F.convolution(x, conv_r, stride=(stride,))
    y_i = F.convolution(x, conv_i, stride=(stride,))

    return y_r, y_i


def compute_mel(wave, basis, hp):
    r"""Compute the mel-spectrogram from the waveform.

    Args:
        wave (nn.Variable): Wavefrom variable of shape (B, 1, L).
        basis (nn.Variable): Basis for mel-spectrogram computation.
        hp (HParams): Hyper-parameters.

    Returns:
        nn.Variable: Output variable.
    """
    reals, imags = stft(wave, window_size=hp.win_length,
                        stride=hp.hop_length, fft_size=hp.n_fft)
    linear = (reals**2 + imags**2)**0.5
    mels = F.batch_matmul(basis, linear)
    mels = F.log(F.clip_by_value(mels, 1e-5, np.inf))

    return mels


# speech-synthesis/WaveGlow/dataset.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pathlib import Path
import sys
sys.path.append(str(Path().cwd().parents[1] / 'utils'))

import librosa as lr
from nnabla.utils.data_source import DataSource
import numpy as np
from tqdm import tqdm


class LJSpeechDataSource(DataSource):
    r"""Data source for LJ Speech dataset.

    Args:
        metadata (str): The metadata containing text and audio indices.
        hparams (HParams): A container containing all hyper parameters.
        shuffle (bool, optional): If `shuffle==True`, data will be loaded in
            a shuffled order. Defaults to False.
        rng (:obj:`numpy.random.RandomState`, optional): Numpy random number
            generator. Defaults to None.
    """

    def __init__(self, metadata, hparams, shuffle=False, rng=None):
        if rng is None:
            rng = np.random.RandomState(hparams.seed)
        super().__init__(shuffle=shuffle, rng=rng)

        # read text and wave files
        texts, waves = list(), list()
        path = Path(hparams.data_dir)
        with open(path / metadata, encoding='utf-8') as f:
            for line in f:
                inputs = line.strip().split('|')
                waves.append(path / 'wavs' / f'{inputs[0]}.wav')
                texts.append(inputs[2])
        # split data
        n = len(waves)
        index = self._rng.permutation(n) if shuffle else np.arange(n)
        if hasattr(hparams, 'comm'):  # distributed learning
            num = n // hparams.comm.n_procs
            index = index[num*hparams.comm.rank:num*(hparams.comm.rank + 1)]

        self._waves = [waves[i] for i in index]
        self._texts = [texts[i] for i in index]
        self._path = Path(hparams.save_data_dir)
        self._size = len(self._waves)
        self._variables = hparams.out_variables
        self.hparams = hparams

        self.reset()

    def reset(self):
        if self._shuffle:
            self._indexes = self._rng.permutation(self._size)
        else:
            self._indexes = np.arange(self._size)
        super().reset()

    def _get_data(self, position):
        r"""Return a tuple of data."""
        hp = self.hparams
        index = self._indexes[position]
        basename = self._waves[index].with_suffix(".npy").name
        w = np.load(self._path / 'audio' / basename)
        if len(w) > hp.segment_length:
            idx = self._rng.randint(0, len(w) - hp.segment_length)
            w = w[idx:idx + hp.segment_length]
        else:
            w = np.pad(w, (0, hp.segment_length - len(w)), mode='constant')
        return [w]

    def _preprocess(self):
        r"""Precomputes waveform of training data."""
        hp = self.hparams
        for f in self._variables:
            self._path.joinpath(f).mkdir(parents=True, exist_ok=True)

        for i in tqdm(range(self._size)):
            file = self._waves[i]
            w, _ = lr.load(file, sr=hp.sr)
            w, _ = lr.effects.trim(w)  # triming
            basename = self._waves[i].with_suffix(".npy").name
            np.save(self._path / 'audio' / basename, w)


def data_split(path):
    r"""Split the LJ dataset into train and validation sets."""
    with open(Path(path) / 'metadata.csv', encoding='utf-8') as f:
        lines = [line.strip() for line in f]
        n = len(lines)
        index = np.random.RandomState(313).permutation(n)
        split = dict(train=[lines[i] for i in index[:-int(0.1 * n)]],
                     valid=[lines[i] for i in index[-int(0.1 * n):]])
        for key in ['train', 'valid']:
            with open(Path(path) / f'metadata_{key}.csv', 'w') as writer:
                writer.write('\n'.join(split[key]))


if __name__ == '__main__':
    from hparams import hparams as hp
    rng = np.random.RandomState(hp.seed)
    data_split(hp.data_dir)
    LJSpeechDataSource('metadata.csv', hp, False, rng)._preprocess()


# speech-synthesis/WaveGlow/hparams.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from neu.tts.hparams import HParams

hparams = HParams(

    # dataset parameters
    data_dir="./data/LJSpeech-1.1/",               # directory to the data
    save_data_dir="./data/LJSpeech-1.1/waveglow",  # directory to store audios
    out_variables=["audio"],            # which variables will be used

    # spectrogram parameters
    sr=22050,                           # sampling rate used to read audios
    # length of the windowed signal after padding with zeros.
    n_fft=1024,
    n_mels=80,                          # number of mel filters
    mel_fmin=0.0,                       # minimum mel bank
    mel_fmax=8000.0,                    # maximum mel bank
    # number of audio samples between adjacent STFT columns
    hop_length=256,
    win_length=1024,                    # window length

    # flow parameters
    n_flows=12,                         # number of flow nets
    # number of samples in a group processed by the steps of flow
    n_samples_per_group=8,
    # Determines how often (i.e., after how many coupling layers)
    n_early_every=4,
                                        # a number of channels are output to the loss function
    n_early_size=2,                     # Number of channels output to the loss function
    sigma=1.0,                          # Standard deviation used for sampling from Gaussian
    # Segment length (audio samples) processed per iteration
    segment_length=8000,
    wn_n_layers=8,                      # Number of layers in WN
    # Kernel size for dialted convolution in the affine coupling layer (WN)
    wn_kernel_size=3,
    wn_n_channels=256,                  # Number of channels in WN
    seed=123456,                        # random seed

    # optimization parameters
    batch_size=4,                       # batch size
    epoch=1001,                         # number of epochs
    print_frequency=20,                 # number of iterations before printing to log file
    epochs_per_checkpoint=50,           # number of epochs for each checkpoint
    output_path="./log/waveglow/",      # directory to save results

    weight_decay=0.0,                   # weight decay
    max_norm=3.4028234663852886e+38,    # maximum norm used in clip_grad_by_norm
    alpha=1e-4,                         # learning rate
    anneal_factor=0.1,                  # factor by which to anneal the learning rate
    anneal_steps=()                     # epoch at which to anneal the learning rate
)


# speech-synthesis/WaveGlow/main.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os
from pathlib import Path

import sys
sys.path.append(str(Path().cwd().parents[1] / 'utils'))
from neu.comm import CommunicatorWrapper
from neu.tts.optimizer import Optimizer

import nnabla as nn
from nnabla.ext_utils import get_extension_context
from nnabla.logger import logger
from nnabla.utils.data_iterator import data_iterator
from nnabla_ext.cuda import StreamEventHandler
import numpy as np

from dataset import LJSpeechDataSource
from hparams import hparams as hp
from model.model import WaveGlow
from train import WaveGlowTrainer
from scheduler import AnnealingScheduler


def run(args):
    """Runs the algorithm."""
    Path(hp.output_path).mkdir(parents=True, exist_ok=True)

    # setup nnabla context
    ctx = get_extension_context(args.context, device_id='0')
    nn.set_default_context(ctx)
    hp.comm = CommunicatorWrapper(ctx)
    hp.event = StreamEventHandler(int(hp.comm.ctx.device_id))

    if hp.comm.n_procs > 1 and hp.comm.rank == 0:
        n_procs = hp.comm.n_procs
        logger.info(f'Distributed training with {n_procs} processes.')

    rng = np.random.RandomState(hp.seed)

    # train data
    train_loader = data_iterator(
        LJSpeechDataSource('metadata_train.csv', hp, shuffle=True, rng=rng),
        batch_size=hp.batch_size, with_memory_cache=False
    )
    # valid data
    valid_loader = data_iterator(
        LJSpeechDataSource('metadata_valid.csv', hp, shuffle=False, rng=rng),
        batch_size=hp.batch_size, with_memory_cache=False
    )
    dataloader = dict(train=train_loader, valid=valid_loader)
    model = WaveGlow(hp)
    # setup optimizer
    anneal_steps = [x*(train_loader.size//hp.batch_size)
                    for x in hp.anneal_steps]
    lr_scheduler = AnnealingScheduler(hp.alpha, anneal_steps=anneal_steps,
                                      anneal_factor=hp.anneal_factor)
    optimizer = Optimizer(
        weight_decay=hp.weight_decay,
        max_norm=hp.max_norm,
        lr_scheduler=lr_scheduler,
        name='Adam', alpha=hp.alpha
    )

    WaveGlowTrainer(model, dataloader, optimizer, hp).run()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--context', '-c', type=str, default='cudnn',
                        help="Extension module. 'cudnn' is highly recommended.")
    parser.add_argument("--device-id", "-d", type=str, default='-1',
                        help='A list of device ids to use.\
                        This is only valid if you specify `-c cudnn`. \
                        Defaults to use all available GPUs.')
    args = parser.parse_args()

    # setup context for nnabla
    if args.device_id != '-1':
        os.environ["CUDA_VISIBLE_DEVICES"] = args.device_id

    run(args)


# speech-synthesis/WaveGlow/scheduler.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import nnabla.solvers as S
from nnabla.utils.learning_rate_scheduler import BaseLearningRateScheduler


class AnnealingScheduler(BaseLearningRateScheduler):
    """Annealing Learning Rate Scheduler.

    Args:
        init_lr (float): Initial learning rate.
        warmup (int): The number of iterations for warm-up.
        anneal_steps (list of int, optional): Iteration at which to anneal the learning rate.
            Defaults to ().
        anneal_factor (float, optional): Factor by which to anneal the learning rate.
            Defaults to 0.1.
    """

    def __init__(self, init_lr, anneal_steps=(), anneal_factor=0.1):
        self.init_lr = init_lr
        self.anneal_steps = anneal_steps
        self.anneal_factor = anneal_factor

    def get_learning_rate(self, cur_iter):
        r"""Get learning rate based on current iteration.

        Args:
            iter (int): Current iteration (starting with 0).

        Returns:
            float: Learning rate
        """
        step = cur_iter + 1
        p = len([x for x in self.anneal_steps if step >= x])
        lr = self.init_lr*(self.anneal_factor ** p)
        return lr


# speech-synthesis/WaveGlow/synthesize.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import argparse
import os
from pathlib import Path

import sys
sys.path.append(str(Path().cwd().parents[1] / 'utils'))

import nnabla as nn
from nnabla.ext_utils import get_extension_context
import numpy as np
from scipy.io import wavfile

from hparams import hparams as hp
from model.denoiser import Denoiser
from model.model import WaveGlow


def synthesize(args):
    hp.batch_size = 1
    # load the model
    model = WaveGlow(hp)
    model.training = False
    model.load_parameters(args.f_model, raise_if_missing=True)

    # build a denoiser
    denoise = Denoiser(model, hp)

    x_mel = nn.Variable.from_numpy_array(np.load(args.f_mel))
    o_aud = model.infer(x_mel)

    o_aud.forward(clear_buffer=True)
    wave = denoise(o_aud.d[0].copy(), strength=args.noise_level)
    wavfile.write(args.f_output, rate=hp.sr, data=wave)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--context', '-c', type=str, default='cudnn',
                        help="Extension module. 'cudnn' is highly recommended.")
    parser.add_argument("--device-id", "-d", type=str, default='-1',
                        help='A list of device ids to use.\
                        This is only valid if you specify `-c cudnn`. \
                        Defaults to use all available GPUs.')
    parser.add_argument("--noise-level", "-l", type=float, default=0.003,
                        help='Noise level.')
    parser.add_argument("--f-model", "-m", type=str,
                        help='File path to the trained model.')
    parser.add_argument("--f-mel", "-f", type=str,
                        help='File path to the mel file.')
    parser.add_argument("--f-output", "-o", type=str, default='sample.wav',
                        help='File path to the synthetic output waveform.')
    args = parser.parse_args()

    # setup context for nnabla
    if args.device_id != '-1':
        os.environ["CUDA_VISIBLE_DEVICES"] = args.device_id

    # setup nnabla context
    ctx = get_extension_context(args.context, device_id='0')
    nn.set_default_context(ctx)

    synthesize(args)


# speech-synthesis/WaveGlow/train.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from abc import ABC
from pathlib import Path

import nnabla as nn
import nnabla.functions as F
import numpy as np
from scipy.io import wavfile
from tqdm import trange

from neu.tts.trainer import Trainer


class WaveGlowTrainer(Trainer):
    r"""Trainer for WaveGlow."""

    def update_graph(self, key='train'):
        r"""Builds the graph and update the placeholder.

        Args:
            key (str, optional): Type of computational graph. Defaults to 'train'.
        """
        assert key in ('train', 'valid')

        self.model.training = key != 'valid'
        hp = self.hparams

        # define input variables
        x_aud = nn.Variable([hp.batch_size, hp.segment_length])

        # output variables
        o_aud, log_s, log_det = self.model(x_aud)
        o_aud = o_aud.apply(persistent=True)
        l_det = sum([F.sum(x) for x in log_det])
        l_log = sum([F.sum(x) for x in log_s])

        l_net = (F.sum(o_aud*o_aud)/(2*hp.sigma**2) -
                 l_det - l_log) / np.prod(o_aud.shape)
        l_net = l_net.apply(persistent=True)

        x_mel = self.model.compute_mel(x_aud)
        s_aud = self.model.infer(x_mel)

        self.placeholder[key] = {'x_aud': x_aud,
                                 'o_aud': o_aud, 'l_net': l_net, 's_aud': s_aud}
        self.out_variables = ['train/l_net']

    def train_on_batch(self):
        r"""Updates the model parameters."""
        batch_size = self.hparams.batch_size
        p, dl = self.placeholder['train'], self.dataloader['train']
        self.optimizer.zero_grad()
        if self.hparams.comm.n_procs > 1:
            self.hparams.event.default_stream_synchronize()
        p['x_aud'].d = dl.next()[0]
        p['l_net'].forward(clear_no_need_grad=True)
        p['l_net'].backward(clear_buffer=True)
        self.monitor.update('train/l_net', p['l_net'].d.copy(), batch_size)
        if self.hparams.comm.n_procs > 1:
            self.hparams.comm.all_reduce(
                self._grads, division=True, inplace=False)
            self.hparams.event.add_default_stream_event()
        self.optimizer.update()

    def valid_on_batch(self):
        r"""Performs validation."""
        batch_size = self.hparams.batch_size
        p, dl = self.placeholder['valid'], self.dataloader['valid']
        if self.hparams.comm.n_procs > 1:
            self.hparams.event.default_stream_synchronize()
        p['x_aud'].d = dl.next()[0]
        p['l_net'].forward(clear_buffer=True)
        self.loss.data += p['l_net'].d.copy() * batch_size
        self.monitor.update('valid/l_net', p['l_net'].d.copy(), batch_size)

    def callback_on_epoch_end(self):
        if self.hparams.comm.n_procs > 1:
            self.hparams.comm.all_reduce(
                [self.loss], division=True, inplace=False)
        self.loss.data /= self.dataloader['valid'].size

        if self.hparams.comm.rank == 0:
            p, hp = self.placeholder['valid'], self.hparams
            self.monitor.info(f'valid/loss={self.loss.data[0]:.5f}\n')

            if self.cur_epoch % hp.epochs_per_checkpoint == 0:
                path = Path(hp.output_path) / 'output' / \
                            f'epoch_{self.cur_epoch}'
                path.mkdir(parents=True, exist_ok=True)
                p['s_aud'].forward(clear_buffer=True)
                wavfile.write(path / 'sample.wav', rate=hp.sr,
                              data=p['s_aud'].d[0].copy())
                self.model.save_parameters(
                    str(path / f'model_{self.cur_epoch}.h5'))

        self.loss.zero()


# speech-synthesis/WaveNet/__init__.py

# Copyright 2019,2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


# speech-synthesis/WaveNet/_checkpoint_nnp_util.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import

import os
import sys

# Import save/load_checkpoint from utils
sys.path.append(os.path.join(
    os.path.dirname(__file__), '../../utils/'))
from neu.checkpoint_util import save_checkpoint, load_checkpoint
from neu.save_nnp import save_nnp


# speech-synthesis/WaveNet/args.py

# Copyright 2019,2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


def get_args(monitor_path='tmp.monitor',
             max_iter=500000,
             model_save_path=None,
             learning_rate=1e-3,
             batch_size=8,
             accum_grad=32,
             train_cachefile_dir=None,
             val_cachefile_dir=None):
    """
    Get command line arguments.

    Arguments set the default values of command line arguments.
    """
    import argparse
    import os
    if model_save_path is None:
        model_save_path = monitor_path

    parser = argparse.ArgumentParser(description="WaveNet example.")
    parser.add_argument("--data-dir", type=str, required=True)
    parser.add_argument("--batch-size", "-b", type=int, default=batch_size)
    parser.add_argument("--accum-grad", "-a", type=int, default=accum_grad,
                        help='Parameters are updated by the gradient accumulated by multiple mini-batches.')
    parser.add_argument("--learning-rate", "-l",
                        type=float, default=learning_rate)

    parser.add_argument("--use-speaker-id", action="store_true")

    parser.add_argument("--monitor-path", "-m",
                        type=str, default=monitor_path,
                        help='Path monitoring logs saved.')
    parser.add_argument("--max-iter", "-i", type=int, default=max_iter,
                        help='Max iteration of training.')
    parser.add_argument("--val-interval", "-v", type=int, default=100,
                        help='Validation interval.')
    parser.add_argument("--val-iter", "-j", type=int, default=10,
                        help='Each validation runs `val_iter mini-batch iteration.')

    parser.add_argument('--context', '-c', type=str,
                        default="cudnn", help="Extension module. 'cudnn' is highly.recommended.")
    parser.add_argument("--device-id", "-d", type=str, default='0',
                        help='Device ID the training run on. This is only valid if you specify `-c cudnn`.')
    parser.add_argument("--type-config", "-t", type=str, default='float',
                        help='Type configuration.')

    parser.add_argument("--model-save-interval", "-s", type=int, default=2000,
                        help='The interval of saving model parameters.')
    parser.add_argument("--model-save-path", "-o",
                        type=str, default=model_save_path,
                        help='Path the model parameters saved.')

    parser.add_argument("--train-cachefile-dir", "-T", type=str, default=train_cachefile_dir,
                        help='Training cache file dir. Create to use create_cache_file.py')
    parser.add_argument("--val-cachefile-dir", "-V", type=str, default=val_cachefile_dir,
                        help='Validation cache file dir. Create to use create_cache_file.py')
    parser.add_argument("--checkpoint", type=str, default=None,
                        help='path to checkpoint file')

    args = parser.parse_args()
    if not os.path.isdir(args.model_save_path):
        os.makedirs(args.model_save_path)

    return args


# speech-synthesis/WaveNet/config.py

# Copyright 2019,2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


class WavenetConfig(object):
    hidden_dims = 32
    skip_dims = 512
    kernel_size = 2
    speaker_dims = 32
    dilations = [
        1, 2, 4, 8, 16, 32, 64, 128, 256, 512,
        1, 2, 4, 8, 16, 32, 64, 128, 256, 512,
        1, 2, 4, 8, 16, 32, 64, 128, 256, 512,
        1, 2, 4, 8, 16, 32, 64, 128, 256, 512,
        1, 2, 4, 8, 16, 32, 64, 128, 256, 512
    ]


class DataConfig(object):
    sample_rate = 16000
    duration = 16000
    shift = 8000
    q_bit_len = 256


wavenet_config = WavenetConfig()
data_config = DataConfig()


# speech-synthesis/WaveNet/dataset.py

# Copyright 2019,2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import division

import os
import fnmatch
import numpy as np
import soundfile as sf
import librosa

from nnabla.logger import logger
from nnabla.utils.data_iterator import data_iterator
from nnabla.utils.data_source import DataSource

from config import data_config


def mu_law_encode(x, quantize=data_config.q_bit_len):
    """
    Applies mu-law encoding algorithm and quantization for an input signal.

    Args:
        x (numpy.ndarray): An array of a signal.
        quantize (int): A bit length of quantization.

    Returns:
         numpy.ndarray: An array of mu-law encoded signal.
    """
    if not isinstance(x, np.ndarray):
        raise ValueError("Input 'x' must be numpy.ndarray")

    if x.max() > 1. or x.min() < -1:
        raise ValueError("All values of 'x' must be between [-1, 1]")

    # The range of return values is between [0, mu] of int32

    mu = quantize - 1

    compressed = np.sign(x) * np.log(1 + mu * np.abs(x)) / \
        np.log(1 + mu)  # [-1, 1]

    # [0, mu]
    return np.asarray((compressed + 1) / 2 * mu + 0.5, dtype=np.int32)


def mu_law_decode(x, quantize=data_config.q_bit_len):
    """
    Applies dequantization and mu-law decoding algorithm for an input signal.

    Args:
        x (numpy.ndarray): A data array of mu-law encoded signal.
        quantize (int): A bit length of quantization.

    Returns:
         numpy.ndarray: An array of a signal (which is mu-law decoded).
    """
    mu = quantize - 1

    if not isinstance(x, np.ndarray):
        raise ValueError("Input 'x' must be numpy.ndarray")

    if x.max() > mu or x.min() < 0:
        raise ValueError(
            "All values of 'x' must be between [0, (quantization bit length) - 1]")

    # The range of return values is between [-1, 1] of float64

    normalized = (x / mu) * 2 - 1

    return np.asarray(np.sign(normalized) * ((1 + mu) ** np.abs(normalized) - 1) / mu, dtype=np.float64)


class LibriSpeechDataSource(DataSource):
    def __init__(self, data_dir, shuffle=True, rng=None):
        super(LibriSpeechDataSource, self).__init__(shuffle=shuffle, rng=rng)

        self._data_list = self.get_all_path_and_label(data_dir)

        speaker_list = sorted(set([x[1] for x in self._data_list]))

        self.n_speaker = len(speaker_list)
        self.speaker2id = dict(zip(speaker_list, np.arange(self.n_speaker)))

        self._size = len(self._data_list)
        self._variables = ('audio_input', 'speaker_id', 'y')

        self.duration = data_config.duration + 1

        self.reset()

    @staticmethod
    def get_all_path_and_label(data_dir):
        ret = []

        for dirpath, dirnames, filenames in os.walk(data_dir):
            flac_files = fnmatch.filter(filenames, "*.flac")
            if len(flac_files) == 0:
                continue

            # each element of ret is (path, speaker_id, script_id)
            for filename in flac_files:
                speaker_id, script_id = filename.split("-")[:2]
                ret.append(
                    (os.path.join(dirpath, filename), speaker_id, script_id))

        return ret

    def _get_data(self, position):
        index = self._indexes[position]

        flac_path, speaker, _ = self._data_list[index]

        # all values of data are between [-1, 1]
        data, sr = sf.read(flac_path)

        data, _ = librosa.effects.trim(data, top_db=20)

        # clip
        if len(data) < self.duration:
            lack = self.duration - len(data)
            before = lack // 2
            after = lack // 2 + lack % 2
            clipped = np.pad(data, pad_width=(before, after), mode="constant")

        else:
            start = np.random.randint(0, len(data) - self.duration)
            clipped = data[start:start + self.duration]

        # shape of clipped == (T,)

        quantized = mu_law_encode(clipped)

        _input = quantized[:-1].reshape(-1, 1)
        _speaker_id = self.speaker2id.get(speaker)

        _output = quantized[1:].reshape(-1, 1)

        return _input, _speaker_id, _output

    def reset(self):
        self._indexes = self._rng.permutation(
            self._size) if self._shuffle else np.arange(self._size)

        super(LibriSpeechDataSource, self).reset()


def data_iterator_librispeech(batch_size, data_dir, shuffle=True, rng=None,
                              with_memory_cache=False, with_file_cache=False):
    return data_iterator(LibriSpeechDataSource(data_dir, shuffle=shuffle, rng=rng),
                         batch_size,
                         rng,
                         with_memory_cache,
                         with_file_cache)


# speech-synthesis/WaveNet/model.py

# Copyright 2019,2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import nnabla as nn
import nnabla.functions as F
import nnabla.parametric_functions as PF

from config import wavenet_config, data_config


def causal_padding(x, kernel_size, dilation):
    # x.shape = (..., channels, time)
    padding_size = (kernel_size - 1) * dilation

    return F.pad(x, (padding_size, 0), mode="constant", constant_value=0)


class WaveNet(object):

    def __init__(self):
        self.hidden_dims = wavenet_config.hidden_dims
        self.skip_dims = wavenet_config.skip_dims
        self.kernel_size = wavenet_config.kernel_size
        self.dilations = wavenet_config.dilations

        self.output_channels = data_config.q_bit_len

    def residual_block(self, x, dilation, speaker_emb=None):
        # x.shape = (Batch, channels, Time)

        # padding
        pad = causal_padding(x, self.kernel_size, dilation)

        # gated convolution
        with nn.parameter_scope("filter"):
            f_audio = PF.convolution(pad, self.hidden_dims, kernel=(
                self.kernel_size, ), dilation=(dilation,), name="audio")

            f_speaker = PF.convolution(speaker_emb, self.hidden_dims, kernel=(
                1, ), name="speaker") if speaker_emb is not None else 0

            f = F.tanh(f_audio + f_speaker)

        with nn.parameter_scope("gate"):
            g_audio = PF.convolution(pad, self.hidden_dims, kernel=(self.kernel_size, ), dilation=(dilation,),
                                     name="audio")

            g_speaker = PF.convolution(speaker_emb, self.hidden_dims, kernel=(
                1, ), name="speaker") if speaker_emb is not None else 0

            g = F.sigmoid(g_audio + g_speaker)

        h = f * g

        # 1 x 1 convolution
        with nn.parameter_scope("skip"):
            skip = PF.convolution(h, self.skip_dims, kernel=(1, ))

        out = x + skip

        return out, skip

    def call(self, x, speaker_emb=None):
        # causal convolution
        with nn.parameter_scope("causal"):
            pad = causal_padding(x, kernel_size=2, dilation=1)
            current = PF.convolution(
                pad, self.skip_dims, kernel=(2, ), dilation=(1,))

        # residual
        skips = []
        for index, dilation in enumerate(self.dilations):
            with nn.parameter_scope("residual_{}".format(index)):
                current, skip = self.residual_block(
                    current, dilation, speaker_emb)
            skips.append(skip)

        # output
        out = F.relu(sum(skips))
        with nn.parameter_scope("out1"):
            y1 = F.relu(PF.convolution(out, self.skip_dims, kernel=(1,)))
        with nn.parameter_scope("out2"):
            y2 = PF.convolution(y1, self.output_channels, kernel=(1, ))

        return y2

    def __call__(self, x, speaker_emd=None):
        with nn.parameter_scope("WaveNet"):
            y = self.call(x, speaker_emd)

        return y


# speech-synthesis/WaveNet/train.py

# Copyright 2019,2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import

import os
import numpy as np
import librosa

import nnabla as nn
import nnabla.functions as F
import nnabla.parametric_functions as PF
import nnabla.solvers as S
import nnabla.utils.save as save

from nnabla.logger import logger
from nnabla.monitor import Monitor, MonitorSeries
from _checkpoint_nnp_util import save_checkpoint, load_checkpoint, save_nnp

from model import WaveNet
from args import get_args
from config import data_config, WavenetConfig
from dataset import data_iterator_librispeech, mu_law_decode


def save_audio(xs, iter, save_dir):
    for index, audio in enumerate(xs):
        librosa.output.write_wav(os.path.join(save_dir, "result_{}_{}.wav".format(iter, index)), audio,
                                 sr=data_config.sample_rate)


def train():
    args = get_args()

    # Set context.
    from nnabla.ext_utils import get_extension_context
    logger.info("Running in {}:{}".format(args.context, args.type_config))
    ctx = get_extension_context(args.context,
                                device_id=args.device_id,
                                type_config=args.type_config)
    nn.set_default_context(ctx)

    data_iterator = data_iterator_librispeech(args.batch_size, args.data_dir)
    _data_source = data_iterator._data_source  # dirty hack...

    # model
    x = nn.Variable(
        shape=(args.batch_size, data_config.duration, 1))  # (B, T, 1)
    onehot = F.one_hot(x, shape=(data_config.q_bit_len, ))  # (B, T, C)
    wavenet_input = F.transpose(onehot, (0, 2, 1))  # (B, C, T)

    # speaker embedding
    if args.use_speaker_id:
        s_id = nn.Variable(shape=(args.batch_size, 1))
        with nn.parameter_scope("speaker_embedding"):
            s_emb = PF.embed(s_id, n_inputs=_data_source.n_speaker,
                             n_features=WavenetConfig.speaker_dims)
            s_emb = F.transpose(s_emb, (0, 2, 1))
    else:
        s_emb = None

    net = WaveNet()
    wavenet_output = net(wavenet_input, s_emb)

    pred = F.transpose(wavenet_output, (0, 2, 1))

    # (B, T, 1)
    t = nn.Variable(shape=(args.batch_size, data_config.duration, 1))

    loss = F.mean(F.softmax_cross_entropy(pred, t))

    # for generation
    prob = F.softmax(pred)

    # Create Solver.
    solver = S.Adam(args.learning_rate)
    solver.set_parameters(nn.get_parameters())

    # load checkpoint
    start_point = 0
    if args.checkpoint is not None:
        # load weights and solver state info from specified checkpoint file.
        start_point = load_checkpoint(args.checkpoint, solver)

    # Create monitor.
    monitor = Monitor(args.monitor_path)
    monitor_loss = MonitorSeries("Training loss", monitor, interval=10)

    # setup save env.
    audio_save_path = os.path.join(os.path.abspath(
        args.model_save_path), "audio_results")
    if audio_save_path and not os.path.exists(audio_save_path):
        os.makedirs(audio_save_path)

    # save_nnp
    contents = save_nnp({'x': x}, {'y': wavenet_output}, args.batch_size)
    save.save(os.path.join(args.model_save_path,
                           'Speechsynthesis_result_epoch0.nnp'), contents)
    # Training loop.
    for i in range(start_point, args.max_iter):
        # todo: validation

        x.d, _speaker, t.d = data_iterator.next()
        if args.use_speaker_id:
            s_id.d = _speaker.reshape(-1, 1)

        solver.zero_grad()
        loss.forward(clear_no_need_grad=True)
        loss.backward(clear_buffer=True)
        solver.update()

        loss.data.cast(np.float32, ctx)
        monitor_loss.add(i, loss.d.copy())

        if i % args.model_save_interval == 0:
            prob.forward()
            audios = mu_law_decode(
                np.argmax(prob.d, axis=-1), quantize=data_config.q_bit_len)  # (B, T)
            save_audio(audios, i, audio_save_path)
            # save checkpoint file
            save_checkpoint(audio_save_path, i, solver)

    # save_nnp
    contents = save_nnp({'x': x}, {'y': wavenet_output}, args.batch_size)
    save.save(os.path.join(args.model_save_path,
                           'Speechsynthesis_result.nnp'), contents)


if __name__ == '__main__':
    train()


# stereo-depth/PSMnet/finetune.py

# Copyright 2019,2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import random
import numpy as np
import scipy.stats as sp
import time
import os
import nnabla as nn
import nnabla.functions as F
import nnabla.parametric_functions as PF
import nnabla.solvers as S
from nnabla.utils.data_iterator import data_iterator
from nnabla.utils.data_iterator import data_iterator_simple
from nnabla.ext_utils import get_extension_context
from nnabla.monitor import Monitor, MonitorSeries, MonitorTimeElapsed
from stackhourglass import psm_net
from datetime import datetime
import scipy.misc
import cv2
import csv
from nnabla.utils.image_utils import imsave

parser = argparse.ArgumentParser(description='psm_net')
parser.add_argument('--maxdisp', type=int, default=192,
                    help='max disparity')
parser.add_argument('--epochs', type=int, default=300,
                    help='number of epochs to train')
parser.add_argument('--loadmodel', default='./psmnet_trained_param_10.h5',
                    help='load model')
parser.add_argument('--savemodel', default="./backup_psm_net_kitti",
                    help='save model')
parser.add_argument('--seed', type=int, default=0, metavar='S',
                    help='random seed (default: 0)')
parser.add_argument('--batchsize_train', type=int, default=4,
                    help='batch size train')
parser.add_argument('--batchsize_test', type=int, default=2,
                    help='batch size test')
parser.add_argument('--crop_width', type=int, default=512,
                    help='set cropped image width')
parser.add_argument('--crop_height', type=int, default=256,
                    help='set cropped image height')
parser.add_argument('--im_width', type=int, default=1232,
                    help='set image width')
parser.add_argument('--im_height', type=int, default=368,
                    help='set  image height')
parser.add_argument('-g', '--gpus', type=str,
                    help='GPU IDs to be used', default="4")
parser.add_argument('--dataset', type=str,
                    help='select dataset from "SceneFlow" or "Kitti"', default="Kitti")
args = parser.parse_args()


def read_csv(path):
    f = open(path, "r")
    reader = csv.reader(f)
    header = next(reader, None)
    column = {}
    for h in header:
        column[h] = []
    for row in reader:
        for k, v in zip(header, row):
            column[k].append(v)
    return column["L"], column["R"], column["disp"]


def adjust_learning_rate(epoch):
    if epoch <= 200:
        lr = 0.001
    else:
        lr = 0.0001
    print(lr, "learning rate")
    return lr


def main():

    random.seed(args.seed)
    np.random.seed(args.seed)

    # Prepare for CUDA.
    ctx = get_extension_context('cudnn', device_id=args.gpus)
    nn.set_default_context(ctx)

    start_full_time = time.time()
    from iterator import data_iterator

    # Data list for sceneflow data set
    train_list = "./dataset/kitti_train.csv"
    test_list = "./dataset/kitti_test.csv"

    train = True
    validation = False

    # Set monitor path.
    monitor_path = './nnmonitor' + str(datetime.now().strftime("%Y%m%d%H%M%S"))

    img_left, img_right, disp_img = read_csv(train_list)
    img_left_test, img_right_test, disp_img_test = read_csv(test_list)
    train_samples = len(img_left)
    test_samples = len(img_left_test)
    train_size = int(len(img_left) / args.batchsize_train)
    test_size = int(len(img_left_test) / args.batchsize_test)

    # Create data iterator.
    data_iterator_train = data_iterator(
        train_samples, args.batchsize_train, img_left, img_right, disp_img, train, shuffle=True, dataset=args.dataset)
    data_iterator_test = data_iterator(
        test_samples, args.batchsize_test, img_left_test, img_right_test, disp_img_test, validation, shuffle=False, dataset=args.dataset)
    # Set data size
    print(train_size, test_size)
    # Clrear patameters
    nn.clear_parameters()
    # Define data shape for training.
    var_left = nn.Variable(
        (args.batchsize_train, 3, args.crop_height, args.crop_width))
    var_right = nn.Variable(
        (args.batchsize_train, 3, args.crop_height, args.crop_width))
    var_disp = nn.Variable(
        (args.batchsize_train, 1, args.crop_height, args.crop_width))
    # Define data shape for testing.
    var_left_test = nn.Variable(
        (args.batchsize_test, 3, args.im_height, args.im_width))
    var_right_test = nn.Variable(
        (args.batchsize_test, 3, args.im_height, args.im_width))
    var_disp_test = nn.Variable(
        (args.batchsize_test, 1, args.im_height, args.im_width))
    if args.loadmodel is not None:
        # Loading CNN pretrained parameters.
        nn.load_parameters(args.loadmodel)
    # === for Training ===
    # Definition of pred
    pred1, pred2, pred3 = psm_net(var_left, var_right, args.maxdisp, True)
    mask_train = F.greater_scalar(var_disp, 0)
    sum_mask = F.maximum_scalar(F.sum(mask_train), 1)
    print(sum_mask.d, "sum_mask_first")
    # Definition of loss
    loss = 0.5 * (0.5 * F.sum(F.huber_loss(pred1, var_disp)*mask_train)/(sum_mask) + 0.7 * F.sum(F.huber_loss(
        pred2, var_disp)*mask_train)/(sum_mask) + F.sum(F.huber_loss(pred3, var_disp)*mask_train)/(sum_mask))
    # === for Testing ===
    # Definition of pred
    pred_test = psm_net(var_left_test, var_right_test, args.maxdisp, False)
    var_gt = var_disp_test + F.less_equal_scalar(var_disp_test, 0) * -1
    var_pred = pred_test + F.less_equal_scalar(pred_test, 0) * -1
    E = F.abs(var_pred - var_gt)
    n_err = F.sum(F.logical_and(F.logical_and(F.greater_scalar(var_gt, 0.0), F.greater_scalar(E, 3.0)),
                                F.greater_scalar(F.div2(E, F.abs(var_gt)), 0.05)))
    n_total = F.sum(F.greater_scalar(var_gt, 0))
    test_loss = F.div2(n_err, n_total)
    # Prepare monitors.
    monitor = Monitor(monitor_path)
    monitor_train = MonitorSeries('Training loss', monitor, interval=1)
    monitor_test = MonitorSeries('Validation loss', monitor, interval=1)
    monitor_time_train = MonitorTimeElapsed(
        "Training time/epoch", monitor, interval=1)
    # Create a solver (parameter updater)
    solver = S.Adam(alpha=0.001, beta1=0.9, beta2=0.999)
    # Set Parameters
    params = nn.get_parameters()
    solver.set_parameters(params)
    params2 = nn.get_parameters(grad_only=False)
    solver.set_parameters(params2)
    for epoch in range(1, args.epochs+1):
        print('This is %d-th epoch' % (epoch))
        total_train_loss = 0
        index = 0
        lr = adjust_learning_rate(epoch)
        ### Training ###
        while index < train_size:
            # Get mini batch
            # Preprocess
            var_left.d, var_right.d, var_disp.d = data_iterator_train.next()
            loss.forward(clear_no_need_grad=True)
            # Initialize gradients
            solver.zero_grad()
            # Backward execution
            loss.backward(clear_buffer=True)
            # Update parameters by computed gradients
            solver.set_learning_rate(lr)
            solver.update()
            print('Iter %d training loss = %.3f' % (index, loss.d))
            total_train_loss += loss.d
            index += 1
        train_error = total_train_loss/train_size
        print('epoch %d total training loss = %.3f' % (epoch, train_error))
        monitor_time_train.add(epoch)
        #  ## teting ##
        total_test_loss = 0
        max_acc = 0
        index_test = 0
        while index_test < test_size:
            var_left_test.d, var_right_test.d, var_disp_test.d = data_iterator_test.next()
            test_loss.forward(clear_buffer=True)
            total_test_loss += test_loss.d
            print('Iter %d test loss = %.3f' % (index_test, test_loss.d*100))
            index_test += 1
        test_error = total_test_loss/test_size
        print('epoch %d total 3-px error in val = %.3f' %
              (epoch, test_error*100))
        if test_error > max_acc:
            max_acc = test_error*100
        print('MAX epoch %d total test error = %.3f' % (epoch, max_acc))
        # Pass validation loss to a monitor.
        monitor_test.add(epoch, test_error*100)
        # Pass training loss to a monitor.
        monitor_train.add(epoch, train_error)
        print('full training time = %.2f HR' %
              ((time.time() - start_full_time)/3600))
        # Save Parameter
        out_param_file = os.path.join(
            args.savemodel, 'psmnet_trained_param_' + str(epoch) + '.h5')
        nn.save_parameters(out_param_file)


if __name__ == '__main__':
    main()


# stereo-depth/PSMnet/inference.py

# Copyright 2019,2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import random
import numpy as np
import scipy.stats as sp
import time
import os
import nnabla as nn
import nnabla.functions as F
import nnabla.parametric_functions as PF
import nnabla.solvers as S
from nnabla.utils.data_iterator import data_iterator
from nnabla.utils.data_iterator import data_iterator_simple
from nnabla.ext_utils import get_extension_context
from nnabla.monitor import Monitor, MonitorSeries, MonitorTimeElapsed
from stackhourglass import psm_net
from datetime import datetime
import cv2
from nnabla.utils.image_utils import imsave, imread

parser = argparse.ArgumentParser(description='psm_net')
parser.add_argument('--maxdisp', type=int, default=192,
                    help='max disparity')
parser.add_argument('--loadmodel', default='./psmnet_trained_param_293.h5',
                    help='load model')
parser.add_argument('--seed', type=int, default=0, metavar='S',
                    help='random seed (default: 0)')
parser.add_argument('--im_width_sf', type=int, default=960,
                    help='set image width')
parser.add_argument('--im_height_sf', type=int, default=512,
                    help='set  image height')
parser.add_argument('--im_width_kt', type=int, default=1232,
                    help='set image width')
parser.add_argument('--im_height_kt', type=int, default=368,
                    help='set  image height')
parser.add_argument('-g', '--gpus', type=str,
                    help='GPU IDs to be used', default="0")
parser.add_argument('-l', '--left_image', type=str,
                    help='path to left input image to model')
parser.add_argument('-r', '--right_image', type=str,
                    help='path to right input image to model')
parser.add_argument('--dataset', type=str,
                    help='select dataset from "SceneFlow" or "Kitti"', default="Kitti")
parser.add_argument('--nnp', type=str, default='psmnet_kitti.nnp')
parser.add_argument('--save-nnp', type=bool, default=False)
args = parser.parse_args()


def preprocess_kitti(image_left, image_right):
    w, h = image_left.shape[1], image_left.shape[0]
    image_l = image_left[h - args.im_heigh_kt:h, w-args.im_width_kt:w, :]
    image_r = image_right[h - args.im_height_kt:h, w-args.im_width_kt:w, :]
    image_l, image_r = preprocess_common(image_l, image_r)
    return image_l, image_r


def preprocess_sceneflow(image_left, image_right):
    image_l = image_left[:args.im_height_sf, :args.im_width_sf, :]
    image_r = image_right[:args.im_height_sf, :args.im_width_sf, :]
    image_l, image_r = preprocess_common(image_l, image_r)
    return image_l, image_r


def preprocess_common(image_left, image_right):
    mean_imagenet = np.asarray([0.485, 0.456, 0.406]).astype(
        np.float32).reshape(3, 1, 1)
    std_imagenet = np.asarray([0.229, 0.224, 0.225]).astype(
        np.float32).reshape(3, 1, 1)
    image_left, image_right = np.rollaxis(
        image_left, 2), np.rollaxis(image_right, 2)
    image_left = (image_left/255).astype(np.float32)
    image_right = (image_right/255).astype(np.float32)
    image_left -= mean_imagenet
    image_left /= std_imagenet
    image_right -= mean_imagenet
    image_right /= std_imagenet
    return image_left, image_right


def main():
    ctx = get_extension_context('cudnn', device_id=args.gpus)
    nn.set_default_context(ctx)
    image_left = imread(args.left_image)
    image_right = imread(args.right_image)

    if args.dataset == 'Kitti':
        var_left = nn.Variable((1, 3, args.im_height_kt, args.im_width_kt))
        var_right = nn.Variable((1, 3, args.im_height_kt, args.im_width_kt))
        img_left, img_right = preprocess_kitti(image_left, image_right)
    elif args.dataset == 'SceneFlow':
        var_left = nn.Variable((1, 3, args.im_height_sf, args.im_width_sf))
        var_right = nn.Variable((1, 3, args.im_height_sf, args.im_width_sf))
        img_left, img_right = preprocess_sceneflow(image_left, image_right)

    var_left.d, var_right.d = img_left, img_right
    if args.loadmodel is not None:
        # Loading CNN pretrained parameters.
        nn.load_parameters(args.loadmodel)
    pred_test = psm_net(var_left, var_right, args.maxdisp, False)
    pred_test.forward(clear_buffer=True)
    pred = pred_test.d
    pred = np.squeeze(pred, axis=1)
    pred = pred[0]
    pred = 2*(pred - np.min(pred))/np.ptp(pred)-1
    imsave('stereo_depth.png', (pred + 1) * 0.5)

    print("Done")

    # Save NNP file (used in C++ inference later.).
    if args.save_nnp:
        runtime_contents = {
            'networks': [
                {'name': 'runtime',
                 'batch_size': 1,
                 'outputs': {'y0': pred_test},
                 'names': {'x0': var_left, 'x1': var_right}}],
            'executors': [
                {'name': 'runtime',
                 'network': 'runtime',
                 'data': ['x0', 'x1'],
                 'output': ['y0']}]}
        import nnabla.utils.save
        nnabla.utils.save.save(args.nnp, runtime_contents)


if __name__ == '__main__':
    main()


# stereo-depth/PSMnet/iterator.py

# Copyright 2019,2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import sys
import argparse
import csv
import nnabla as nn
from read_pfm import readPFM
from nnabla.utils.data_iterator import data_iterator_simple
from nnabla.utils.image_utils import imread
import numpy as np
import random


def data_iterator(num_examples, batch_size, img_left, img_right, img_disp, train, shuffle, dataset, rng=None):
    def dataset_load_func(i):
        # get images from the list
        image_left = imread(img_left[i]).astype('float32')
        image_right = imread(img_right[i]).astype('float32')
        # print(img_disp)
        if dataset == "SceneFlow":
            from main import parser
            args = parser.parse_args()
            image_disp, scale = readPFM(img_disp[i])
            image_disp = np.ascontiguousarray(image_disp, dtype=np.float32)
        elif dataset == "Kitti":
            from finetune import parser
            args = parser.parse_args()
            image_disp = imread(img_disp[i]).astype('float32')
        image_disp = image_disp.reshape(
            image_disp.shape[0], image_disp.shape[1], 1)

        mean_imagenet = np.asarray([0.485, 0.456, 0.406]).astype(
            np.float32).reshape(3, 1, 1)
        std_imagenet = np.asarray([0.229, 0.224, 0.225]).astype(
            np.float32).reshape(3, 1, 1)

        if train:
            w, h = image_left.shape[1], image_left.shape[0]
            th, tw = args.crop_height, args.crop_width
            x1 = random.randint(0, w - tw)
            y1 = random.randint(0, h - th)
            # crop
            image_left = image_left[y1:y1 + th, x1:x1 + tw]
            image_right = image_right[y1:y1 + th, x1:x1 + tw]
            if dataset == "Kitti":
                image_disp = np.ascontiguousarray(
                    image_disp, dtype=np.float32)/256
            image_disp = image_disp[y1:y1 + th, x1:x1 + tw]
            # normalize with mean and std
            image_left, image_right, image_disp = np.rollaxis(
                image_left, 2), np.rollaxis(image_right, 2), np.rollaxis(image_disp, 2)
            image_left = (image_left/255).astype(np.float32)
            image_right = (image_right/255).astype(np.float32)
            image_left -= mean_imagenet
            image_left /= std_imagenet
            image_right -= mean_imagenet
            image_right /= std_imagenet
        else:
            # crop
            if dataset == "SceneFlow":
                image_left = image_left[:args.im_height, :args.im_width, :]
                image_right = image_right[:args.im_height, :args.im_width, :]
                image_disp = image_disp[:args.im_height, :args.im_width, :]
            elif dataset == "Kitti":
                w, h = image_left.shape[1], image_left.shape[0]
                image_left = image_left[h -
                                        args.im_height:h, w-args.im_width:w, :]
                image_right = image_right[h -
                                          args.im_height:h, w-args.im_width:w, :]
                image_disp = image_disp[h -
                                        args.im_height:h, w-args.im_width:w, :]
                image_disp = np.ascontiguousarray(
                    image_disp, dtype=np.float32)/256
            # normalize
            image_left, image_right, image_disp = np.rollaxis(
                image_left, 2), np.rollaxis(image_right, 2), np.rollaxis(image_disp, 2)
            image_left = (image_left/255).astype(np.float32)
            image_right = (image_right/255).astype(np.float32)
            image_left -= mean_imagenet
            image_left /= std_imagenet
            image_right -= mean_imagenet
            image_right /= std_imagenet

        return image_left, image_right, image_disp
    return data_iterator_simple(dataset_load_func, num_examples, batch_size, shuffle=shuffle, rng=rng,
                                with_file_cache=False, with_memory_cache=False)


# stereo-depth/PSMnet/main.py

# Copyright 2019,2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import random
import numpy as np
import scipy.stats as sp
import time
import os
import nnabla as nn
import nnabla.functions as F
import nnabla.parametric_functions as PF
import nnabla.solvers as S
from nnabla.utils.data_iterator import data_iterator, data_iterator_csv_dataset
from nnabla.utils.data_iterator import data_iterator_simple
from nnabla.ext_utils import get_extension_context
from nnabla.monitor import Monitor, MonitorSeries, MonitorTimeElapsed
from stackhourglass import psm_net
from datetime import datetime
import scipy.misc
import cv2
import csv
from nnabla.utils.image_utils import imsave

parser = argparse.ArgumentParser(description='psm_net')
parser.add_argument('--maxdisp', type=int, default=192,
                    help='max disparity')
parser.add_argument('--epochs', type=int, default=10,
                    help='number of epochs to train')
parser.add_argument('--loadmodel', default='./psmnet_trained_param_10.h5',
                    help='load model')
parser.add_argument('--savemodel', default='./output_psm_net',
                    help='save model')
parser.add_argument('--seed', type=int, default=0, metavar='S',
                    help='random seed (default: 0)')
parser.add_argument('--batchsize_train', type=int, default=4,
                    help='batch size train')
parser.add_argument('--batchsize_test', type=int, default=2,
                    help='batch size test')
parser.add_argument('--crop_width', type=int, default=512,
                    help='set cropped image width')
parser.add_argument('--crop_height', type=int, default=256,
                    help='set cropped image height')
parser.add_argument('--im_width', type=int, default=960,
                    help='set image width')
parser.add_argument('--im_height', type=int, default=512,
                    help='set  image height')
parser.add_argument('-g', '--gpus', type=str,
                    help='GPU IDs to be used', default="4")
parser.add_argument('--dataset', type=str,
                    help='select dataset from "SceneFlow" or "Kitti"', default="SceneFlow")
args = parser.parse_args()


def read_csv(path):
    f = open(path, "r")
    reader = csv.reader(f)
    header = next(reader, None)
    column = {}
    for h in header:
        column[h] = []
    for row in reader:
        for k, v in zip(header, row):
            column[k].append(v)
    return column["L"], column["R"], column["disp"]


def main():

    random.seed(args.seed)
    np.random.seed(args.seed)

    # Prepare for CUDA.
    ctx = get_extension_context('cudnn', device_id=args.gpus)
    nn.set_default_context(ctx)

    start_full_time = time.time()
    from iterator import data_iterator

    # Data list for sceneflow data set
    train_list = "./dataset/sceneflow_train.csv"
    test_list = "./dataset/sceneflow_test.csv"
    train = True
    validation = True

    # Set monitor path.
    monitor_path = './nnmonitor' + str(datetime.now().strftime("%Y%m%d%H%M%S"))

    img_left, img_right, disp_img = read_csv(train_list)
    img_left_test, img_right_test, disp_img_test = read_csv(test_list)
    train_samples = len(img_left)
    test_samples = len(img_left_test)
    train_size = int(len(img_left) / args.batchsize_train)
    test_size = int(len(img_left_test) / args.batchsize_test)

    # Create data iterator.
    data_iterator_train = data_iterator(
        train_samples, args.batchsize_train, img_left, img_right, disp_img, train=True, shuffle=True, dataset=args.dataset)
    data_iterator_test = data_iterator(
        test_samples, args.batchsize_test, img_left_test, img_right_test, disp_img_test, train=False, shuffle=False, dataset=args.dataset)

    # Set data size

    print(train_size, test_size)

    # Define data shape for training.
    var_left = nn.Variable(
        (args.batchsize_train, 3, args.crop_height, args.crop_width))
    var_right = nn.Variable(
        (args.batchsize_train, 3, args.crop_height, args.crop_width))
    var_disp = nn.Variable(
        (args.batchsize_train, 1, args.crop_height, args.crop_width))
    # Define data shape for testing.
    var_left_test = nn.Variable(
        (args.batchsize_test, 3, args.im_height, args.im_width))
    var_right_test = nn.Variable(
        (args.batchsize_test, 3, args.im_height, args.im_width))
    var_disp_test = nn.Variable(
        (args.batchsize_test, 1, args.im_height, args.im_width))
    mask_test = nn.Variable(
        (args.batchsize_test, 1, args.im_height, args.im_width))

    if args.loadmodel is not None:
        # Loading CNN pretrained parameters.
        nn.load_parameters(args.loadmodel)

    # === for Training ===
    # Definition of pred
    pred1, pred2, pred3 = psm_net(var_left, var_right, args.maxdisp, True)
    mask_train = F.less_scalar(var_disp, args.maxdisp)
    sum_mask = F.maximum_scalar(F.sum(mask_train), 1)
    # Definition of loss
    loss = 0.5 * (0.5 * F.sum(F.huber_loss(pred1, var_disp)*mask_train)/(sum_mask) + 0.7 * F.sum(F.huber_loss(
        pred2, var_disp)*mask_train)/(sum_mask) + F.sum(F.huber_loss(pred3, var_disp)*mask_train)/(sum_mask))

    # === for Testing ===
    # Definition of pred
    mask_test = F.less_scalar(var_disp_test, args.maxdisp)
    sum_mask_test = F.maximum_scalar(F.sum(mask_test), 1)
    pred_test = psm_net(var_left_test, var_right_test, args.maxdisp, False)
    test_loss = F.sum(F.abs(pred_test - var_disp_test)*mask_test)/sum_mask_test

    # Prepare monitors.
    monitor = Monitor(monitor_path)
    monitor_train = MonitorSeries('Training loss', monitor, interval=1)
    monitor_test = MonitorSeries('Validation loss', monitor, interval=1)
    monitor_time_train = MonitorTimeElapsed(
        "Training time/epoch", monitor, interval=1)

    # Create a solver (parameter updater)
    solver = S.Adam(alpha=0.001, beta1=0.9, beta2=0.999)

    # Set Parameters
    params = nn.get_parameters()
    solver.set_parameters(params)
    params2 = nn.get_parameters(grad_only=False)
    solver.set_parameters(params2)

    for epoch in range(1, args.epochs+1):
        print('This is %d-th epoch' % (epoch))

        if validation:
            ## teting ##
            total_test_loss = 0

            index_test = 0
            while index_test < test_size:
                var_left_test.d, var_right_test.d, var_disp_test.d = data_iterator_test.next()
                test_loss.forward(clear_no_need_grad=True)
                total_test_loss += test_loss

                print('Iter %d test loss = %.3f' % (index_test, test_loss.d))
                index_test += 1
            test_error = total_test_loss/test_size
            print('epoch %d total 3-px error in val = %.3f' %
                  (epoch, test_error.d))
            # Pass validation loss to a monitor.
            monitor_test.add(epoch, test_error)

        if train:
            ## training ##
            total_train_loss = 0
            index = 0

            while index < train_size:

                # Get mini batch
                # Preprocess
                var_left.d, var_right.d, var_disp.d = data_iterator_train.next()
                loss.forward(clear_no_need_grad=True)
                # Initialize gradients
                solver.zero_grad()
                # Backward execution
                loss.backward(clear_buffer=True)
                # Update parameters by computed gradients
                solver.update()
                print('Iter %d training loss = %.3f' %
                      (index, loss.d))
                total_train_loss += loss.d
                index += 1
            train_error = total_train_loss/train_size
            monitor_time_train.add(epoch)
            print('epoch %d total training loss = %.3f' % (epoch, train_error))

            # Pass training loss to a monitor.
            monitor_train.add(epoch, train_error)
            print('full training time = %.2f HR' %
                  ((time.time() - start_full_time)/3600))

            # Save Parameter
            out_param_file = os.path.join(
                args.savemodel, 'psmnet_trained_param_' + str(epoch) + '.h5')
            nn.save_parameters(out_param_file)


if __name__ == '__main__':
    main()


# stereo-depth/PSMnet/read_pfm.py

# Copyright 2019,2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import re
import numpy as np
import sys


def readPFM(file):
    with open(file, "rb") as file:

        color = None
        width = None
        height = None
        scale = None
        endian = None
        header = file.readline().rstrip().decode('latin-1')
        if header == 'PF':
            color = True
        elif header == 'Pf':
            color = False
        else:
            raise Exception('Not a PFM file.')

        dim_match = re.match(r'^(\d+)\s(\d+)\s$',
                             file.readline().decode('latin-1'))
        if dim_match:
            width, height = map(int, dim_match.groups())
        else:
            raise Exception('Malformed PFM header.')
        scale = float(file.readline().rstrip())
        if scale < 0:  # little-endian
            endian = '<'
            scale = -scale
        else:
            endian = '>'  # big-endian
        data = np.fromfile(file, endian + 'f')
        shape = (height, width, 3) if color else (height, width)
        data = np.reshape(data, shape)
        data = np.flipud(data)
    return data, scale


# stereo-depth/PSMnet/stackhourglass.py

# Copyright 2019,2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import nnabla as nn
import nnabla.functions as F
import nnabla.parametric_functions as PF
import numpy as np
import math
import nnabla.initializer as I


def convbn(x, planes, kernel, pad, stride, dilation, batch_stat):
    h = conv(x, planes, kernel=kernel, pad=pad, stride=stride,
             dilation=dilation, with_bias=False)
    h = PF.batch_normalization(h, batch_stat=batch_stat)
    return h


def conv(x, planes, kernel, pad, stride, dilation, with_bias):
    inchannels = x.shape[1]
    outchannels = planes
    s = I.calc_normal_std_he_backward(inchannels, outchannels, kernel)
    w_init = I.NormalInitializer(s)
    if dilation[0] > 1:
        pad2 = dilation
    else:
        pad2 = pad
    h = PF.convolution(x, planes, kernel=kernel, pad=pad2, stride=stride,
                       dilation=dilation, with_bias=with_bias, w_init=w_init)

    return h


def deconvbn(x, planes, kernel, pad, stride, batch_stat):
    h = deconv(x, planes, kernel=kernel, pad=pad,
               stride=stride, with_bias=False)
    h = PF.batch_normalization(h, batch_stat=batch_stat)
    return h


def deconv(x, planes, kernel, pad, stride, with_bias):
    h = PF.deconvolution(x, planes, kernel=kernel, pad=pad,
                         stride=stride, with_bias=with_bias)
    return h


def upsample(x, factor, training, left_shape=None):
    if len(x.shape) == 4:
        if training:
            h = F.interpolate(x, scale=(factor, factor),
                              mode='linear', align_corners=True)
        else:
            h = F.interpolate(x, output_size=(
                left_shape[2]//4, left_shape[3]//4), mode='linear', align_corners=True)
    elif len(x.shape) == 5:
        planes = x.shape[1]
        kernel_size = 2*factor - factor % 2
        stride = int(factor)
        pad = int(math.ceil((factor-1) / 2.))
        scale_factor = (kernel_size + 1) // 2
        if kernel_size % 2 == 1:
            center = scale_factor - 1
        else:
            center = scale_factor - 0.5
        bilinear_kernel = np.zeros(
            [kernel_size, kernel_size, kernel_size], dtype=np.float32)
        for i in range(kernel_size):
            for j in range(kernel_size):
                for d in range(kernel_size):
                    bilinear_kernel[i, j, d] = (1 - abs(i - center) / scale_factor) * (
                        1 - abs(j - center) / scale_factor) * (1 - abs(d - center) / scale_factor)
        w_filter = np.zeros([1, planes, kernel_size, kernel_size, kernel_size])
        for i in range(planes):
            w_filter[:, i, :, :, :] = bilinear_kernel
        h = PF.deconvolution(x, planes, kernel=(kernel_size, kernel_size, kernel_size), pad=(pad, pad, pad), stride=(
            stride, stride, stride), w_init=w_filter, fix_parameters=True, group=planes)
    return h


def conv1(x, batch_stat):
    with nn.parameter_scope('conv1'):
        cv1 = convbn(x, 64, kernel=(3, 3, 3), pad=(1, 1, 1), stride=(
            2, 2, 2), dilation=(1, 1, 1), batch_stat=batch_stat)
        cv1 = F.relu(cv1)
    return cv1


def conv2(x, batch_stat):
    with nn.parameter_scope('conv2'):
        cv2 = convbn(x, 64, kernel=(3, 3, 3), pad=(1, 1, 1), stride=(
            1, 1, 1), dilation=(1, 1, 1), batch_stat=batch_stat)
    return cv2


def conv3(x, batch_stat):
    with nn.parameter_scope('conv3'):
        cv3 = convbn(x, 64, kernel=(3, 3, 3), pad=(1, 1, 1), stride=(
            2, 2, 2), dilation=(1, 1, 1), batch_stat=batch_stat)
        cv3 = F.relu(cv3)
    return cv3


def conv4(x, batch_stat):
    with nn.parameter_scope('conv4'):
        cv4 = convbn(x, 64, kernel=(3, 3, 3), pad=(1, 1, 1), stride=(
            1, 1, 1), dilation=(1, 1, 1), batch_stat=batch_stat)
        cv4 = F.relu(cv4)
    return cv4


def conv5(x, batch_stat):
    with nn.parameter_scope('conv5'):
        cv5 = deconvbn(x, 64, kernel=(2, 2, 2), pad=(0, 0, 0), stride=(
            2, 2, 2), batch_stat=batch_stat)
    return cv5


def conv6(x, batch_stat):
    with nn.parameter_scope('conv6'):
        cv6 = deconvbn(x, 32, kernel=(2, 2, 2), pad=(0, 0, 0), stride=(
            2, 2, 2), batch_stat=batch_stat)
    return cv6


def hourglass(x, presqu, postsqu, batch_stat):
    out = conv1(x, batch_stat)  # in:1/4 out:1/8
    pre = conv2(out, batch_stat)  # in:1/8 out:1/8
    if postsqu is not None:
        pre = F.relu(pre + postsqu)
    else:
        pre = F.relu(pre)
    out = conv3(pre, batch_stat)  # in:1/8 out:1/16
    out = conv4(out, batch_stat)  # in:1/16 out:1/16
    if presqu is not None:
        post = F.relu(conv5(out, batch_stat) + presqu)  # in:1/16 out:1/8
    else:
        post = F.relu(conv5(out, batch_stat) + pre)
    out = conv6(post, batch_stat)  # in:1/8 out:1/4
    return out, pre, post


def dres0(x, batch_stat):
    with nn.parameter_scope('dres0_conv1'):
        dr0_conv = convbn(x, 32, kernel=(3, 3, 3), pad=(1, 1, 1), stride=(
            1, 1, 1), dilation=(1, 1, 1), batch_stat=batch_stat)
        dr0_conv = F.relu(dr0_conv)
    with nn.parameter_scope('dres0_conv2'):
        dr0 = convbn(dr0_conv, 32, kernel=(3, 3, 3), pad=(1, 1, 1), stride=(
            1, 1, 1), dilation=(1, 1, 1), batch_stat=batch_stat)
        dr0 = F.relu(dr0)
    return dr0


def dres1(x, batch_stat):
    with nn.parameter_scope('dres1_conv1'):
        dr1_conv1 = convbn(x, 32, kernel=(3, 3, 3), pad=(1, 1, 1), stride=(
            1, 1, 1), dilation=(1, 1, 1), batch_stat=batch_stat)
        dr1_conv1 = F.relu(dr1_conv1)
    with nn.parameter_scope('dres1_conv2'):
        dr1 = convbn(dr1_conv1, 32, kernel=(3, 3, 3), pad=(1, 1, 1), stride=(
            1, 1, 1), dilation=(1, 1, 1), batch_stat=batch_stat)
    return dr1


def classif1(x, batch_stat):
    with nn.parameter_scope('classif1_conv1'):
        cl1_conv1 = convbn(x, 32, kernel=(3, 3, 3), pad=(1, 1, 1), stride=(
            1, 1, 1), dilation=(1, 1, 1), batch_stat=batch_stat)
        cl1_conv1 = F.relu(cl1_conv1)
    with nn.parameter_scope('classif1_conv2'):
        cl1 = conv(cl1_conv1, 1, kernel=(3, 3, 3), pad=(1, 1, 1),
                   stride=(1, 1, 1), dilation=(1, 1, 1), with_bias=True)
    return cl1


def classif2(x, batch_stat):
    with nn.parameter_scope('classif2_conv1'):
        cl2_conv1 = convbn(x, 32, kernel=(3, 3, 3), pad=(1, 1, 1), stride=(
            1, 1, 1), dilation=(1, 1, 1), batch_stat=batch_stat)
        cl2_conv1 = F.relu(cl2_conv1)
    with nn.parameter_scope('classif2_conv2'):
        cl2 = conv(cl2_conv1, 1, kernel=(3, 3, 3), pad=(1, 1, 1),
                   stride=(1, 1, 1), dilation=(1, 1, 1), with_bias=True)
    return cl2


def classif3(x, batch_stat):
    with nn.parameter_scope('classif3_conv1'):
        cl3_conv1 = convbn(x, 32, kernel=(3, 3, 3), pad=(1, 1, 1), stride=(
            1, 1, 1), dilation=(1, 1, 1), batch_stat=batch_stat)
        cl3_conv1 = F.relu(cl3_conv1)
    with nn.parameter_scope('classif3_conv2'):
        cl3 = conv(cl3_conv1, 1, kernel=(3, 3, 3), pad=(1, 1, 1),
                   stride=(1, 1, 1), dilation=(1, 1, 1), with_bias=True)
    return cl3


def build_cost_volume(limg, rimg, maxdisp):
    left_stack = []
    right_stack = []
    for i in range(int(maxdisp/4)):
        sliced_limg = limg[:, :, :, i:]
        sliced_rimg = rimg[:, :, :, :limg.shape[3] - i]
        if i == 0:
            padded_limg = sliced_limg
            padded_rimg = sliced_rimg
        else:
            # Padd i pixels on the left edge
            # The shape of padded_* becomes [B, C, H, W]
            padded_limg = F.pad(sliced_limg, (i, 0))
            padded_rimg = F.pad(sliced_rimg, (i, 0))

        left_stack.append(padded_limg)
        right_stack.append(padded_rimg)

    left_stacked = F.stack(*left_stack, axis=2)  # [B, C, D, H, W]
    right_stacked = F.stack(*right_stack, axis=2)  # [B, C, D, H, W]

    cost_volume = F.concatenate(
        left_stacked, right_stacked, axis=1)  # [B, 2C, D, H, W]
    return cost_volume


def psm_net(left, right, maxdisp, training):
    print(training)
    if training:
        batch_stat = True
    else:
        batch_stat = False

    # feature extraction
    refimg_fea = feature_extraction(left, batch_stat, training)
    targetimg_fea = feature_extraction(right, batch_stat, training)

    # matching
    cost = build_cost_volume(refimg_fea, targetimg_fea, maxdisp)

    cost0 = dres0(cost, batch_stat)
    cost0 = dres1(cost0, batch_stat) + cost0

    out1, pre1, post1 = hourglass(cost0, None, None, batch_stat)
    out1 = out1 + cost0

    out2, pre2, post2 = hourglass(out1, pre1, post1, batch_stat)
    out2 = out2 + cost0

    out3, pre3, post3 = hourglass(out2, pre1, post2, batch_stat)
    out3 = out3 + cost0

    cost1 = classif1(out1, batch_stat)
    cost2 = classif2(out2, batch_stat) + cost1
    cost3 = classif3(out3, batch_stat) + cost2

    if training:
        with nn.parameter_scope('cost1_upsample'):
            cost1_upsample = upsample(cost1, 4, True)
            cost1_upsample = F.softmax(cost1_upsample, axis=2)
            pred1 = disparityregression(cost1_upsample, maxdisp)

        with nn.parameter_scope('cost2_upsample'):
            cost2_upsample = upsample(cost2, 4, True)
            cost2_upsample = F.softmax(cost2_upsample, axis=2)
            pred2 = disparityregression(cost2_upsample, maxdisp)

    with nn.parameter_scope('cost3_upsample'):
        cost3_upsample = upsample(cost3, 4, True)
        cost3_upsample = F.softmax(cost3_upsample, axis=2)
        pred3 = disparityregression(cost3_upsample, maxdisp)

    if training:
        return pred1, pred2, pred3
    else:
        return pred3


def make_layer(x, planes, blocks, pad, stride, dilation, batch_stat):
    if stride != (1, 1) or x.shape[1] != planes:
        ml = BasicBlock(x, planes, pad, stride, dilation, True, batch_stat)
    else:
        ml = BasicBlock(x, planes, pad, stride, dilation, False, batch_stat)
    for i in range(1, blocks):
        with nn.parameter_scope('bb_loop_' + str(i)):
            ml = BasicBlock(ml, planes, pad, (1, 1),
                            dilation, False, batch_stat)
    return ml


def downsample(x, planes, stride, batch_stat):
    # convbn
    ds_conv = convbn(x, planes=planes, kernel=(1, 1), pad=(
        0, 0), stride=stride, dilation=(1, 1), batch_stat=batch_stat)
    return ds_conv


def BasicBlock(x, planes, pad, stride, dilation, isDownsample, batch_stat):
    # conv1
    with nn.parameter_scope('bb_conv1'):
        bb_conv1 = convbn(x, planes, kernel=(3, 3), pad=pad, stride=stride,
                          dilation=dilation, batch_stat=batch_stat)
        bb_conv1 = F.relu(bb_conv1)
    # conv2
    with nn.parameter_scope('bb_conv2'):
        bb_conv2 = convbn(bb_conv1, planes, kernel=(3, 3), pad=pad, stride=(
            1, 1), dilation=dilation, batch_stat=batch_stat)
    # downsample
    if isDownsample is True:
        with nn.parameter_scope('bb_downsample'):
            x = downsample(x, planes, stride, batch_stat)
    bb_conv2 += x
    return bb_conv2


def feature_extraction(x, batch_stat, training):
    # conv0_1
    with nn.parameter_scope('conv0_1'):
        fe_conv01 = convbn(x, 32, kernel=(3, 3), pad=(1, 1), stride=(
            2, 2), dilation=(1, 1), batch_stat=batch_stat)
        fe_conv01 = F.relu(fe_conv01)
    # conv0_2
    with nn.parameter_scope('conv0_2'):
        fe_conv02 = convbn(fe_conv01, 32, kernel=(3, 3), pad=(1, 1), stride=(
            1, 1), dilation=(1, 1), batch_stat=batch_stat)
        fe_conv02 = F.relu(fe_conv02)
    # conv0_3
    with nn.parameter_scope('conv0_3'):
        fe_conv03 = convbn(fe_conv02, 32, kernel=(3, 3), pad=(1, 1), stride=(
            1, 1), dilation=(1, 1), batch_stat=batch_stat)
        fe_conv03 = F.relu(fe_conv03)
    # conv1_x
    with nn.parameter_scope('conv1_x'):
        fe_conv1x = make_layer(fe_conv03, 32, 3, pad=(1, 1), stride=(
            1, 1), dilation=(1, 1), batch_stat=batch_stat)
    # conv2_x
    with nn.parameter_scope('conv2_x'):
        fe_conv2x = make_layer(fe_conv1x, 64, 16, pad=(1, 1), stride=(
            2, 2), dilation=(1, 1), batch_stat=batch_stat)
    # conv3_x
    with nn.parameter_scope('conv3_x'):
        fe_conv3x = make_layer(fe_conv2x, 128, 3, pad=(1, 1), stride=(
            1, 1), dilation=(1, 1), batch_stat=batch_stat)
    # conv4_x
    with nn.parameter_scope('conv4_x'):
        fe_conv4x = make_layer(fe_conv3x, 128, 3, pad=(1, 1), stride=(
            1, 1), dilation=(2, 2), batch_stat=batch_stat)
    # SPP
    with nn.parameter_scope('SPP'):
        fe = SPP(fe_conv4x, fe_conv2x, batch_stat, training, x.shape)

    return fe


def branch(x, kernel, batch_stat):
    # AveragePooling
    br_avg = F.average_pooling(x, kernel, kernel)
    # convbn
    br_conv = convbn(br_avg, 32, kernel=(1, 1), pad=(0, 0), stride=(
        1, 1), dilation=(1, 1), batch_stat=batch_stat)
    # ReLU
    br_conv = F.relu(br_conv)
    return br_conv


def lastconv(x, batch_stat):
    # convbn
    with nn.parameter_scope('lastconv_conv1'):
        lc_conv1 = convbn(x, 128, kernel=(3, 3), pad=(1, 1), stride=(
            1, 1), dilation=(1, 1), batch_stat=batch_stat)
        # ReLU
        lc_conv1 = F.relu(lc_conv1)
    # conv
    with nn.parameter_scope('lastconv_conv2'):
        lc = conv(lc_conv1, 32, kernel=(1, 1), pad=(0, 0),
                  stride=(1, 1), dilation=(1, 1), with_bias=True)
    return lc


def SPP(x0, x1, batch_stat, training, left_shape):

    # branch_1
    with nn.parameter_scope('branch_1_conv'):
        branch_1_avg = branch(x0, (64, 64), batch_stat)

    with nn.parameter_scope('branch_1_upsample'):
        branch_1 = upsample(branch_1_avg, 64, training, left_shape)

    # branch_2
    with nn.parameter_scope('branch_2_conv'):
        branch_2_avg = branch(x0, (32, 32), batch_stat)
    with nn.parameter_scope('branch_2_upsample'):
        branch_2 = upsample(branch_2_avg, 32, training, left_shape)

    # branch_3
    with nn.parameter_scope('branch_3_conv'):
        branch_3_avg = branch(x0, (16, 16), batch_stat)
    with nn.parameter_scope('branch_3_upsample'):
        branch_3 = upsample(branch_3_avg, 16, training, left_shape)

    # branch
    with nn.parameter_scope('branch_4_conv'):
        branch_4_avg = branch(x0, (8, 8), batch_stat)
    with nn.parameter_scope('branch_4_upsample'):
        branch_4 = upsample(branch_4_avg, 8, training, left_shape)

    # Concatenate
    spp_concat = F.concatenate(
        x1, x0, branch_4, branch_3, branch_2, branch_1, axis=1)

    # Lastconv
    with nn.parameter_scope('spp_lastconv'):
        spp = lastconv(spp_concat, batch_stat)
    return spp


def disparityregression(x, maxdisp):
    disp = nn.Variable((x.shape), need_grad=False)
    for i in range(0, maxdisp):
        disp.d[:, :, i, :, :] = i
    dispx = F.mul2(disp, x)
    out = F.sum(dispx, axis=2)
    return out


# utils/neu/__init__.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import


# utils/neu/callbacks.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import nnabla.functions as F
import nnabla.parametric_functions as PF


############################################################
# nnabla callback helpers. (used in apply_w, apply_b and so on)
############################################################


def spectral_norm_callback(dim=0):
    """
    :args
        dim: axis along which normalization is taken. default is 0, for convolution.
    :returns
        callback function.
    """

    def callback(x):
        return PF.spectral_norm(x, dim=dim)

    return callback


def weitgh_standardization_callback(dim=0):
    """
    :args
        dim: axis along which normalization is taken. default is 0, for convolution.
    :returns
        callback function.
    """

    def callback(x):
        return F.weight_standardization(x, channel_axis=dim)

    return callback


# utils/neu/checkpoint_util.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
import os
import queue

import nnabla as nn
from nnabla.logger import logger

global prev_save_paths
prev_save_paths = queue.Queue()


def save_checkpoint(path, current_iter, solvers, n_keeps=-1, split_h5_per_solver=False):
    """Saves the checkpoint file which contains the params and its state info.

        Args:
            path: Path to the directory the checkpoint file is stored in.
            current_iter: Current iteretion of the training loop.
            solvers: A dictionary about solver's info, which is like;
                     solvers = {"identifier_for_solver_0": solver_0,
                               {"identifier_for_solver_1": solver_1, ...}
                     The keys are used just for state's filenames, so can be anything.
                     Also, you can give a solver object if only one solver exists.
                     Then, the "" is used as an identifier.
            n_keeps: Number of latest checkpoints to keep. If -1, all checkpoints are kept.
                     Note that we assume save_checkpoint is called from a single line in your script.
                     When you have to call this from multiple lines, n_keeps must be -1 (you have to disable n_keeps).
            split_h5_per_solver: If True, save several h5 files for parameters.
                                 Each h5 file contains subset of parameters that each solver has.

        Examples:
            # Create computation graph with parameters.
            pred = construct_pred_net(input_Variable, ...)

            # Create solver and set parameters.
            solver = S.Adam(learning_rate)
            solver.set_parameters(nn.get_parameters())

            # If you have another_solver like,
            # another_solver = S.Sgd(learning_rate)
            # another_solver.set_parameters(nn.get_parameters())

            # Training loop.
            for i in range(start_point, max_iter):
                pred.forward()
                pred.backward()
                solver.zero_grad()
                solver.update()
                save_checkpoint(path, i, solver)

                # If you have another_solver,
                # save_checkpoint(path, i,
                      {"solver": solver, "another_solver": another})

        Notes:
            It generates the checkpoint file (.json) which is like;
            checkpoint_1000 = {
                    "":{
                        "states_path": <path to the states file>
                        "params_names":["conv1/conv/W", ...],
                        "num_update":1000
                       },
                    "current_iter": 1000
                    }

            If you have multiple solvers.
            checkpoint_1000 = {
                    "generator":{
                        "states_path": <path to the states file>,
                        "params_names":["deconv1/conv/W", ...],
                        "num_update":1000
                       },
                    "discriminator":{
                        "states_path": <path to the states file>,
                        "params_names":["conv1/conv/W", ...],
                        "num_update":1000
                       },
                    "current_iter": 1000
                    }

    """

    if isinstance(solvers, nn.solver.Solver):
        solvers = {"": solvers}

    assert isinstance(solvers, dict), \
        "`solvers` must be either Solver object or dict of { `name`: Solver }."

    checkpoint_info = dict()
    save_paths = []

    for solvername, solver_obj in solvers.items():
        prefix = "{}_".format(solvername.replace(
            "/", "_")) if solvername else ""
        partial_info = dict()

        # save solver states.
        states_fname = prefix + 'states_{}.h5'.format(current_iter)
        states_path = os.path.join(path, states_fname)
        solver_obj.save_states(states_path)
        save_paths.append(states_path)
        # save relative path to support moving a saved directory
        partial_info["states_path"] = states_fname

        # save registered parameters' name. (just in case)
        params_names = [k for k in solver_obj.get_parameters().keys()]
        partial_info["params_names"] = params_names

        # save the number of solver update.
        num_update = getattr(solver_obj.get_states()[params_names[0]], "t")
        partial_info["num_update"] = num_update

        # save parameters per solver
        if split_h5_per_solver:
            solver_params_fname = f"{solvername}_params_{current_iter}.h5"
            solver_params_path = os.path.join(path, solver_params_fname)
            nn.save_parameters(path=solver_params_path,
                               params=solver_obj.get_parameters())
            save_paths.append(solver_params_path)

            # save relative path so to support moving a saved directory
            partial_info["params_path"] = solver_params_fname

        checkpoint_info[solvername] = partial_info

    # save parameters.
    if not split_h5_per_solver:
        params_fname = 'params_{}.h5'.format(current_iter)
        params_path = os.path.join(path, params_fname)
        nn.parameter.save_parameters(params_path)
        save_paths.append(params_path)

        # save relative path so to support moving a saved directory
        checkpoint_info["params_path"] = params_fname

    checkpoint_info["current_iter"] = current_iter

    # save checkpoint
    checkpoint_fname = 'checkpoint_{}.json'.format(current_iter)
    filename = os.path.join(path, checkpoint_fname)

    with open(filename, 'w') as f:
        json.dump(checkpoint_info, f)

    logger.info("Checkpoint save (.json): {}".format(filename))
    save_paths.append(filename)

    # keep only n_keeps latest checkpoints.
    if n_keeps > 0:
        global prev_save_paths
        prev_save_paths.put(save_paths)
        if prev_save_paths.qsize() > n_keeps:
            oldest = prev_save_paths.get()
            for path in oldest:
                os.remove(path)

    return


def _get_full_path(path, base_path):
    # for backward compatibility

    # case1: path is absolute
    if os.path.exists(path):
        return path

    # case2: path is reative based on base_path
    path1 = os.path.join(base_path, path)
    if os.path.exists(path1):
        return path1

    # otherwise: raise
    raise ValueError(
        f"Given path doesn't exist. (path: {path}, base_path: {base_path})")


def load_checkpoint(path, solvers):
    """Given the checkpoint file, loads the parameters and solver states.

        Args:
            path: Path to the checkpoint file.
            solvers: A dictionary about solver's info, which is like;
                     solvers = {"identifier_for_solver_0": solver_0,
                               {"identifier_for_solver_1": solver_1, ...}
                     The keys are used for retrieving proper info from the checkpoint.
                     so must be the same as the one used when saved.
                     Also, you can give a solver object if only one solver exists.
                     Then, the "" is used as an identifier.

        Returns:
            current_iter: The number of iteretions that the training resumes from.
                          Note that this assumes that the numbers of the update for
                          each solvers is the same.

        Examples:
            # Create computation graph with parameters.
            pred = construct_pred_net(input_Variable, ...)

            # Create solver and set parameters.
            solver = S.Adam(learning_rate)
            solver.set_parameters(nn.get_parameters())

            # AFTER setting parameters.
            start_point = load_checkpoint(path, solver)

            # Training loop.

        Notes:
            It requires the checkpoint file. For details, refer to save_checkpoint;
            checkpoint_1000 = {
                    "":{
                        "states_path": <path to the states file>
                        "params_names":["conv1/conv/W", ...],
                        "num_update":1000
                       },
                    "current_iter": 1000
                    }

            If you have multiple solvers.
            checkpoint_1000 = {
                    "generator":{
                        "states_path": <path to the states file>,
                        "params_names":["deconv1/conv/W", ...],
                        "num_update":1000
                       },
                    "discriminator":{
                        "states_path": <path to the states file>,
                        "params_names":["conv1/conv/W", ...],
                        "num_update":1000
                       },
                    "current_iter": 1000
                    }

    """

    if isinstance(solvers, nn.solver.Solver):
        solvers = {"": solvers}
    assert isinstance(solvers, dict), \
        "`solvers` must be either Solver object or dict of { `name`: Solver }."

    assert os.path.isfile(path), "checkpoint file not found"
    base_path = os.path.dirname(path)

    # load checkpoint
    with open(path, 'r') as f:
        checkpoint_info = json.load(f)
    logger.info("Checkpoint load (.json): {}".format(path))

    # load parameters (stored in global).
    if "params_path" in checkpoint_info:
        params_path = _get_full_path(checkpoint_info["params_path"], base_path)
        assert os.path.isfile(params_path), "parameters file not found."

        nn.parameter.load_parameters(params_path)

    for solvername, solver_obj in solvers.items():
        partial_info = checkpoint_info[solvername]
        if set(solver_obj.get_parameters().keys()) != set(partial_info["params_names"]):
            logger.warning("Detected parameters do not match.")

        # load solver states.
        states_path = _get_full_path(partial_info["states_path"], base_path)
        assert os.path.isfile(states_path), "states file not found."

        # set solver states.
        if solvername == "ema":
            try:
                solver_obj.load_states(states_path)
            except:
                logger.info("load state for ema is failed.")
        else:
            solver_obj.load_states(states_path)

        # load parameters belonging to this solver if exists
        if "params_path" in partial_info:
            solver_params_path = _get_full_path(
                partial_info["params_path"], base_path)
            nn.load_parameters(solver_params_path)

    # get current iteration. note that this might differ from the numbers of update.
    current_iter = checkpoint_info["current_iter"]

    return current_iter


# utils/neu/cli.py


import argparse


def setup_fid_cli(subparsers, callbacks):
    from neu.metrics.gan_eval import fid
    parser = fid.get_parser()
    subparsers.add_parser("fid", parents=[parser], conflict_handler='resolve')
    callbacks["fid"] = fid.main


def setup_inception_score_cli(subparsers, callbacks):
    from neu.metrics.gan_eval import inception_score
    parser = inception_score.get_parser()
    subparsers.add_parser("inception_score", parents=[
                          parser], conflict_handler='resolve')
    callbacks["inception_score"] = inception_score.main


def main():
    parser = argparse.ArgumentParser("neu_cli")
    subparsers = parser.add_subparsers(dest="command", required=True)

    callbacks = {}

    # fid
    setup_fid_cli(subparsers, callbacks)

    # inception score
    setup_inception_score_cli(subparsers, callbacks)

    args = parser.parse_args()

    if args.command not in callbacks:
        raise ValueError(
            f"The command '{args.command}' is not supported. Must be one of {callbacks.keys()}.")

    callbacks[args.command](args)


# utils/neu/comm.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import nnabla as nn

from nnabla.logger import logger


def create_float_context(ctx):
    from nnabla.ext_utils import get_extension_context
    ctx_float = get_extension_context(ctx.backend[0].split(':')[
                                      0], device_id=ctx.device_id)
    return ctx_float


class CommunicatorWrapper(object):
    def __init__(self, ctx):
        try:
            import nnabla.communicators as C
            comm = C.MultiProcessDataParallelCommunicator(ctx)
            comm.init()
            self.n_procs = comm.size
            self.rank = comm.rank
            self.local_rank = comm.local_rank
            self.comm = comm
        except Exception as e:
            print(e)
            print('No communicator found. Running with a single process. If you run this with MPI processes,'
                  ' all processes will perform totally same.')
            self.n_procs = 1
            self.rank = 0
            self.local_rank = 0
            self.comm = None

        if len(ctx.device_id) > 0:
            ctx.device_id = str(int(ctx.device_id) + int(self.local_rank))
        self.ctx = ctx
        self.ctx_float = create_float_context(ctx)

        logger.info("[Communicator] Using gpu_id = {} as rank = {}".format(
            self.ctx.device_id, self.rank))

    def barrier(self):
        if self.n_procs == 1:
            # skip all reduce since no processes have to be all-reduced
            return
        self.comm.barrier()

    def broadcast(self, x):
        assert isinstance(x, nn.Variable)

        if self.n_procs == 1:
            # skip all reduce since no processes have to be all-reduced
            return

        self.comm.bcast([x.data], src=0, inplace=True)

    def all_reduce(self, params, division, inplace):
        if self.n_procs == 1:
            # skip all reduce since no processes have to be all-reduced
            return
        self.comm.all_reduce(params, division=division, inplace=inplace)

    def all_reduced_solver_update(self, solver, division=False, inplace=True):
        if self.n_procs > 1:
            params = [
                x.grad for x in solver.get_parameters().values()]
            self.all_reduce(params, division=division, inplace=inplace)

        solver.update()

    def all_reduced_solver_update_all(self, *solvers, division=False, inplace=True):
        for solver in solvers:
            self.all_reduced_solver_update(
                solver, division=division, inplace=inplace)

    def get_all_reduce_callback(self, params=None, packing_size=2 << 20):
        if self.n_procs == 1:
            return None

        if params is None:
            params = nn.get_parameters().values()

        return self.comm.all_reduce_callback([x.grad for x in params], packing_size)


# utils/neu/debug.py


from nnabla.ext_utils import import_extension_module
from nnabla.logger import logger

from time import perf_counter
from contextlib import contextmanager


@contextmanager
def timer(device: str, tag: str, ctx: str = "cuda"):
    ext = import_extension_module(ctx)

    # perform host-device sync.
    ext.device_synchronize(device)

    start = perf_counter()

    yield

    logger.info(f"[{tag}] elapsed time: {perf_counter() - start} sec")


@contextmanager
def nvtx(name):
    from nnabla_ext.cuda.nvtx import range_push, range_pop

    range_push(name)
    yield
    range_pop()


# utils/neu/gan_losses.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import nnabla.functions as F
__all__ = [
    'GanLossContainer',
    'GanLoss',
    'LsGanLoss',
    'WassersteinGanLoss',
    'GeometricGanLoss',
    'HingeGanLoss',
    'SymmetricHingeGanLoss',
    'RelativisticAverageGanLoss',
]


class GanLossContainer(object):
    '''
    A container class of GAN outputs from `GanLoss` classes.
    Attributes:
        loss_dr (~nnabla.Variable): Discriminator loss for real data.
        loss_df (~nnabla.Variable): Discriminator loss for fake data.
        loss_gr (~nnabla.Variable):
            Generator loss for real data. This is usually set as `None`
            because generator parameters are independent of real data in
            standard GAN losses. Exceptions for examples are relativistic
            GAN losses.
        loss_gf (~nnabla.Variable): Generator loss for fake data.
    Property:
        discriminator_loss (~nnabla.Variable):
            Returns `loss_dr + loss_df`. It is devided by 2 if loss_gr is
            not `None`.
        generator_loss (~nnabla.Variable): Returns `loss_gr + loss_gf`
    This class implements `+` (add) operators (radd as well) for the
    following operands.
    * `GanLossContainer + GanLossContainer`: Losses from containers are
      added up for each of loss_dr, loss_df, loss_gr, and loss_gf
      and a new GanLossContainer object is returned.
    * `GanLossContainer + None`: `None` is ignored and the given
      `GanLossContainer` is returned.
    '''

    def __init__(self, loss_dr, loss_df, loss_gr, loss_gf):
        self.loss_dr = loss_dr
        self.loss_df = loss_df
        self.loss_gr = loss_gr
        self.loss_gf = loss_gf
        self.d_div_factor = 2.0
        self._discriminator_loss = None
        self._generator_loss = None

    def set_persistent(self, b):
        self.loss_dr.persistent = b
        self.loss_df.persistent = b
        if self.loss_gr is not None:
            self.loss_gr.persistent = b
        self.loss_gf.persistent = b

    @property
    def generator_loss(self):
        if self._generator_loss is not None:
            return self._generator_loss
        self._generator_loss = self.loss_gf
        if self.loss_gr is not None:
            self._generator_loss += self.loss_gr
            self._generator_loss /= 2
        return self._generator_loss

    @property
    def discriminator_loss(self):
        if self._discriminator_loss is not None:
            return self._discriminator_loss
        self._discriminator_loss = \
            (self.loss_dr + self.loss_df) \
            / self.d_div_factor
        return self._discriminator_loss

    def __add__(self, rhs):
        if rhs is None:
            return self
        assert isinstance(rhs, GanLossContainer)
        loss_dr = self.loss_dr + rhs.loss_dr
        loss_df = self.loss_df + rhs.loss_df
        loss_gr = None
        if self.loss_gr is not None:
            loss_gr = self.loss_gr
        if rhs.loss_gr is not None:
            if loss_gr is None:
                loss_gr = rhs.loss_gr
            else:
                loss_gr += rhs.loss_gr
        loss_gf = self.loss_gf + rhs.loss_gf
        return GanLossContainer(loss_dr, loss_df, loss_gr, loss_gf)

    def __radd__(self, rhs):
        return self.__add__(rhs)


class BaseGanLoss(object):
    '''
    A base class of GAN loss functions.
    This class object offers a callable method which takes discriminator
    output variables from both real and fake, and returns a `GanLossContainer` which
    holds discreminator and generator loss values as computation graph
    variables.
    GAN loss functions for discriminator :math:`L_D` and generator :math:`L_G` can be written in the following generalized form
    .. math::
        L_D &= \mathop{\mathbb{E}}_{x_r \sim P_{\rm data}} \left[L^r_D \left(D(x_r)\right)\right] + \mathop{\mathbb{E}}_{x_f \sim G}\left[L^f_D \left( D(x_f) \right)\right] \\
        L_G &= \mathop{\mathbb{E}}_{x_r \sim P_{\rm data}} \left[L^r_G \left( D(x_r) \right)\right] + \mathop{\mathbb{E}}_{x_f \sim G} \left[L^f_G \left( D(x_f) \right)\right]
    where :math:`L^r_D` and :math:`L^r_G` are loss functions of real data :math:`x_r` sampled from dataset :math:`P_{\rm data}` for discriminator and generator respectively, and :math:`L^f_D` and :math:`L^f_G` are for fake data :math:`x_f` generated from the current generator :math:`G`. Those functions take discriminator outputs :math:`D(\cdot)` as inputs.
    In most of GAN variants (with some exceptions), those loss functions can be defined as the following symmetric form
    .. math::
        L^r_D(d) = l^+(d)
        L^f_D(d) = l^-(d)
        L^r_G(d) = l^-(d)
        L^f_G(d) = l^+(d)
    where :math:`l^+` is a loss function which encourages the discriminator
    output to be high while :math:`l^-` to be low.
    Different :math:`l^+` and :math:`l^-` give different types of GAN losses.
    For example, the Least Square GAN (LSGAN) is derived from
    .. math::
        l^+(d) &= (d - 1)^2 \\
        l^-(d) &= (d - 0)^2.
    Any derived class must implement both :math:`l^+(d)` :math:`l^-(d)` as `def _loss_plus(self, d)` and `def _loss_minus(self, d)` for :math:`l^+(d)` and :math:`l^-(d)` respectively, then the overall loss function is defined by the symmetric form explained above.
    Note:
        The loss term for real data of generator loss :math:`\mathop{\mathbb{E}}_{x_r \sim P_{\rm data}} \left[l^- \left( D(x_r) \right)\right]` is usually omitted at computation graph because generator model is not dependent of that term. If you want to obtain it as outputs, call a method `use_generator_loss_real(True)` to enable it.
    '''

    def __init__(self):

        self._use_generator_loss_for_real = False

    def use_generator_loss_for_real(self, use):
        '''
        Whether or not to compute generator loss for real data. This is
        originally set as False, because the generator model which
        we would like to train is not dependent of real data, so we don't
        have to consider the generator loss for real data as training loss.
        Args:
            use (bool):
                If True, the generator loss for real data is taken into
                account.
        '''
        self._use_generator_loss_for_real = use

    def _mean(self, loss):
        '''
        Reduction function to obtain a scalar value for each GAN loss.
        '''
        return F.mean(loss)

    def _loss_plus(self, dout):
        raise NotImplementedError('Not implemented.')

    def _loss_minus(self, dout):
        raise NotImplementedError('Not implemented.')

    def _loss_dis_real(self, dout):
        return self._loss_plus(dout)

    def _loss_dis_fake(self, dout):
        return self._loss_minus(dout)

    def _loss_gen_real(self, dout):
        return self._loss_minus(dout)

    def _loss_gen_fake(self, dout):
        return self._loss_plus(dout)

    def __call__(self, d_r, d_f):
        '''
        Get GAN losses given disriminator outputs of both real and fake.
        Args:
            d_r (~nnabla.Variable): Discriminator output of real data, `D(x_real)`.
            d_f (~nnabla.Variable): Discriminator output of fake data, `D(x_fake)`.
        Note:
            The discriminator scores which are fed into this must be
            pre-activation values, that is `[-inf, inf]`.
        Returns: GanLossContainer
        '''
        # L_D
        loss_dr = self._mean(self._loss_dis_real(d_r))
        loss_df = self._mean(self._loss_dis_fake(d_f))

        # L_G
        loss_gr = None
        if self._use_generator_loss_for_real:
            loss_gr = self._mean(self._loss_gen_real(d_r))
        loss_gf = self._mean(self._loss_gen_fake(d_f))
        return GanLossContainer(loss_dr, loss_df, loss_gr, loss_gf)


class GanLoss(BaseGanLoss):
    '''
    Standard GAN loss defined as

    .. math::
        l^+(d) &= \ln \sigma (d) \\
        l^-(d) &= \ln \left(1 - \sigma (d )\right)

    in a generalized form described in `BaseGanLoss` documentation. Here, :math:`\sigma` is Sigmoid function :math:`\sigma(d) = \frac{1}{1 + e^{-d}}` to interpret input as probability.

    References:

        `Ian J. Goodfellow et. al.
        Generative Adversarial Networks
        <https://arxiv.org/abs/1406.2661>`_

    '''

    def _loss_plus(self, dout):
        return -F.log_sigmoid(dout)

    def _loss_minus(self, dout):
        return -F.log_sigmoid(-dout)


class LsGanLoss(BaseGanLoss):
    '''
    Least Square GAN loss defined as

    .. math::
        l^+(d) &= (d - 1)^2 \\
        l^-(d) &= (d - 0)^2

    in a generalized form described in `BaseGanLoss` documentation.

    References:

        `Xudong Mao, Qing Li, Haoran Xie, Raymond Y.K. Lau, Zhen Wang, Stephen Paul Smolley.
        Least Squares Generative Adversarial Networks.
        <https://arxiv.org/abs/1611.04076>`_

    '''

    def _loss_plus(self, dout):
        return F.squared_error(dout,
                               F.constant(1., shape=dout.shape))

    def _loss_minus(self, dout):
        return F.squared_error(dout,
                               F.constant(0., shape=dout.shape))


class WassersteinGanLoss(BaseGanLoss):
    '''
    Wasserstein GAN loss defined as

    .. math::
        l^+(d) &= d \\
        l^-(d) &= -d

    in a generalized form described in `BaseGanLoss` documentation.

    References:

        `Martin Arjovsky, Soumith Chintala, Leon Bottou.
        Wasserstein GAN.
        <https://arxiv.org/abs/1701.07875>`_

    '''

    def _loss_plus(self, dout):
        return -dout

    def _loss_minus(self, dout):
        return dout


class GeometricGanLoss(BaseGanLoss):
    '''
    Geometric GAN loss with SVM hyperplane defined as

    .. math::
        L^r_D &= \max(0, 1 - d) \\
        L^f_D &= \max(0, 1 + d) \\
        L^r_G &= d \\
        L^f_G &= -d

    in a generalized form described in `BaseGanLoss` documentation, but note that it's not symmetric.

    References:

        `Jae Hyun Lim, Jong Chul Ye.
        Geometric GAN.
        <https://arxiv.org/abs/1705.02894>`_

    Note:
        This is somestimes called as Hinge GAN.

    '''

    def _loss_plus(self, dout):
        return F.relu(1. - dout)

    def _loss_minus(self, dout):
        return F.relu(1. + dout)

    def _loss_gen_real(self, dout):
        return dout

    def _loss_gen_fake(self, dout):
        return -1 * dout


class HingeGanLoss(GeometricGanLoss):
    '''
    An alias of `GeometricGanLoss`.

    '''
    pass


class SymmetricHingeGanLoss(BaseGanLoss):
    '''
    Symmetric hinge GAN loss defined as

    .. math::
        l^+ &= \max(0, 1 - d) \\
        l^- &= \max(0, 1 + d)

    in a generalized form described in `BaseGanLoss` documentation.

    The loss function of RaHingeGAN can be created by passing this to `RelativisticAverageGanLoss`.

    References:
        `Alexia Jolicoeur-Martineau.
        Relativistic Average GAN.
        <https://arxiv.org/pdf/1807.00734.pdf>`_

    '''

    def _loss_plus(self, dout):
        return F.relu(1. - dout)

    def _loss_minus(self, dout):
        return F.relu(1. + dout)


class RelativisticAverageGanLoss(object):
    '''
    Relativistic Average GAN (RaGAN) Loss.

    Args:
        gan_loss (BaseGanLoss): A GAN loss.
        average (bool): If False, averaging is omitted. Hence it becomes Relativistic GAN.

    References:
        `Alexia Jolicoeur-Martineau.
        Relativistic Average GAN.
        <https://arxiv.org/pdf/1807.00734.pdf>`_

    '''

    def __init__(self, gan_loss, average=True):
        import copy
        assert isinstance(gan_loss, BaseGanLoss)
        gan_loss = copy.copy(gan_loss)
        # Relativistic GANs require generator loss for real
        gan_loss.use_generator_loss_for_real(True)
        self._gan_loss = gan_loss
        self._average = average

    def _average_func(self, d_values):
        if not self._average:
            return d_values
        return F.mean(d_values, keepdims=True)

    def __call__(self, d_r, d_f):

        rel_d_r = d_r - self._average_func(d_f)
        rel_d_f = d_f - self._average_func(d_r)
        return self._gan_loss(rel_d_r, rel_d_f)


# utils/neu/html_creator.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

###############################################################################
# Code from
# https://github.com/NVIDIA/pix2pixHD/blob/master/util/html.py
# Modified so it fits nnabla usage more.
###############################################################################

import dominate
from dominate.tags import *
import os


def check_dir_exist_and_create(dir_path):
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)


class HtmlCreator(object):
    def __init__(self, root, page_title="no title", redirect_interval=0):
        self.root = root
        self.img_dir = os.path.join(os.path.abspath(root), "imgs")

        check_dir_exist_and_create(self.img_dir)

        self.doc = dominate.document(title=page_title)

        if redirect_interval > 0:
            with self.doc.head:
                meta(http_equiv="refresh", content=str(redirect_interval))

    def add_text(self, text):
        with self.doc:
            h3(text)

    def add_images(self, imgs, im_titles, width=512):
        t = table(border=1, style="table-layout: fixed;")
        self.doc.add(t)
        with t:
            for im, title in zip(imgs, im_titles):
                href = os.path.join("imgs", im)
                with td(style="overflow-wrap: break-word", align="center", valign="top"):
                    with p():
                        with a(href=href):
                            img(style="width:{}px".format(width), src=href)

                        br()
                        p(title)

    def save(self):
        with open(os.path.join(self.root, "index.html"), "w") as f:
            f.write(str(self.doc))


# utils/neu/initializer.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np

import nnabla.initializer as I


def w_init(x, out_dims, gain=0.02, type="xavier"):
    if type == "xavier":
        return I.NormalInitializer(sigma=I.calc_normal_std_glorot(x.shape[1], out_dims) * gain)

    raise ValueError("unsupported init type: {}.".format(type))


def pytorch_conv_init(inmaps, kernel):
    scale = 1 / np.sqrt(inmaps * np.prod(kernel))

    return I.UniformInitializer(lim=(-scale, scale))


# utils/neu/layers.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np

import nnabla as nn
import nnabla.functions as F
import nnabla.parametric_functions as PF
import nnabla.initializer as I

from .initializer import w_init
from .losses import get_gan_loss, mae
from .callbacks import spectral_norm_callback

# alias
ps = nn.parameter_scope


#################################################################
# normalizations
#################################################################
def _normalize(x, norm_type, channel_axis=1):
    if norm_type.lower() == "in":
        return F.instance_normalization(x, gamma=None, beta=None, channel_axis=channel_axis)
    elif norm_type.lower() == "bn":
        return F.batch_normalization(x, gamma=None, beta=None, mean=None, variance=None, axes=channel_axis)
    else:
        raise ValueError("unknown norm_type: {}".format(norm_type))


def spade(x, m, hidden_dim=128, kernel=(3, 3), norm_type="in"):
    """
    Spatially-Adaptive Normalization proposed in Semantic Image Synthesis with Spatially-Adaptive Normalization (https://arxiv.org/pdf/1903.07291.pdf).


    Args:
        x (nn.Variable): Input variable for spade layer.
        m (nn.Variable):
            Spatial condition variable like object_id mask segmentation.
            This is for generating adaptive scale(gamma) and adaptice bias(beta) applied after normalization.
        hidden_dim (int): Hidden dims for first convolution applied to m.
        kernel (list of int): Kernel shapes for convolutions.
        norm_type (str) : A type of normalization. ["in", "bn"] are supported now.
    """
    # x: (N, Cx, H, W), m: (N, Cm, H, W)
    assert len(x.shape) == 4 and len(m.shape) == 4

    pad = tuple(i // 2 for i in kernel)
    c_dim = x.shape[1]
    conv_args = dict(kernel=kernel, pad=pad)
    with ps("spatial_adaptive_normalization"):
        normalized = _normalize(x, norm_type)

        m = F.interpolate(m, output_size=x.shape[2:], mode="nearest")

        with ps("shared"):
            actv = F.relu(PF.convolution(
                m, hidden_dim, w_init=w_init(m, hidden_dim), **conv_args))

        with ps("gamma"):
            gamma = PF.convolution(
                actv, c_dim, w_init=w_init(actv, c_dim), **conv_args)

        with ps("beta"):
            beta = PF.convolution(
                actv, c_dim, w_init=w_init(actv, c_dim), **conv_args)

    return normalized * gamma + beta


def rescale_values(x, input_min=-1, input_max=1, output_min=0, output_max=255):
    """
    Rescale the range of values of `x` from [input_min, input_max] -> [output_min, output_max].
    """

    assert input_max > input_min
    assert output_max > output_min

    return (x - input_min) * (output_max - output_min) / (input_max - input_min) + output_min


##############################################
# GAN Generator/Discriminator
##############################################

class PatchGAN(object):
    def __init__(self, n_layers=4, base_ndf=64, n_scales=2,
                 use_sigmoid=False, use_spectral_normalization=True):
        """
        PatchGAN discriminator.

        Args:
            n_layers:
        :param base_ndf:
        :param n_scales:
        :param use_sigmoid:
        """

        self.n_layers = n_layers
        self.base_ndf = base_ndf
        self.n_scales = n_scales
        self.use_sigmoid = use_sigmoid

        self.conv_opts = dict(w_init=I.NormalInitializer(0.02))
        if use_spectral_normalization:
            self.conv_opts["apply_w"] = spectral_norm_callback(dim=0)

    def instance_norm_lrelu(self, x, alpha=0.2):
        norm = PF.instance_normalization(x, no_scale=True, no_bias=True)
        return F.leaky_relu(norm, alpha=alpha)

    def pad_conv(self, x, fdim, stride):
        h = PF.convolution(x, fdim, (4, 4), stride=stride,
                           pad=(2, 2), **self.conv_opts)

        return h

    @staticmethod
    def get_loss(real_out, real_feats, fake_out, fake_feats,
                 use_fm=True, fm_lambda=10., gan_loss_type="ls"):
        """
        Get losses from patchGAN outputs.
        If use_fm is False, F.constant(0) is returned as g_feat.

        Args:
            real_out, real_feats, fake_out, fake_feats (nn.Variable): Variables returned from PatchGAN.__call__().
            use_fm (bool): Use feature matching loss or not.
            fm_lambda (float): Coefficient applied for feature matching loss.
            gan_loss_type (str): GAN loss type. One of ["ls", "hinge"].

        usage:
            real = nn.Variable(...)
            fake = generator(...)

            d = PatchGAN(n_layers=4, n_scales=2)
            real_out, real_feats = d(real)
            fake_out, fake_feats = d(fake)
            g_gan, g_feat, d_real, d_fake = d.get_loss(real_out, real_feats, fake_out, fake_feats,
                                                       use_fm=True, fm_lambda=10., gan_loss_type="ls")

        """

        g_gan = 0
        g_feat = 0 if use_fm else F.constant(0)
        d_real = 0
        d_fake = 0

        gan_loss = get_gan_loss(gan_loss_type)

        n_disc = len(real_out)

        for disc_id in real_out.keys():
            r_out = real_out[disc_id]
            f_out = fake_out[disc_id]

            # define GAN loss
            _d_real, _d_fake, _g_gan = gan_loss(r_out, f_out)

            d_real += _d_real
            d_fake += _d_fake
            g_gan += _g_gan

            # feature matching
            if use_fm:
                assert r_out.shape == f_out.shape
                r_feats = real_feats[disc_id]
                f_feats = fake_feats[disc_id]

                for layer_id, r_feat in r_feats.items():
                    g_feat += mae(r_feat,
                                  f_feats[layer_id]) * fm_lambda / n_disc

        return g_gan, g_feat, d_real, d_fake

    def __call__(self, x):
        outs = {}
        feats = {}
        with ps("discriminator/patch_gan"):
            # Create all scale inputs first.
            inputs = [x]
            for i in range(self.n_scales - 1):
                inputs.append(F.average_pooling(
                    inputs[-1], (3, 3), (2, 2), pad=(1, 1), including_pad=False))

            for i in range(self.n_scales):
                # Get input in reverse order of its scale (from coarse to fine) to preserve discriminator indexes.
                h = inputs.pop()
                d_name = "d_{}".format(i)

                feats[d_name] = {}
                with ps(d_name):
                    # first layer
                    fdim = self.base_ndf
                    with ps("layer_0"):
                        h = self.pad_conv(h, fdim, stride=(2, 2))
                        h = F.leaky_relu(h, alpha=0.2)

                    feats[d_name]["layer_0"] = h

                    # middle
                    for j in range(1, self.n_layers - 1):
                        fdim = min(fdim * 2, 512)
                        layer_name = "layer_{}".format(j)
                        with ps(layer_name):
                            h = self.pad_conv(h, fdim, stride=(2, 2))
                            h = self.instance_norm_lrelu(h)

                        feats[d_name][layer_name] = h

                    # last 2 layers
                    fdim = min(fdim * 2, 512)
                    layer_name = "layer_{}".format(self.n_layers - 1)
                    with ps(layer_name):
                        h = self.pad_conv(h, fdim, stride=(1, 1))
                        h = self.instance_norm_lrelu(h)

                    feats[d_name][layer_name] = h

                    with nn.parameter_scope("last_layer"):
                        h = self.pad_conv(h, 1, stride=(1, 1))
                        if self.use_sigmoid:
                            h = F.sigmoid(h)

                    outs[d_name] = h

        return outs, feats


# utils/neu/learning_rate_scheduler.py

# Copyright 2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
from .misc import DictInterfaceFactory

# Creater function
_lr_sched_factory = DictInterfaceFactory()


def create_learning_rate_scheduler(cfg):
    '''
    Create a learning rate scheduler from config.

    Args:
        cfg (dict-like object):
            It must contain ``scheduler_type`` to specify a learning rate scheduler class.

    Returns:
        Learning rate scheduler object.

    Example:

    With the following yaml file (``example.yaml``),

    .. code-block:: yaml

        learning_rate_scheduler:
            scheduler_type: EpochStepLearningRateScheduler
            base_lr: 1e-2
            decay_at: [40, 65]
            decay_rate: 0.1
            power: 1  # Ignored in EpochStepLearningRateScheduler

    you can create a learning rate scheduler class as following.

    .. code-block:: python

        from neu.yaml_wrapper import read_yaml
        cfg = read_yaml('example.yaml)

        lr_sched = create_learning_rate_scheduler(cfg.learning_rate_scheduler)

    '''

    return _lr_sched_factory.call(cfg.scheduler_type, cfg)


class BaseLearningRateScheduler(object):
    '''
    Base class of Learning rate scheduler.

    This gives a current learning rate according to a scheduling logic
    implemented as a method `_get_lr` in a derived class. It internally
    holds the current epoch and the current iteration to calculate a
    scheduled learning rate. You can get the current learning rate by
    calling `get_lr`. You have to set the current epoch which will be
    used in `_get_lr` by manually calling  `set_epoch(self, epoch)`
    while it updates the current iteration when you call
    `get_lr_and_update`.

    Example:

    .. code-block:: python
        class EpochStepLearningRateScheduler(BaseLearningRateScheduler):
            def __init__(self, base_lr, decay_at=[30, 60, 80], decay_rate=0.1, warmup_epochs=5):
                self.base_learning_rate = base_lr
                self.decay_at = np.asarray(decay_at, dtype=np.int)
                self.decay_rate = decay_rate
                self.warmup_epochs = warmup_epochs

            def _get_lr(self, current_epoch, current_iter):
                # This scheduler warmups and decays using current_epoch
                # instead of current_iter
                lr = self.base_learning_rate
                if current_epoch < self.warmup_epochs:
                    lr = lr * (current_epoch + 1) / (self.warmup_epochs + 1)
                    return lr

                p = np.sum(self.decay_at <= current_epoch)
                return lr * (self.decay_rate ** p)

        def train(...):

            ...

            solver = Momentum()
            lr_sched = EpochStepLearningRateScheduler(1e-1)
            for epoch in range(max_epoch):
                lr_sched.set_epoch(epoch)
                for it in range(max_iter_in_epoch):
                    lr = lr_sched.get_lr_and_update()
                    solver.set_learning_rate(lr)
                    ...


    '''

    def __init__(self):
        self._iter = 0
        self._epoch = 0

    def set_iter_per_epoch(self, it):
        pass

    def set_epoch(self, epoch):
        '''Set current epoch number.
        '''
        self._epoch = epoch

    def get_lr_and_update(self):
        '''
        Get current learning rate and update itereation count.

        The iteration count is calculated by how many times this method is called.

        Returns: Current learning rate

        '''
        lr = self.get_lr()
        self._iter += 1
        return lr

    def get_lr(self):
        '''
        Get current learning rate according to the schedule.
        '''
        return self._get_lr(self._epoch, self._iter)

    def _get_lr(self, current_epoch, current_iter):
        '''
        Get learning rate by current iteration.

        Args:
            current_epoch(int): Epoch count.
            current_iter(int):
                Current iteration count from the beginning of training.

        Note:
            A derived class must override this method. 

        '''
        raise NotImplementedError('')


@_lr_sched_factory.register
class EpochStepLearningRateScheduler(BaseLearningRateScheduler):
    '''
    Learning rate scheduler with step decay.

    Args:
        base_lr (float): Base learning rate
        decay_at (list of ints): It decays the lr by a factor of `decay_rate`.
        decay_rate (float): See above.
        warmup_epochs (int): It performs warmup during this period.
        legacy_warmup (bool):
            We add 1 in the denominator to be consistent with previous
            implementation.

    '''

    def __init__(self, base_lr, decay_at=[30, 60, 80], decay_rate=0.1, warmup_epochs=5, legacy_warmup=False):
        super().__init__()
        self.base_learning_rate = base_lr
        self.decay_at = np.asarray(decay_at, dtype=np.int)
        self.decay_rate = decay_rate
        self.warmup_epochs = warmup_epochs
        self.legacy_warmup_denom = 1 if legacy_warmup else 0

    def _get_lr(self, current_epoch, current_iter):
        lr = self.base_learning_rate
        # Warmup
        if current_epoch < self.warmup_epochs:
            lr = lr * (current_epoch + 1) \
                / (self.warmup_epochs + self.legacy_warmup_denom)
            return lr

        p = np.sum(self.decay_at <= current_epoch)
        return lr * (self.decay_rate ** p)


@_lr_sched_factory.register
class EpochCosineLearningRateScheduler(BaseLearningRateScheduler):
    '''
    Cosine Annealing Decay with warmup.

    The learning rate gradually increases linearly towards `base_lr` during
    `warmup_epochs`, then gradually decreases with cosine decay towards 0 for
    `epochs - warmup_epochs`.


    Args:
        base_lr (float): Base learning rate
        epochs (int): See description above.
        warmup_epochs (int): It performs warmup during this period.

    '''

    def __init__(self, base_lr, epochs, warmup_epochs=5):

        from nnabla.utils.learning_rate_scheduler import CosineScheduler
        super().__init__()
        self.base_lr = base_lr
        self.epochs = epochs
        self.warmup_epochs = warmup_epochs
        self.cosine = CosineScheduler(
            self.base_lr, self.epochs - self.warmup_epochs)

    def _get_lr(self, current_epoch, current_iter):

        # Warmup
        if current_epoch < self.warmup_epochs:
            return self.base_lr * (current_epoch + 1) / self.warmup_epochs

        # Cosine decay
        return self.cosine.get_learning_rate(
            current_epoch - self.warmup_epochs)


@_lr_sched_factory.register
class PolynomialLearningRateScheduler(BaseLearningRateScheduler):
    def __init__(self, base_lr, epochs, warmup_epochs=5, power=0.1):
        super().__init__()
        self.base_lr = base_lr
        self.epochs = epochs
        self.warmup_epochs = warmup_epochs
        self.power = power
        self.iter_per_epoch = None
        self.max_iters = None
        self.warmup_iters = None
        self.poly = None

    def set_iter_per_epoch(self, it):
        from nnabla.utils.learning_rate_scheduler import PolynomialScheduler
        self.iter_per_epoch = it
        self.max_iters = it * self.epochs
        self.warmup_iters = it * self.warmup_epochs
        self.poly = PolynomialScheduler(
            self.base_lr, self.max_iters - self.warmup_iters, self.power)

    def _get_lr(self, current_epoch, current_iter):
        if self.iter_per_epoch is None:
            raise ValueError(
                'You must call a method `set_iter_per_epoch(it)` before using this scheduler.')
        if current_iter < self.warmup_iters:
            return self.base_lr * (current_iter + 1) / self.warmup_iters
        return self.poly.get_learning_rate(current_iter - self.warmup_iters)


# utils/neu/lms.py

# Copyright 2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from contextlib import contextmanager

import nnabla.logger as logger
from nnabla.lms import SwapInOutScheduler


@contextmanager
def sechdule_scope(scheduler):
    scheduler.update_pre_hook()
    yield scheduler
    scheduler.update_post_hook()


def lms_scheduler(ctx, use_lms, gpu_memory_size=8 << 30, window_length=12 << 30):
    _check_list = [x.split(":")[0] for x in ctx.backend]
    if "cudnn" not in _check_list and "cuda" not in _check_list:
        logger.warn(
            "ctx passed to scheduler doesn't have cuda/cudnn backend. lms scheduler will not be used.")
        use_lms = False

    if use_lms:
        logger.info("[LMS] gpu_memory_limit: {}GB, prefetch_window_length: {}GB".format(float(gpu_memory_size) / (1 << 30),
                                                                                        float(window_length) / (1 << 30)))

        # Change array preference so that lms works well.
        # import nnabla_ext.cuda.init as cuda_init
        # cuda_init.prefer_cpu_pinned_array()
        # cuda_init.prefer_cuda_virtual_array()
        #
        from nnabla.ext_utils import get_extension_context
        # from nnabla import set_default_context
        be, tc = ctx.backend[0].split(":")
        # ctx = get_extension_context(be, device_id=ctx.device_id, type_config=tc)
        # set_default_context(ctx)

        cpu_ctx = get_extension_context("cpu", device_id="", type_config=tc)
        return SwapInOutScheduler(cpu_ctx, ctx, gpu_memory_size, window_length)
    else:
        class DummyScheduler(object):
            function_pre_hook = None
            function_post_hook = None
            update_pre_hook = None
            update_post_hook = None

            def start_scheduling(self):
                return None

            def end_scheduling(self):
                return None

            def __enter__(self):
                return self

            def __exit__(self, exc_type, exc_val, exc_tb):
                pass

        return DummyScheduler()


# utils/neu/losses.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np

import nnabla as nn
import nnabla.functions as F


##############################################
# classification/regression
##############################################
def sigmoid_ce(logits, value, mask=None, eps=1e-5):
    # sigmoid cross entropy and reduce_mean
    sce = F.sigmoid_cross_entropy(
        logits, F.constant(val=value, shape=logits.shape))

    if mask is not None:
        assert sce.shape[:2] == mask.shape[:2]

        sce *= F.reshape(mask, sce.shape)

        return F.sum(sce) / (F.sum(mask) + eps)

    return F.mean(sce)


def softmax_ce(logits, targets, mask=None, eps=1e-5):
    # softmax cross entropy and reduce_mean
    assert logits.shape[:-1] == targets.shape[:-1]
    assert targets.shape[-1] == 1

    sce = F.softmax_cross_entropy(logits, targets)

    if mask is not None:
        assert sce.shape[:2] == mask.shape[:2]

        sce *= F.reshape(mask, sce.shape)

        return F.sum(sce) / (F.sum(mask) + eps)

    return F.mean(sce)


def mae(x, y, mask=None, eps=1e-5):
    # l1 distance and reduce mean
    ae = F.absolute_error(x, y)

    if mask is not None:
        assert ae.shape[:2] == mask.shape[:2]

        ae *= F.reshape(mask, ae.shape)

        return F.sum(ae) / (F.sum(mask) + eps)

    return F.mean(ae)


def mse(x, y, mask=None, eps=1e-5):
    # l2 distance and reduce mean
    se = F.squared_error(x, y)

    if mask is not None:
        assert se.shape[:2] == mask.shape[:2]

        se *= F.reshape(mask, se.shape)

        return F.sum(se) / (F.sum(mask) + eps)

    return F.mean(se)


###############################################
# likelihood-based losses
###############################################

def kl_snd(mu, logvar):
    # kl divergence from standard normal gaussian to the gaussian of given parameter, mu and logvar.
    # Namely, it computes D_kl(N(mu, exp(logvar)) || N(0, 1)).

    return F.sum(F.pow_scalar(mu, 2) + F.exp(logvar) - logvar - 1) / 2


"""
Functions below are ported from https://github.com/openai/guided-diffusion/blob/main/guided_diffusion/losses.py, 
and are re-implemented with nnabla.
"""


def kl_normal(mean1, logvar1, mean2, logvar2):
    # kl divergence between two gaussians.

    return 0.5 * (((mean1 - mean2) ** 2) * F.exp(-logvar2) + F.exp(logvar1 - logvar2) + logvar2 - logvar1 - 1.0)


def approx_standard_normal_cdf(x):
    """
    A fast approximation of the cumulative distibution function of the standard normal.
    """

    return 0.5 * (1.0 + F.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * F.pow_scalar(x, 3))))


def gaussian_log_likelihood(x, mean, logstd, orig_max_val=255):
    """
    Compute the log-likelihood of a Gaussian distribution for given data `x`.

    Args:
        x (nn.Variable): Target data. It is assumed that the values are ranged [-1, 1],
                         which are originally [0, orig_max_val].
        means (nn.Variable): Gaussian mean. Must be the same shape as x.
        logstd (nn.Variable): Gaussian log standard deviation. Must be the same shape as x.
        orig_max_val (int): The maximum value that x originally has before being rescaled.

    Return:
        A log probabilies of x in nats.
    """
    assert x.shape == mean.shape == logstd.shape
    centered_x = x - mean
    inv_std = F.exp(-logstd)
    half_bin = 1.0 / orig_max_val

    def clamp(val):
        # Here we don't need to clip max
        return F.clip_by_value(val, min=1e-12, max=1e8)

    # x + 0.5 (in original scale)
    plus_in = inv_std * (centered_x + half_bin)
    cdf_plus = approx_standard_normal_cdf(plus_in)
    log_cdf_plus = F.log(clamp(cdf_plus))

    # x - 0.5 (in original scale)
    minus_in = inv_std * (centered_x - half_bin)
    cdf_minus = approx_standard_normal_cdf(minus_in)
    log_one_minus_cdf_minus = F.log(clamp(1.0 - cdf_minus))

    log_cdf_delta = F.log(clamp(cdf_plus - cdf_minus))

    log_probs = F.where(
        F.less_scalar(x, -0.999),
        log_cdf_plus,  # Edge case for 0. It uses cdf for -inf as cdf_minus.
        F.where(F.greater_scalar(x, 0.999),
                # Edge case for orig_max_val. It uses cdf for +inf as cdf_plus.
                log_one_minus_cdf_minus,
                log_cdf_delta  # otherwise
                )
    )

    assert log_probs.shape == x.shape
    return log_probs

##############################################
# GAN loss
##############################################


def ls_gan_loss(r_out, f_out):
    # todo: set constant arbitrary
    # D
    d_gan_real = F.mean(F.squared_error(r_out,
                                        F.constant(1., shape=r_out.shape)))
    d_gan_fake = F.mean(F.squared_error(f_out,
                                        F.constant(0., shape=f_out.shape)))

    # G
    g_gan = F.mean(F.squared_error(f_out,
                                   F.constant(1., shape=f_out.shape)))

    return d_gan_real, d_gan_fake, g_gan


def hinge_gan_loss(r_out, f_out):
    # D
    d_gan_real = F.mean(F.relu(1. - r_out))
    d_gan_fake = F.mean(F.relu(1. + f_out))

    # G
    g_gan = -1 * F.mean(f_out)

    return d_gan_real, d_gan_fake, g_gan


def get_gan_loss(type):
    gan_loss_dict = {
        "ls": ls_gan_loss,
        "hinge": hinge_gan_loss
    }

    gan_loss = gan_loss_dict.get(type, None)

    if gan_loss is None:
        raise ValueError("unsupported gan loss type: {}".format(type))

    return gan_loss


def vgg16_perceptual_loss(fake, real):
    '''VGG perceptual loss based on VGG-16 network.

    Assuming the values in fake and real are in [0, 255].

    Features are obtained from all ReLU activations of the first convolution
    after each downsampling (maxpooling) layer
    (including the first convolution applied to an image).
    '''
    from nnabla.models.imagenet import VGG16

    class VisitFeatures(object):
        def __init__(self):
            self.features = []
            self.relu_counter = 0
            self.features_at = set([0, 2, 4, 7, 10])

        def __call__(self, f):
            # print(f.name, end='')
            if not f.name.startswith('ReLU'):
                # print('')
                return
            if self.relu_counter in self.features_at:
                self.features.append(f.outputs[0])
                # print('*', end='')
            # print('')
            self.relu_counter += 1

    # We use VGG16 model instead of VGG19 because VGG19
    # is not in nnabla.models.
    vgg = VGG16()

    def get_features(x):
        o = vgg(x, use_up_to='lastconv')
        f = VisitFeatures()
        o.visit(f)
        return f

    with nn.parameter_scope("vgg16_loss"):
        fake_features = get_features(fake)
        real_features = get_features(real)

    volumes = np.array([np.prod(f.shape)
                        for f in fake_features.features], dtype=np.float32)
    weights = volumes[-1] / volumes
    return sum([w * F.mean(F.absolute_error(ff, fr)) for w, ff, fr in zip(weights, fake_features.features, real_features.features)])


# utils/neu/lr_scheduler.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import numpy as np


class LinearDecayScheduler(object):
    """
    Linear decaying learning rate decay scheduler.

          | start_lr  iter <= start
    lr =  | end_lr    iter >= end
          | start_lr + (end_lr - start_lr) * (iter - start) / (end - start)  otherwise

    `start_lr` stands for the learning rate before decaying, and decaying starts at `start_iter` iteration.
    `end_lr` stands for the final learning rate in the end of decaying, and it is reached at `end_iter` iteration.

    """

    def __init__(self, start_lr, end_lr, start_iter, end_iter):
        self.start_lr = start_lr
        self.end_lr = end_lr
        self.start = start_iter
        self.end = end_iter

    def __call__(self, iter):
        if iter <= self.start:
            return self.start_lr

        if iter >= self.end:
            return self.end_lr

        return self.start_lr + (self.end_lr - self.start_lr) * (iter - self.start) / (self.end - self.start)


# utils/neu/misc.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys

import numpy as np
from nnabla.logger import logger


def set_random_pseed(comm):
    import nnabla as nn
    tc = comm.ctx.backend[0].split(":")[1]

    x = nn.Variable.from_numpy_array(np.random.randint(
        low=0, high=1 << 16 if tc == "half" else 1 << 30))
    comm.broadcast(x)

    from nnabla.random import set_parameter_seed
    set_parameter_seed(int(x.d))

    from nnabla.random import pseed
    logger.info(f"[rank {comm.rank}] seed: {pseed}")


def init_nnabla(conf=None, ext_name=None, device_id=None, type_config=None, random_pseed=True):
    import nnabla as nn
    from nnabla.ext_utils import get_extension_context

    from .comm import CommunicatorWrapper
    if conf is None:
        conf = AttrDict()
    if ext_name is not None:
        conf.ext_name = ext_name
    if device_id is not None:
        conf.device_id = device_id
    if type_config is not None:
        conf.type_config = type_config

    # set context
    ctx = get_extension_context(
        ext_name=conf.ext_name, device_id=conf.device_id, type_config=conf.type_config)

    # init communicator
    comm = CommunicatorWrapper(ctx)
    nn.set_default_context(comm.ctx)

    # set random seed for parameter
    if random_pseed:
        set_random_pseed(comm)

    # disable outputs from logger except rank==0
    if comm.rank > 0:
        import logging

        logger.setLevel(logging.ERROR)

    return comm


class AttrDict(dict):
    # special internal variable used for error message.
    _parent = []

    def __setattr__(self, key, value):
        if key == "_parent":
            self.__dict__["_parent"] = value
            return

        self[key] = value

    def __getattr__(self, key):
        if key not in self:
            raise AttributeError(
                "dict (AttrDict) has no chain of attributes '{}'".format(".".join(self._parent + [key])))

        if isinstance(self[key], dict):
            self[key] = AttrDict(self[key])
            self[key]._parent = self._parent + [key]

        return self[key]

    def dump_to_stdout(self):
        self.dump()

    def dump(self, file=sys.stdout, sort_keys=True):
        if not hasattr(file, "write"):
            assert isinstance(file, str)
            file = open(file, 'w')

        out = "\n================================configs================================\n"

        iterator = self.items()
        if sort_keys:
            iterator = sorted(iterator, key=lambda x: x[0])

        for k, v in iterator:
            out += "{}: {}\n".format(k, v)

        out += "======================================================================="

        print(out, file=file)


class DictInterfaceFactory(object):
    '''Creating a single dict interface of any function or class.

    Example:

    .. code-block:: python
        # Define a function.
        def foo(a, b=1, c=None):
            for k, v in locals():
                print(k, v)

        # Register the function to the factory.
        dictif = DictInterfaceFactory()
        dictif.register(foo)

        # You can call the registered function by name and a dict representing the arguments.
        cfg = dict(a=1, c='hello')
        dictif.call('foo', cfg)

        # The following will fail because the `foo` function requires `a`.
        #     cfg = dict(c='hello')
        #     dictif.call('foo', cfg)

        # Any argument not required will be just ignored.
        cfg = dict(a=1, aaa=0)
        dictif.call('foo', cfg)

        # You can also use it for class initializer (we use it as a class decorator).
        @dictif.register
        class Bar:
            def __init__(self, a, b, c=None):
                for k, v in locals():
                    print(k, v)

        bar = dictif.call('Bar', dict(a=0, b=0))

    '''

    def __init__(self):
        self._factory = {}

    def register(self, cls):
        import inspect

        # config interface function
        def func(cfg):
            sig = inspect.signature(cls)
            # Handle all arguments of the created class
            args = {}
            for p in sig.parameters.values():
                # Positional argument
                if p.default is p.empty and p.name not in cfg:
                    raise ValueError(
                        f'`{cls.__name__}`` requires an argument `{p.name}`. Not found in cfg={cfg}.')
                args[p.name] = cfg.get(p.name, p.default)
            return cls(**args)

        # Register config interface function
        self._factory[cls.__name__] = func
        return cls

    def call(self, name, cfg):
        if name in self._factory:
            return self._factory[name](cfg)
        raise ValueError(
            f'`{name}`` not found in `{list(self._factory.keys())}`.')


def makedirs(dirpath):
    if os.path.exists(dirpath):
        if os.path.isdir(dirpath):
            return
        else:
            raise ValueError(
                "{} already exists as a file not a directory.".format(dirpath))

    os.makedirs(dirpath)


def get_current_time():
    from datetime import datetime

    return datetime.now().strftime('%y%m%d_%H%M%S')


def get_iteration_per_epoch(dataset_size, batch_size, round="ceil"):
    """
    Calculate a number of iterations to see whole images in dataset (= 1 epoch).

    Args:
     dataset_size (int): A number of images in dataset
     batch_size (int): A number of batch_size.
     round (str): Round method. One of ["ceil", "floor"].

    return: int
    """
    import numpy as np

    round_func = {"ceil": np.ceil, "floor": np.floor}
    if round not in round_func:
        raise ValueError("Unknown rounding method {}. must be one of {}.".format(round,
                                                                                 list(round_func.keys())))

    ipe = float(dataset_size) / batch_size

    return int(round_func[round](ipe))


# utils/neu/mixed_precision.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from nnabla.solvers import Solver
from nnabla import Variable
from nnabla.logger import logger


DEFAULT_INITIAL_LOG_LOSS_SCALE = 20  # 2 to the power of 20


class MixedPrecisionManager(object):
    """
    Helper class for mixed precision training.
    This class automatically scales loss value to prevent overflow during backward.
    See https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html to know the fundametal about mixed training training.

    Example:

    .. code-block:: python
        from neu.mixed_precision import MixedPrecisionManager

        mpm = MixedPrecisionManager(use_fp16=True)

        x = nn.Variable(...)
        loss = model(x, ...)
        solver = S.Sgd()
        solver.set_parameters(nn.get_parameters())

        # communicator
        comm = ...

        # data iterator
        data = ...

        for i in range(max_iter):
            solver.zero_grad()

            for accum in range(accum_cnt):
                x.d = data.next()
                loss.forward(...)
                mpm.backward(loss, ...)

            # When model parallel, do all_reduce here
            comm.all_reduce(...)

            if mpm.is_grad_overflow(solver):
                # When overflow happens, do not update parameter
                # If detecting overflow, mpm decreases loss scale.
                continue 

            # After calling solver.update, mpm increases loss scale.
            mpm.update(solver)

            # do postprocesses after parameter update like ema update.

    """

    def __init__(self,
                 use_fp16,
                 initial_log_loss_scale=DEFAULT_INITIAL_LOG_LOSS_SCALE,
                 inc_scale_factor=1e-3,
                 dec_scale_factor=1):
        self.use_fp16 = use_fp16
        self.log_loss_scale = initial_log_loss_scale
        self.inc_scale_factor = inc_scale_factor
        self.dec_scale_factor = dec_scale_factor

    @property
    def loss_scale(self):
        if not self.use_fp16:
            return 1.0

        return 2 ** self.log_loss_scale

    def is_grad_overflow(self, solver: Solver) -> bool:
        # solver.check_inf_or_nan_grad() returns True if inf or nan is detected.
        if solver.check_inf_or_nan_grad():
            # skip update and make loss scale smaller
            logger.info("[MixedPrecisionManager] Detects overflow. "
                        "Update log_loss_scale from {:.3g} to {:.3g}.".format(self.log_loss_scale, self.log_loss_scale - self.dec_scale_factor))
            self.log_loss_scale -= self.dec_scale_factor
            solver.zero_grad()
            return True

        return False

    def backward(self, loss: Variable, **kwargs):
        if not self.use_fp16:
            loss.backward(**kwargs)
            return

        loss.backward(grad=self.loss_scale, **kwargs)

    def scale_grad(self, solver):
        if not self.use_fp16:
            return

        solver.scale_grad(1. / self.loss_scale)

    def update(self, solver: Solver, *, clip_grad=None, **kwargs):
        if clip_grad is not None:
            assert isinstance(
                clip_grad, float), "clip_grad must be None or float."
            solver.clip_grad_by_norm(clip_grad)

        solver.update(**kwargs)

        if self.use_fp16:
            # increment loss scale
            self.log_loss_scale += self.inc_scale_factor


# utils/neu/pointnet2_utils.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Optional, Tuple, Iterable
import numpy as np

import nnabla as nn
import nnabla.functions as F
import nnabla.parametric_functions as PF


def compute_square_distance(source_point_cloud: nn.Variable, destination_point_cloud: nn.Variable) -> nn.Variable:
    """compute square distance between given point clouds

    Args:
        source_point_cloud (nn.Variable): shape (batch_size_source, num_points_source, dim_source)
        destination_point_cloud (nn.Variable): shape (batch_size_destination, num_points_destination, dim_destination)

    Returns:
        nn.Variable: distance from source point cloud to destination_point_cloud
            (batch_size, num_points_source, num_points_destination)
    """
    batch_size_source, num_points_source, dim_source = source_point_cloud.shape
    batch_size_destination, num_points_destination, dim_destination = destination_point_cloud.shape

    assert dim_source == 3
    assert dim_destination == 3
    assert batch_size_source == batch_size_destination

    # distance.shape = (batch_size, num_points_source, num_points_destination)
    source_squared = F.sum(
        source_point_cloud**2, axis=2).reshape((batch_size_source, num_points_source, 1))
    destination_squared = F.sum(destination_point_cloud**2, axis=2).reshape(
        (batch_size_destination, 1, num_points_destination)
    )
    both_squared = -2.0 * \
        F.batch_matmul(source_point_cloud, F.transpose(
            destination_point_cloud, (0, 2, 1)))

    return source_squared + destination_squared + both_squared


def farthest_point_sample(point_cloud: nn.Variable, num_samples: int) -> Tuple[nn.Variable, nn.Variable]:
    """sample farthest points from given point clouds

    Args:
        point_cloud (nn.Variable): shape (batch_size, num_points, 3)
        num_sample_points (int): number of samples

    Return:
        Tuple[nn.Variable, nn.Variable]: sample index and sample points, shape(batch_size, num_samples, 3)
    """
    batch_size, num_points, dim = point_cloud.shape
    assert dim == 3

    sample_points_idx = []
    # Initialize with huge value
    batch_distance = nn.Variable.from_numpy_array(
        np.ones((batch_size, num_points)) * 1e6)
    # Initialize with random value
    farthest_points_idx = F.randint(0, high=num_points, shape=(batch_size,))

    for _ in range(num_samples):
        sample_points_idx.append(farthest_points_idx.reshape((batch_size, 1)))

        # new_base_point.shape = (batch_size, 1, 3)
        new_base_point = index_batch_variable(
            point_cloud, farthest_points_idx.reshape((batch_size, 1)))

        distance = F.sum((point_cloud - new_base_point) ** 2, axis=2)

        # mask.shape = (batch_size, num_points)
        mask = F.less(distance, batch_distance)
        not_mask = F.logical_not(mask)
        batch_distance = batch_distance * not_mask + distance * mask

        farthest_points_idx = F.max(batch_distance, axis=1, only_index=True)

    # sample_points_idx.shape = (batch_size, num_samples)
    sample_points_idx = F.concatenate(*sample_points_idx, axis=1)
    sample_points_idx.need_grad = False

    sample_points = index_batch_variable(point_cloud, sample_points_idx)
    sample_points.need_grad = False
    return sample_points, sample_points_idx


def query_ball_point(
    radius: float, num_samples: int, point_cloud: nn.Variable, query_point_cloud: nn.Variable
) -> Tuple[nn.Variable, nn.Variable]:
    """compute ball point centered on query_point_cloud

    Args:
        radius (float): local region radius
        num_samples (int): max sample number in local region
        point_cloud (nn.Variable): shape (batch_size, num_points, 3)
        query_point_cloud (nn.Variable): shape (batch_size, num_query_points, 3)

    Return:
        Tuple[nn.Variable, nn.Variable]: sample_group_points shape (batch_size, num_queries, num_samples),
            sorted_group_indices, shape (batch_size, num_queries, num_samples)
    """
    batch_size, num_points, dim = point_cloud.shape
    assert dim == 3

    query_batch_size, num_queries, dim = query_point_cloud.shape
    assert dim == 3
    assert batch_size == query_batch_size
    assert num_samples <= num_points

    # group_indices.shape = (batch_size, num_queries, num_points)
    group_indices_array = np.tile(np.arange(num_points).reshape(
        (1, 1, num_points)), (batch_size, num_queries, 1))
    group_indices = nn.Variable.from_numpy_array(
        group_indices_array.astype(np.int32))

    # compute distance from query point cloud
    # square_distance.shape = (batch_size, num_queries, num_points)
    squared_distance = compute_square_distance(query_point_cloud, point_cloud)

    # if greater than radius**2, replace it with max index
    max_indices_array = np.ones(
        (batch_size, num_queries, num_points)) * num_points
    max_indices_var = nn.Variable.from_numpy_array(
        max_indices_array.astype(np.int32))
    greater_mask = F.greater_scalar(squared_distance, radius**2)
    greater_not_mask = F.logical_not(greater_mask)
    group_indices = group_indices * greater_not_mask + max_indices_var * greater_mask

    sorted_group_indices = F.sort(group_indices, axis=2)
    sorted_group_indices = sorted_group_indices[:, :, :num_samples]

    group_first_indices = F.tile(
        sorted_group_indices[:, :, 0].reshape(
            (batch_size, num_queries, 1)), (1, 1, num_samples)
    )

    max_index_mask = F.equal_scalar(sorted_group_indices, num_points)
    max_index_not_mask = F.logical_not(max_index_mask)
    sorted_group_indices = group_first_indices * \
        max_index_mask + sorted_group_indices * max_index_not_mask
    sorted_group_indices.need_grad = False

    # sample (batch_size, num_queries, num_samples, 3)
    sample_group_points = index_batch_variable(
        point_cloud, sorted_group_indices.reshape(
            (batch_size, num_queries * num_samples))
    )
    sample_group_points = F.reshape(
        sample_group_points, (batch_size, num_queries, num_samples, 3))
    sample_group_points.need_grad = False
    return sample_group_points, sorted_group_indices


def index_batch_variable(batch_variable, batch_indices) -> nn.Variable:
    """index batch variable according to given batch indices

    Args:
        point_cloud (nn.Variable): shape (batch_size, num_points, dim)
        indices (nn.Variable): shape (batch_size, num_samples)

    Returns:
        nn.Variable: shape (batch_size, num_samples, dim)
    """
    batch_size, num_points, dim = batch_variable.shape

    indice_batch_size, num_indices = batch_indices.shape
    assert batch_size == indice_batch_size

    arange_array = np.tile(np.arange(batch_size)[:, np.newaxis], (1, num_indices))[
                           np.newaxis, :, :]
    arange_indices = nn.Variable.from_numpy_array(arange_array)

    indices = F.concatenate(arange_indices, batch_indices.reshape(
        (1, batch_size, num_indices)), axis=0)
    # indexed_variable.shape = (batch_size, num_samples, dim)
    indexed_variable = F.gather_nd(batch_variable, indices)

    return indexed_variable


def sample_and_group_all(point_cloud: nn.Variable, point_feature: nn.Variable) -> Tuple[nn.Variable, nn.Variable]:
    """sample and group all points

    Args:
        point_cloud (nn.Variable): shape (batch_size, num_points, 3)
        point_feature (nn.Variable): shape (batch_size, num_points, feature_dim)

    Returns:
        Tuple[nn.Variable, nn.Variable]: new_point_cloud (all elements are zeros), new_point_feature

    """
    batch_size, num_points, dim = point_cloud.shape
    assert dim == 3

    feature_batch_size, num_features, feature_dim = point_feature.shape
    assert num_features == num_points
    assert batch_size == feature_batch_size

    new_point_cloud_array = np.zeros((batch_size, 1, dim))
    new_point_cloud = nn.Variable.from_numpy_array(new_point_cloud_array)

    grouped_point_cloud = F.reshape(
        point_cloud, (batch_size, 1, num_points, dim))

    if point_feature is not None:
        new_point_feature = F.concatenate(
            grouped_point_cloud, point_feature.reshape((batch_size, 1, num_points, -1)), axis=3
        )
    else:
        new_point_feature = grouped_point_cloud

    return new_point_cloud, new_point_feature


def sample_and_group(
    num_new_point_cloud: int,
    radius: float,
    num_samples: int,
    point_cloud: nn.Variable,
    point_feature: Optional[nn.Variable],
) -> Tuple[nn.Variable, nn.Variable]:
    """sample and group points

     Args:
        num_new_point_clouds (int): number of new points
        radius (float): radius for query ball point
        num_samples (int): num samples for query ball point
        point_cloud (nn.Variable): shape (batch_size, num_points, 3)
        point_feature (nn.Variable): shape (batch_size, num_points, feature_dim)

    Returns:
        Tuple[nn.Variable, nn.Variable]: new_point_cloud shape (batch_size, num_new_point_cloud, dim), \
            new_point_feature (batch_size, num_new_point_cloud, num_samples, feature_dim+dim)
    """
    batch_size, num_points, dim = point_cloud.shape
    assert dim == 3
    if point_feature is not None:
        feature_batch_size, num_features, feature_dim = point_feature.shape

    new_point_cloud, _ = farthest_point_sample(
        point_cloud, num_new_point_cloud)
    # grouped_point_cloud.shape = (batch_size, num_new_point_cloud, num_query_ball_samples, 3)
    grouped_point_cloud, grouped_idx = query_ball_point(
        radius, num_samples, point_cloud, new_point_cloud)
    grouped_point_cloud = grouped_point_cloud - \
        F.reshape(new_point_cloud, (batch_size, num_new_point_cloud, 1, dim))

    if point_feature is not None:
        # grouped_feature.shape = (batch_size, num_new_point_cloud*num_samples, feature_dim)
        grouped_feature = index_batch_variable(
            point_feature, grouped_idx.reshape((batch_size, -1)))
        # grouped_feature.shape = (batch_size, num_new_point_cloud, num_samples, feature_dim)
        grouped_feature = F.reshape(
            grouped_feature, (batch_size, num_new_point_cloud, num_samples, feature_dim))
        new_point_feature = F.concatenate(
            grouped_point_cloud, grouped_feature, axis=3)
    else:
        new_point_feature = grouped_point_cloud

    return new_point_cloud, new_point_feature


def set_abstraction_msg(
    point_cloud: nn.Variable,
    point_feature: Optional[nn.Variable],
    num_new_point_cloud: int,
    radius_list: Iterable[float],
    num_samples_list: Iterable[int],
    conv_channels: Iterable[Iterable[int]],
    train: bool,
) -> Tuple[nn.Variable, nn.Variable]:
    """set abstraction layer using multi scale grouping (msg).

    Args:
        point_cloud (nn.Variable): shape (batch_size, num_points, dim)
        point_feature (Union[nn.Variable, None]): shape (batch_size, num_new_point_cloud, feature_dim)
        num_new_point_cloud (int): number of new point cloud
        radius_list (Iterable[float]): list of radius
        num_samples_list (Iterable[int]): list of number of samples
        conv_channels (Iterable[Iterable[int]]): list of convolutional channel
        train (bool): training flag

    Returns:
        Tuple[nn.Variable, nn.Variable]: new_point_cloud, shape (batch_size, len(radius_list)*num_new_point_cloud, 3)
            new_point_features, shape (batch_size, len(radius_list)*num_new_point_cloud, feature_dim)
    """
    batch_size, num_points, dim = point_cloud.shape
    assert dim == 3

    if point_feature is not None:
        feature_batch_size, num_features, feature_dim = point_feature.shape
        assert feature_batch_size == batch_size

    # sample new point cloud
    new_point_cloud, _ = farthest_point_sample(
        point_cloud, num_new_point_cloud)
    new_point_features = []

    assert len(num_samples_list) == len(radius_list)
    for i, (num_samples, radius) in enumerate(zip(num_samples_list, radius_list)):
        # grouped_idx.shape = (batch_size, num_new_point_cloud, num_samples)
        # gropued_point_cloud.shape = (batch_size, num_new_point_cloud, num_samples, 3)
        grouped_point_cloud, grouped_idx = query_ball_point(
            radius, num_samples, point_cloud, new_point_cloud)
        grouped_point_cloud = grouped_point_cloud - F.reshape(
            new_point_cloud, (batch_size, num_new_point_cloud, 1, dim)
        )

        if point_feature is not None:
            # grouped_feature.shape = (batch_size, num_new_point_cloud*num_samples, feature_dim)
            grouped_feature = index_batch_variable(
                point_feature, grouped_idx.reshape((batch_size, -1)))
            # grouped_feature.shape = (batch_size, num_new_point_cloud, num_samples, feature_dim)
            grouped_feature = F.reshape(
                grouped_feature, (batch_size, num_new_point_cloud, num_samples, feature_dim))
            grouped_feature = F.concatenate(
                grouped_feature, grouped_point_cloud, axis=-1)
        else:
            grouped_feature = grouped_point_cloud

        # grouped_feature.shape = (batch_size, num_new_point_cloud, num_samples, feature_dim)
        # to (batch_size, feature_dim, num_samples, num_new_point_cloud)
        grouped_feature = F.transpose(grouped_feature, (0, 3, 2, 1))
        conv_h = grouped_feature

        for j, conv_channel in enumerate(conv_channels[i]):
            with nn.parameter_scope(f"abstraction{i}/conv{j}"):
                conv_h = PF.convolution(
                    conv_h,
                    conv_channel,
                    (1, 1),
                    stride=(1, 1),
                    with_bias=False,
                )
                conv_h = PF.batch_normalization(conv_h, batch_stat=train)
                conv_h = F.relu(conv_h)

        # take max axis of num_samples
        # to (batch_size, feature_dim, num_new_point_cloud)
        new_point_feature = F.max(conv_h, axis=2)
        new_point_features.append(new_point_feature)

    new_point_features = F.concatenate(*new_point_features, axis=1)
    # to (batch_size, len(radius_list)*num_new_point_cloud, feature_dim)
    new_point_features = F.transpose(new_point_features, (0, 2, 1))
    return new_point_cloud, new_point_features


def set_abstraction(
    point_cloud: nn.Variable,
    point_feature: Optional[nn.Variable],
    num_new_point_cloud: Optional[int],
    radius: Optional[float],
    num_samples: int,
    conv_channels: Iterable[int],
    group_all: bool,
    train: bool,
) -> Tuple[nn.Variable, nn.Variable]:
    """set abstraction layer

    Args:
        point_cloud (nn.Variable): shape (batch_size, num_points, dim)
        point_feature (Union[nn.Variable, None]): shape (batch_size, num_new_point_cloud, feature_dim)
        num_new_point_cloud (int): number of new point cloud
        radius (float): radius
        num_samples (int): number of samples
        conv_channels (Iterable[int]): list of convolutional channel
        group_all (bool): grouping the point
        train (bool): training flag

    Returns:
        Tuple[nn.Variable, nn.Variable]: new_point_cloud, shape (batch_size, num_new_point_cloud, 3),
            new_point_features, shape (batch_size, num_new_point_cloud, feature_dim)
    """
    if group_all:
        new_point_cloud, new_point_feature = sample_and_group_all(
            point_cloud, point_feature)
    else:
        new_point_cloud, new_point_feature = sample_and_group(
            num_new_point_cloud, radius, num_samples, point_cloud, point_feature
        )

    # new_point_feature.shape = (batch_size, num_new_point_cloud, num_samples, feature_dim)
    # to (batch_size, feature_dim, num_samples, num_new_point_cloud)
    new_point_feature = F.transpose(new_point_feature, (0, 3, 2, 1))
    conv_h = new_point_feature

    for i, conv_channel in enumerate(conv_channels):
        with nn.parameter_scope(f"conv{i}"):
            conv_h = PF.convolution(
                conv_h,
                conv_channel,
                (1, 1),
                stride=(1, 1),
                with_bias=False,
            )
            conv_h = PF.batch_normalization(conv_h, batch_stat=train)
            conv_h = F.relu(conv_h)

    # take max axis of num_samples
    # to (batch_size, feature_dim, num_new_point_cloud)
    new_point_feature = F.max(conv_h, axis=2)

    # to (batch_size, num_new_point_cloud, feature_dim)
    new_point_feature = F.transpose(new_point_feature, (0, 2, 1))
    return new_point_cloud, new_point_feature


def feature_propagation(
    point_cloud_base: nn.Variable,
    point_feature_base: nn.Variable,
    point_cloud_next: nn.Variable,
    point_feature_next: nn.Variable,
    conv_channels: Iterable[int],
    train: bool,
    k: int = 3,
) -> nn.Variable:
    """feature propagation layer

    Args:
        point_cloud_base (nn.Variable): shape (batch_size, num_base_points, dim)
        point_feature_base (nn.Variable): shape (batch_size, num_base_points, feature_dim)
        point_cloud_next (nn.Variable): shape (batch_size, num_next_points, dim)
        point_feature_next (nn.Variable): shape (batch_size, num_next_points, feature_dim)
        train (bool): training flag
        k (int): knn parameter. Defaults to 3.

    Returns:
        nn.Variable: new_features, shape (batch_size, num_new_features, feature_dim)
    """
    batch_size, num_base_points, base_dim = point_cloud_base.shape
    _, num_next_points, _ = point_cloud_next.shape

    if num_next_points == 1:
        interpolated_point_feature = F.tile(
            point_feature_next, (1, num_base_points, 1))
    else:
        # distance.shape(batch_size, num_base_points, num_next_points)
        distance = compute_square_distance(point_cloud_base, point_cloud_next)
        sorted_distance, sorted_distance_idx = F.sort(
            distance, axis=2, with_index=True)
        sorted_distance = sorted_distance[:, :, :k]
        sorted_distance_idx = sorted_distance_idx[:, :, :k]
        batch_size, num_base_points, _ = sorted_distance_idx.shape

        inv_distance_batch = 1.0 / (sorted_distance + 1e-8)
        norm_batch = F.sum(inv_distance_batch, axis=2, keepdims=True)
        weight = inv_distance_batch / norm_batch
        # weight.shape(batch_size, num_base_points, 3, 1)
        weight = F.reshape(weight, (batch_size, num_base_points, 3, 1))

        sorted_distance_idx = F.reshape(
            sorted_distance_idx, (batch_size, num_base_points * 3))
        # selected_point_feature_next.shape(batch_size, num_base_points*3, point_next_feature_dim)
        selected_point_feature_next = index_batch_variable(
            point_feature_next, sorted_distance_idx)
        # selected_point_feature_next.shape(batch_size, num_base_points, 3, point_next_feature_dim)
        selected_point_feature_next = F.reshape(
            selected_point_feature_next, (batch_size, num_base_points, 3, -1))

        # sum according to num_next_points axis
        # selected_point_feature_next.shape(batch_size, num_base_points, 3, point_next_feature_dim)
        interpolated_point_feature = F.sum(
            selected_point_feature_next * weight, axis=2)

    if point_feature_base is not None:
        new_point_feature = F.concatenate(
            point_feature_base, interpolated_point_feature, axis=2)
    else:
        new_point_feature = interpolated_point_feature

    # new_point_feature.shape = (batch_size, num_base_points, feature_dim+point_next_feature_dim)
    # to (batch_size, feature_dim+point_next_feature_dim, num_base_points)
    new_point_feature = F.transpose(new_point_feature, (0, 2, 1))
    # add dimensions, to (batch_size, feature_dim, 1, num_base_points)
    new_point_feature = F.reshape(
        new_point_feature, (batch_size, -1, 1, num_base_points))

    conv_h = new_point_feature
    for i, conv_channel in enumerate(conv_channels):
        with nn.parameter_scope(f"conv{i}"):
            conv_h = PF.convolution(
                conv_h,
                conv_channel,
                (1, 1),
                stride=(1, 1),
            )
            conv_h = PF.batch_normalization(conv_h, batch_stat=train)
            conv_h = F.relu(conv_h)

    new_point_feature = F.reshape(conv_h, (batch_size, -1, num_base_points))
    new_point_feature = F.transpose(new_point_feature, (0, 2, 1))
    return new_point_feature


# utils/neu/post_processing.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


###############################################################################
# Code from
# https://github.com/ycszen/pytorch-seg/blob/master/transform.py
# Modified so it complies with the Citscape label map colors
# and is able to be used in nnabla
###############################################################################

import numpy as np


def uint82bin(n, count=8):
    """returns the binary of integer n, count refers to amount of bits"""
    return ''.join([str((n >> y) & 1) for y in range(count - 1, -1, -1)])


def labelcolormap(N):
    if N == 35:  # cityscape
        cmap = np.array([(0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (111, 74, 0), (81, 0, 81),
                         (128, 64, 128), (244, 35, 232), (250,
                                                          170, 160), (230, 150, 140), (70, 70, 70),
                         (102, 102, 156), (190, 153, 153),
                         (180, 165, 180), (150, 100, 100), (150,
                                                            120, 90), (153, 153, 153), (153, 153, 153),
                         (250, 170, 30), (220, 220, 0),
                         (107, 142, 35), (152, 251, 152), (70, 130,
                                                           180), (220, 20, 60), (255, 0, 0), (0, 0, 142),
                         (0, 0, 70),
                         (0, 60, 100), (0, 0, 90), (0, 0, 110), (0, 80, 100), (0, 0, 230), (119, 11, 32), (0, 0, 142)],
                        dtype=np.uint8)
    else:
        cmap = np.zeros((N, 3), dtype=np.uint8)
        for i in range(N):
            r, g, b = 0, 0, 0
            id = i
            for j in range(7):
                str_id = uint82bin(id)
                r = r ^ (np.uint8(str_id[-1]) << (7 - j))
                g = g ^ (np.uint8(str_id[-2]) << (7 - j))
                b = b ^ (np.uint8(str_id[-3]) << (7 - j))
                id = id >> 3
            cmap[i, 0] = r
            cmap[i, 1] = g
            cmap[i, 2] = b
    return cmap


class Colorize(object):
    def __init__(self, n=35):
        self.cmap = labelcolormap(n)

    def __call__(self, label_image, channel_first=False):
        import nnabla as nn

        if isinstance(label_image, np.ndarray):
            im = label_image
        elif isinstance(label_image, nn.NdArray):
            im = label_image.get_data("r")
        elif isinstance(label_image, nn.Variable):
            im = label_image.data.get_data("r")
        else:
            raise ValueError("Error to colorize label_ids")

        color_image = np.zeros(shape=im.shape + (3,))

        for label in range(0, len(self.cmap)):
            mask = label_image == label
            color_image[mask] = self.cmap[label]

        if channel_first:
            color_image = color_image.transpose((0, 3, 1, 2))

        return color_image


# utils/neu/reporter.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import

import os
import sys
from collections import OrderedDict, defaultdict

import nnabla as nn
import nnabla.functions as F
import nnabla.monitor as M
import numpy as np
from nnabla.monitor import tile_images
from nnabla.utils.image_utils import imsave

from .html_creator import HtmlCreator


def get_value(val, dtype=float, reduction=True):
    """
    get float value from nn.NdArray / nn.Variable / np.ndarray / float
    """

    # get NdArray from Variable
    if isinstance(val, nn.Variable):
        val = val.data

    # get value as float
    if isinstance(val, nn.NdArray):
        assert not val.clear_called

        # take average if val has more than one element
        if reduction and val.size > 1:
            with nn.auto_forward(), nn.no_grad():
                val = F.mean(val)

        v = val.get_data("r")
    elif isinstance(val, np.ndarray):
        if reduction and val.size > 1:
            val = np.mean(val)

        v = val
    else:
        assert isinstance(val, (int, float, np.generic))
        v = val

    return dtype(v)


def get_tiled_image(img, channel_last=False):
    assert len(img.shape) == 4
    assert isinstance(img, np.ndarray)

    if channel_last:
        # nnabla.monitor.tile_images requests (B, C, H, W)
        # (B, H, W, C) -> (B, C, H, W)
        img = img.transpose(0, 3, 1, 2)

    B, C, H, W = img.shape

    # create tiled image. channel last image will be returned.
    tiled_image = tile_images(img)
    _, _, Ct = tiled_image.shape
    assert C == Ct

    return tiled_image


def save_tiled_image(img, path, channel_last=False):
    """
    Save given batched images as tiled image.
    The first axis will be handled as batch.

    Args:
        img (np.ndarray):
            Images to save. The shape should be (B, C, H, W) or (B, H, W, C) depending on `channel_last`. dtype must be np.uint8.
        path (str):
            Path to save.
        channel_last (bool):
            If True, the last axis (=3) will be handled as channel.
    """
    tiled_image = get_tiled_image(img, channel_last=channel_last)

    assert tiled_image.dtype == np.uint8

    # create directory if needed
    os.makedirs(os.path.dirname(path), exist_ok=True)

    # save
    imsave(path, tiled_image)


class MonitorWrapper(object):
    """
    Usage:
        monitor = MonitorWrapper("path/to/save", interval=10, save_time=True)

        # Currently supports three types below
        # v1: nn.Variable()
        # v2: nn.NdArray()
        # v3: np.ndarray() or float
        vars = {"name1": v1, "name2": v2, "name3", v3}

        for epoch in range(max_epoch):
            monitor(vars, epoch)
    """

    def __init__(self, save_path, interval=1, save_time=True, silent=False):
        self.monitor = M.Monitor(save_path)
        self.interval = interval
        self.silent = silent

        self.series_monitors = {}

        if save_time:
            self.monitor_time = M.MonitorTimeElapsed(
                "Epoch time", self.monitor, interval=interval, verbose=not self.silent)

    def set_series(self, name):
        """
        Set a new series name.
        """
        assert isinstance(name, str), f"name must be string but {type(name)}."
        self.series_monitors[name] = M.MonitorSeries(name, self.monitor,
                                                     interval=self.interval,
                                                     verbose=not self.silent)

    def __call__(self, name, series_val, epoch):
        assert isinstance(name, str), "name must be string."

        if name not in self.series_monitors:
            self.set_series(name)

        self.series_monitors[name].add(epoch, get_value(series_val))

        if self.monitor_time is not None:
            self.monitor_time.add(epoch)


class Reporter(object):
    def __init__(self, comm, losses, save_path=None, nimage_per_epoch=1,
                 show_interval=20, show_keys=None):
        # losses: {"loss_name": loss_Variable, ...} or list of tuple(key, value)

        self.batch_cnt = 0
        self.piter = None
        self.comm = comm
        self.save_path = save_path
        self.nimage_per_epoch = nimage_per_epoch
        self.show_interval = show_interval

        self.losses = OrderedDict(losses)  # fix loss order
        self.epoch_losses = {k: 0. for k in losses.keys()}
        self.buff = {k: nn.NdArray() for k in losses.keys()}
        self.show_keys = list(
            losses.keys()) if show_keys is None else show_keys

        is_master = comm.rank == 0
        self.monitor = MonitorWrapper(save_path) if (
            save_path is not None and is_master) else None

        self._reset_buffer()

    def set_losses(self, losses, update=True):
        # {"loss_name": loss_Variable, ...}
        assert isinstance(losses, dict)

        for key, value in losses.items():
            if key in self.losses and not update:
                continue

            self.losses[key] = value

    def _reset_buffer(self):
        # reset buff
        for loss_name, loss in self.losses.items():
            if loss is None:
                continue
            self.buff[loss_name] = nn.NdArray()
            self.buff[loss_name].zero()

        self.flushed = True

    def _save_image(self, file_name, image):
        if isinstance(image, nn.Variable):
            img = image.data.get_data("r")
        elif isinstance(image, nn.NdArray):
            img = image.get_data("r")
        else:
            assert isinstance(image, np.ndarray)
            img = image

        dir_path = os.path.join(self.save_path, "html", "imgs")

        if not os.path.exists(dir_path):
            os.makedirs(dir_path)

        save_path = os.path.join(dir_path, file_name)

        img = (img - img.min()) / (img.max() - img.min())
        imsave(save_path, img)

    def _render_html(self, epoch, image_names):
        # currently dominate only supports to add dom attrs from top to bottom.
        # So we have to re-create html from scratch at each epoch in order to place results in reverse order.
        self.html = HtmlCreator(os.path.join(
            self.save_path, "html"), redirect_interval=300)

        for e in range(epoch, -1, -1):
            self.html.add_text("epoch {}".format(e))
            image_files = []
            for i in range(self.nimage_per_epoch):
                for image_name in image_names:
                    image_files.append(
                        "_".join([image_name, str(e), str(i)]) + ".png")

            self.html.add_images(image_files,
                                 [x.split("_")[0] for x in image_files])

        self.html.save()

    def start(self, progress_iter):
        self.batch_cnt = 0
        for k in self.losses.keys():
            self.epoch_losses[k] = 0.

        self.piter = progress_iter

    def _flush_losses(self):
        if self.flushed:
            return
        self.flushed = True

        desc = "[report]"

        for loss_name, loss in self.losses.items():
            if loss is None:
                continue
            self.epoch_losses[loss_name] += self.buff[loss_name].get_data("r")

            if loss_name in self.show_keys:
                desc += " {}: {:.3f}".format(loss_name,
                                             self.epoch_losses[loss_name] / self.batch_cnt)

        # show current values
        self.piter.set_description(desc)

        self._reset_buffer()

    def __call__(self):
        # update state
        self.batch_cnt += 1
        self.flushed = False

        for loss_name, loss in self.losses.items():
            if loss is None:
                continue
            self.buff[loss_name] += loss.data

        if self.batch_cnt % self.show_interval == 0:
            self._flush_losses()

    def step(self, iteration, images=None):
        # images = {"image_name": image, ...}

        self._flush_losses()

        comm_values = {k: nn.NdArray.from_numpy_array(np.asarray(x / self.batch_cnt, dtype=np.float32))
                       for k, x in self.epoch_losses.items()}

        self.comm.all_reduce(list(comm_values.values()),
                             division=True, inplace=True)

        if self.comm.rank == 0:
            if self.monitor is not None:
                for name, val in comm_values.items():
                    self.monitor(name, val, iteration)

            if images is not None:
                # write images to files.
                images_as_tuple = tuple(images.items())

                for image_name, image in images_as_tuple:
                    assert len(image) >= self.nimage_per_epoch

                    for i in range(self.nimage_per_epoch):
                        file_name = "_".join(
                            [image_name, str(iteration), str(i)]) + ".png"
                        self._save_image(file_name, image[i])

                self._render_html(iteration, tuple(images.keys()))


class AverageLogger(object):
    def __init__(self):
        self.avg_val = 0.
        self.cnt = 0

    def reset(self):
        self.avg_val = 0.
        self.cnt = 0.

    def update(self, val):
        val = get_value(val, dtype=float)
        oldcnt = self.cnt
        self.cnt += 1

        # Should always be true, but just in case.
        assert self.cnt > 0, "zero division"

        # take moving average
        self.avg_val = self.avg_val * oldcnt / self.cnt + val / self.cnt

    @property
    def val(self):
        return self.avg_val


class KVReporter(object):
    """
    Usage:

        import KVReporter

        comm = init_nnabla(ext_name="cudnn", device_id=args.device_id, type_config="float")
        reporter = KVReporter(comm, save_path="path/to/logdir", show_interval=20, force_persistent=True)

        x = nn.Variable(...)
        h1 = F.affine(x, c1, name="aff1")
        h2 = F.affine(h1, c2, name="aff2")
        loss = F.mean(F.squared_error(h2, t))

        solver = S.Adam()
        solver.set_paramters(nn.get_parameters())

        # show loss
        for i in range(max_iter):
            loss.forward()
            loss.backward()
            solver.zero_grad()
            solver.update()

            # KVReporter can handle nn.Variable and nn.NdArray as well as np.ndarry.
            # Using kv_mean(name, val), you can calculate moving average.
            reporter.kv_mean("loss", loss)

            # If you don't need to take moving average for the value, use kv(name, val) instead.
            reporter.kv("iterations", i)

            # get all traced values.
            # If sync=True, all values will be synced across devices via `comm`.
            # If reset=True, KVReporter resets the current average for all values as zero.
            reporter.dump(file=sys.stdout, sync=True, reset=True)

            # save values through nnabla.Monitor
            reporter.flush_monitor(i)

    """

    def __init__(self, comm=None, save_path=None,
                 monitor_silent=True, skip_kv_to_monitor=True):
        """
        Arguments:
            comm: Communicator or CommunicatorWrapper
            save_path: str. Where to save.
            show_interval: int. Interval to report.
        """
        self.batch_cnt = 0
        self.comm = comm
        self.is_synced = False

        self.name2logger = defaultdict(AverageLogger)
        self.name2val = defaultdict(float)

        self._monitor = None
        self.is_master = comm is None or comm.rank == 0
        if self.is_master:
            self._monitor = MonitorWrapper(save_path,
                                           interval=1,
                                           silent=monitor_silent) if save_path else None

        self.monitor_cnt = 0
        self.skip_kv_to_monitor = skip_kv_to_monitor

    def get_val(self, name):
        return self.name2logger[name].val

    def reset(self, names=None):
        # Do not reset monitor_cnt

        if isinstance(names, str):
            names = [names]

        for key, logger in self.name2logger.items():
            if (names is not None) and (key not in names):
                continue

            logger.reset()

        self.is_synced = True  # No need to sync initial values

    def sync_all(self, reset=True, sync=True):
        if self.is_synced:
            return

        for name, logger in sorted(self.name2logger.items(), key=lambda x: x):
            if sync:
                # sync across all devices
                synced_val = nn.NdArray.from_numpy_array(
                    np.asarray(get_value(logger.val)))
                synced_cnt = nn.NdArray.from_numpy_array(
                    np.asarray(get_value(logger.cnt)))
                if self.comm is not None:
                    try:
                        synced_val *= synced_cnt
                        self.comm.all_reduce(
                            [synced_val, synced_cnt], division=False, inplace=True)
                        synced_val /= synced_cnt
                    except:
                        raise ValueError(
                            f"Sync error. rank: {self.comm.rank}, key: {name}")

                # update
                self.name2val[name] = get_value(synced_val.get_data("r"))
            else:
                self.name2val[name] = get_value(logger.val)

        self.is_synced = True

        if reset:
            self.reset()

    def set_key(self, key):
        """
        set key before to prevent synchronization error.
        """

        self.name2logger[key] = self.name2logger[key]

    def desc(self, reset=True, sync=True):
        self.sync_all(reset=reset, sync=sync)

        desc = "[report]"

        for name, val in sorted(self.name2val.items(), key=lambda x: x):
            desc += " {}: {:-8.3g}".format(name, val)

        return desc

    def dump(self, file=sys.stdout, reset=True, sync=True):
        self.sync_all(reset=reset, sync=sync)

        key2str = OrderedDict()
        max_key_width = 0
        max_val_width = 0
        for name, val in sorted(self.name2val.items(), key=lambda x: x):
            key2str[name] = "{:-8.3g}".format(val)
            max_key_width = max(max_key_width, len(name))
            max_val_width = max(max_val_width, len(key2str[name]))

        if len(key2str) == 0:
            return

        line = "=" * (max_key_width + max_val_width + 5)

        out = [line]
        for name, val in key2str.items():
            key_pad = " " * (max_key_width - len(name))
            out.append("{:<}{} : {:<}".format(name, key_pad, val))
        out.append(line)

        if file is not None and hasattr(file, "write"):
            file.write("\n".join(out) + "\n")
            file.flush()

    def kv(self, name, val):
        """
        For additional information. This is not dumped to monitor. 
        """
        self.name2val[name] = val

    def kv_mean(self, name, val):
        """
        update average value by `val` for the key `name`.
        """
        # take average
        logger = self.name2logger[name]
        logger.update(get_value(val))

        self.is_synced = False

    def flush_monitor(self, iter, names=None):
        """
        write kv_mean values to monitor
        """
        self.sync_all(reset=False)

        if not self.is_master or self._monitor is None:
            return

        for name, val in self.name2val.items():
            #  skip values set by kv()
            if self.skip_kv_to_monitor and name not in self.name2logger:
                continue

            if (names is not None) and (name not in names):
                continue

            self._monitor(name, val, iter)


# utils/neu/safe_extract.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import tarfile


def safe_extract(tar, path=".", members=None, *, numeric_owner=False):
    def is_within_directory(directory, target):
        abs_directory = os.path.abspath(directory)
        abs_target = os.path.abspath(target)

        prefix = os.path.commonprefix([abs_directory, abs_target])
        return prefix == abs_directory

    for member in tar.getmembers():
        member_path = os.path.join(path, member.name)
        if not is_within_directory(path, member_path):
            raise Exception("Attempted Path Traversal in Tar File")

    tar.extractall(path, members, numeric_owner=numeric_owner)


# utils/neu/save_args.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
from .yaml_wrapper import write_yaml


def save_args(args, config=None):
    from nnabla import logger
    if not os.path.exists(args.monitor_path):
        os.makedirs(args.monitor_path)

    path = os.path.join(args.monitor_path, "Arguments.txt")
    logger.info("Arguments are saved to {}.".format(path))
    with open(path, "w") as fp:
        for k, v in sorted(vars(args).items()):
            logger.info("{}={}".format(k, v))
            fp.write("{}={}\n".format(k, v))

            if config is not None and isinstance(v, str) and v.endswith(".yaml"):
                if os.path.isabs(v):
                    path = os.path.join(args.monitor_path, os.path.basename(v))
                else:
                    path = os.path.join(args.monitor_path, v)
                write_yaml(path, config)


# utils/neu/save_nnp.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


def save_nnp(input, output, batchsize):
    runtime_contents = {
        'networks': [
            {'name': 'Validation',
             'batch_size': batchsize,
             'outputs': output,
             'names': input}],
        'executors': [
            {'name': 'Runtime',
             'network': 'Validation',
             'data': [k for k, _ in input.items()],
             'output':[k for k, _ in output.items()]}]}
    return runtime_contents


# utils/neu/solvers.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from contextlib import contextmanager
from collections import defaultdict
import numpy as np

import nnabla as nn
import nnabla.solvers as S
import nnabla.functions as F
from nnabla.parameter import get_parameter_or_create


@contextmanager
def context_scope(type_config):
    assert type_config in ["float", "half"], \
        "type_config must be either of ['float', 'half']"

    from nnabla.ext_utils import get_extension_context
    current_ctx = nn.get_current_context()
    ctx = get_extension_context(ext_name=current_ctx.backend[0].split(":")[0],
                                device_id=current_ctx.device_id,
                                type_config=type_config)

    nn.set_default_context(ctx)

    yield ctx

    nn.set_default_context(current_ctx)


def force_float(func):
    def wrapped_func(*args, **kwargs):
        with context_scope("float"):
            return func(*args, **kwargs)

    return wrapped_func


def create_ema_parameter(base_name, base_param):
    # create ema parameter whose name is "ema/{base_name}".
    with nn.parameter_scope("ema"):
        ema_param = get_parameter_or_create(base_name,
                                            shape=base_param.shape,
                                            need_grad=False)

    return ema_param


@force_float
def ema_update(ema_param: nn.Variable, train_param: nn.Variable, decay=0.9999):
    assert isinstance(ema_param, nn.Variable)
    assert isinstance(train_param, nn.Variable)

    # apply ema update by inplace operation
    # p_ema = decay * p_ema + (1 - decay) * p_train <- needs copy to assign the result
    #       = p_ema - (1 - decay) * (p_ema - p_train) <- performs the same thing by isub operation

    sub = (1. - decay) * (ema_param.data - train_param.data)
    ema_param.data -= sub


class PackedParameterSolver(object):

    def __init__(self,
                 solver: S.Solver,
                 use_ema: bool):
        self.solver = solver

        # keep original params and states for saving
        self.orig_params = solver.get_parameters()
        self.orig_states = solver.get_states()

        # big params and states
        # Basically we assume params.data is float32, but prams.grad can be either float32 or float16.
        # Therefore, we separately pack params depending on the dtype of grad.
        self._num_pp = 2  # for easy implementation
        self.packed_params = None
        self.packed_states = None
        self.is_packed = False

        # cudnn requires 128 bit and recommends 1024 bit alignment
        # Align size by a multiple of 64 is enough since at least data has a bit size of multiple of 16 bit (float16).
        self.alignment = 64

        # for ema parameters
        self.use_ema = use_ema
        self.packed_ema_params = None

    @staticmethod
    def _get_param_type_based_on_dtype(x: nn.Variable):
        assert x.grad.dtype in (np.float32, np.float16), \
            f"datatype {x.grad.dtype} is not supported."
        return int(x.grad.dtype == np.float32)

    def _align_size(self, size):
        return (size + self.alignment - 1) // self.alignment * self.alignment

    def _create_packed_param_state(self):
        orig_params = self.solver.get_parameters()
        orig_states = self.solver.get_states()

        # 1. compute total size for both parameter and state
        p_total_sizes = [0 for _ in range(self._num_pp)]
        s_total_sizes = [defaultdict(int) for _ in range(self._num_pp)]
        solver_step = None
        for p_key, p in orig_params.items():
            idx = self._get_param_type_based_on_dtype(p)

            # For packed param
            p_total_sizes[idx] += self._align_size(p.size)

            # For packed states
            state = orig_states[p_key]
            if solver_step is None:
                solver_step = state.t
            else:
                # All solver steps must be identical
                assert solver_step == state.t

            # Each param could have more than one corresponding states.
            for s_name, pstate in state.pstate.items():
                s_total_sizes[idx][s_name] += self._align_size(pstate.size)

        # 2. create large variable for packing
        # parameter
        self.packed_params = [nn.Variable(shape=(x, ))
                              for x in p_total_sizes]

        if self.use_ema:
            self.packed_ema_params = [nn.Variable(shape=(x, ))
                                      for x in p_total_sizes]

        # state
        packed_pstates = [dict() for _ in range(self._num_pp)]
        for i, sname_to_size_dict in enumerate(s_total_sizes):
            for sname, size in sname_to_size_dict.items():
                packed_pstates[i][sname] = nn.Variable(shape=(size, ))

        # 3. assign a segment to each param/state.
        p_cur_sizes = [0 for _ in range(self._num_pp)]
        s_cur_sizes = [defaultdict(int) for _ in range(self._num_pp)]

        for p_key, p in orig_params.items():
            # get packed parameter id (0: grad is float16, 1: grad is float32)
            idx = self._get_param_type_based_on_dtype(p)

            # 3.1 replace parameter
            # select a packed param for the current one
            packed_param = self.packed_params[idx]
            p_cur_size = p_cur_sizes[idx]

            # get narrowed array for the current parameter
            data_narrowed = packed_param.data.narrow(
                0, p_cur_size, p.size).view(p.shape)
            grad_narrowed = packed_param.grad.narrow(
                0, p_cur_size, p.size).view(p.shape)

            # copy data and grad from the original parameter
            if p.data.dtype != "float32":
                raise ValueError(f"p.data.dtype != float32 for {p_key}.")
            # set use_current_context as False to keep the dtype of the original parameter.
            data_narrowed.copy_from(p.data, use_current_context=False)
            grad_narrowed.copy_from(p.grad, use_current_context=False)

            # replace ndarray
            p.data = data_narrowed
            p.grad = grad_narrowed

            # ema
            if self.use_ema:
                # create (or get) ema param
                p_ema = create_ema_parameter(p_key, p)

                # replace ndarray
                packed_ema_param = self.packed_ema_params[idx]
                ema_data_narrowed = packed_ema_param.data.narrow(
                    0, p_cur_size, p.size).view(p.shape)

                ema_data_narrowed.copy_from(
                    p_ema.data, use_current_context=False)
                p_ema.data = ema_data_narrowed

            # update total size
            p_cur_sizes[idx] += self._align_size(p.size)

            # 3.2 replace state
            # get states for the current param
            state = orig_states[p_key]

            # replace ndarray for each state
            for s_name, pstate in state.pstate.items():
                s_cur_size = s_cur_sizes[idx][s_name]
                pstate_narrowed = packed_pstates[idx][s_name].data.narrow(0,
                                                                          s_cur_size,
                                                                          pstate.size)
                pstate_narrowed = pstate_narrowed.view(pstate.shape)

                # (20221103) Always use float context since there is no way to check whether pstate has array or not.
                with context_scope("float"):
                    pstate_narrowed.copy_from(pstate.data)
                pstate.data = pstate_narrowed

                s_cur_sizes[idx][s_name] += self._align_size(pstate.size)

        # create SolverState object from packed states
        self.packed_states = []
        for idx in range(self._num_pp):
            solver_state = S.SolverState()
            solver_state.t = solver_step
            solver_state.pstate = {}

            for s_name, pstate in packed_pstates[idx].items():
                solver_state.pstate[s_name] = pstate

            self.packed_states.append(solver_state)

    def _packing(self):
        if self.is_packed:
            return

        # create packed param and state
        self._create_packed_param_state()

        # keep original parameters and states for saving
        self.orig_params = self.solver.get_parameters()
        self.orig_states = self.solver.get_states()

        # reset solver to have new packed parameters
        self.solver.clear_parameters()
        packed_params_dict = {}
        packed_states_dict = {}
        for idx in range(self._num_pp):
            p = self.packed_params[idx]
            s = self.packed_states[idx]

            # skip if empty, otherwise solver method raises due to empty array.
            if p.size == 0:
                assert len(s.pstate) == 0
                continue

            packed_params_dict[f"PackedParam{idx}"] = p
            packed_states_dict[f"PackedParam{idx}"] = s

        self.solver.set_parameters(packed_params_dict)
        self.solver.set_states(packed_states_dict)

        self.is_packed = True

    def updata_ema_params(self, decay=0.9999):
        if not self.use_ema:
            return

        if not self.is_packed:
            self._packing()

        for p_ema, p_train in zip(self.packed_ema_params, self.packed_params):
            if p_ema.size == 0:
                assert p_train.size == 0
                continue

            # make sure to keep data
            p_ema.persistent = True
            p_train.persistent = True

            ema_update(p_ema, p_train, decay=decay)

    def set_parameters(self, *args, **kwargs):
        raise NotImplementedError(
            "All parameters must be set to solver object before passing solver to PackedParameterSolver.")

    def set_learning_rate(self, lr):
        self.solver.set_learning_rate(lr)

    def zero_grad(self):
        self.solver.zero_grad()

    def check_inf_or_nan_grad(self):
        self._packing()
        return self.solver.check_inf_or_nan_grad()

    def update(self, *args, **kwargs):
        self._packing()
        self.solver.update(*args, **kwargs)

    def scale_grad(self, *args, **kwargs):
        self._packing()
        self.solver.scale_grad(*args, **kwargs)

    def clip_by_norm(self, *args, **kwargs):
        self._packing()
        self.solver.clip_by_norm(*args, **kwargs)

    # Functions which need to convert the packed param to the original params
    def save_states(self, path):
        # keep packed states
        packed_states = self.solver.get_states()
        packed_params = self.solver.get_parameters()

        # reset original states to solver
        self.solver.set_parameters(self.orig_params, reset=True)

        # set current timestep for all state
        packed_t = None
        for idx in range(self._num_pp):
            s = self.packed_states[idx]
            if packed_t is None:
                packed_t = s.t
            else:
                assert packed_t == s.t

        for k, s in self.orig_states.items():
            s.t = packed_t

        self.solver.set_states(self.orig_states)
        self.solver.save_states(path)

        # revert packed params
        self.solver.set_parameters(packed_params, reset=True)
        self.solver.set_states(packed_states)

    def get_parameters(self):
        return self.orig_params

    def get_states(self):
        # set current timestep for all state
        packed_t = None
        for idx in range(self._num_pp):
            s = self.packed_states[idx]
            if packed_t is None:
                packed_t = s.t
            else:
                assert packed_t == s.t

        for k, s in self.orig_states.items():
            if s.t == packed_t:
                # no need to restore timestep
                break
            s.t = packed_t

        return self.orig_states

    def load_states(self, path):
        if not self.is_packed:
            # solver must have origianl parameters and states.
            # call solver.load_stats directely
            self.solver.load_states(path)
            return

        # otherwise, replace states and re-packing again

        # reset original states to solver
        self.solver.set_parameters(self.orig_params, reset=True)
        self.solver.load_states(path)

        # re-pack parameters
        self.is_packed = False
        self._packing()


# utils/neu/variable_utils.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import nnabla as nn


def get_params_startswith(str):
    return {k: v for k, v in nn.get_parameters().items() if k.startswith(str)}


def set_persistent_all(*variables):
    for var in variables:
        if var is None:
            continue

        if not isinstance(var, nn.Variable):
            raise ValueError("all variables must be nn.Variable")

        var.persistent = True


def set_need_grad_all(*variables, need_grad):
    assert isinstance(need_grad, bool)

    for var in variables:
        if var is None:
            continue

        if not isinstance(var, nn.Variable):
            raise ValueError("all variables must be nn.Variable")

        var.need_grad = need_grad


def get_unlinked_all(*variables):
    ret = []
    for var in variables:
        if var is None:
            ret.append(None)
            continue

        if not isinstance(var, nn.Variable):
            raise ValueError("all variables must be nn.Variable")

        ret.append(var.get_unlinked_variable())

    return ret


def zero_grads_all(*variables):
    for var in variables:
        if var is None:
            continue

        if not isinstance(var, nn.Variable):
            raise ValueError("all variables must be nn.Variable")

        var.grad.zero()


def fill_all(*variables, value=0):
    for var in variables:
        if var is None:
            continue

        if not isinstance(var, nn.Variable):
            raise ValueError("all variables must be nn.Variable")

        var.data.fill(value)


# utils/neu/yaml_wrapper.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import os
import yaml
import re

from .misc import AttrDict, makedirs


# Resolve scientific notation: https://github.com/yaml/pyyaml/pull/174/files
yaml.resolver.Resolver.add_implicit_resolver(
    u'tag:yaml.org,2002:float',
    re.compile(r'''^(?:[-+]?(?:[0-9][0-9_]*)\.[0-9_]*(?:[eE][-+]?[0-9]+)?
                    |[-+]?(?:[0-9][0-9_]*)(?:[eE][-+]?[0-9]+)
                    |\.[0-9_]+(?:[eE][-+]?[0-9]+)?
                    |[-+]?[0-9][0-9_]*(?::[0-5]?[0-9])+\.[0-9_]*
                    |[-+]?\.(?:inf|Inf|INF)
                    |\.(?:nan|NaN|NAN))$''', re.X),
    list('-+0123456789.'))


# Convert AttrDict tag to standard dict
def _represent_attrdict(dumper, instance):
    return dumper.represent_mapping('tag:yaml.org,2002:map', instance.items())


yaml.add_representer(AttrDict, _represent_attrdict)


def read_yaml(filepath):
    with open(filepath, "r") as f:
        data = yaml.load(f, Loader=yaml.FullLoader)

    return AttrDict(data)


def write_yaml(filepath, obj):
    dirname = os.path.dirname(filepath)

    makedirs(dirname)

    with open(filepath, 'w') as f:
        yaml.dump(obj, f, default_flow_style=False)


__all__ = [read_yaml, write_yaml]


# video-colorization/deep-exemplar-based-video-colorization/args.py


# Copyright 2021 Sony Group Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os
from utils import *


def get_config():
    """
    Get command line arguments.
    Arguments set the default values of command line arguments
    """

    parser = argparse.ArgumentParser()
    parser.add_argument('--cfg', default="./config.yaml")
    args, _ = parser.parse_known_args()
    conf = read_yaml(args.cfg)

    parser.add_argument('--device-id', '-d',
                        type=int,
                        default=0,
                        help='device id (default: 0)')
    parser.add_argument(
        '--context',
        '-c',
        type=str,
        default='cudnn',
        help="Extension path: cpu or cudnn.")
    parser.add_argument(
        "--frame_propagation",
        default=False,
        type=bool,
        help="color propogation mode")
    parser.add_argument(
        "--image_size",
        type=int,
        default=[
            216 * 2,
            384 * 2],
        nargs="+",
        help="the image size, eg. 432 768")
    parser.add_argument(
        "--input_path",
        type=str,
        default="./images/input/v32",
        help="path of input clips")
    parser.add_argument(
        "--ref_path",
        type=str,
        default="./images/ref/v32",
        help="path of reference images")
    parser.add_argument(
        "--output_path",
        type=str,
        default="./images/output",
        help="path of output clips")
    parser.add_argument(
        "--output_video",
        type=str,
        default="video.avi",
        help="Video output in *.avi for example, video.avi")
    parser.add_argument(
        "--checkpoint",
        type=str,
        default="./devc_weights/",
        help="path to checkpoint folder"
    )
    args = parser.parse_args()

    conf.data.image_size = args.image_size
    conf.data.input_path = args.input_path
    conf.data.ref_path = args.ref_path
    conf.data.output_path = args.output_path
    conf.data.output_video = args.output_video
    conf.data.frame_propagation = args.frame_propagation
    conf.nnabla_context.context = args.context
    conf.nnabla_context.device_id: args.device_id
    conf.checkpoint.path: args.checkpoint

    return conf


# video-colorization/deep-exemplar-based-video-colorization/colornet.py

# Copyright 2021 Sony Group Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import nnabla as nn
import nnabla.functions as F
import nnabla.parametric_functions as PF


def conv_rl(
    x, out_ch, kernel=(
        3, 3), stride=(
            1, 1), pad=(
                1, 1), dilation=None, name=None, act=True):
    conv = PF.convolution(
        x,
        out_ch,
        kernel=kernel,
        stride=stride,
        pad=pad,
        dilation=dilation,
        name=name)
    if act:
        conv = F.relu(conv)
    return conv


def ins_norm(x):
    conv_i_norm = F.instance_normalization(
        x, gamma=None, beta=None, channel_axis=1)
    return conv_i_norm


def conv_norm_ss(
        x,
        out_ch,
        kernel=(
            1,
            1),
    stride=(
            2,
            2),
        with_bias=False,
        group=1,
        name=None):
    conv = PF.convolution(
        x,
        out_ch,
        kernel=kernel,
        stride=stride,
        with_bias=with_bias,
        group=group,
        name=name)
    return conv


def conv_up(x, out_ch, kernel=(3, 3), stride=(1, 1), pad=(1, 1), name=None):
    upsample = F.interpolate(
        x,
        scale=(
            2,
            2),
        mode='nearest',
        align_corners=False)
    conv = PF.convolution(
        upsample,
        out_ch,
        kernel=kernel,
        pad=pad,
        stride=stride,
        name=name)
    return conv


def colorvidnet(x):
    '''
    Colorization Network 
    Args:
        x : NNabla Variable 
    Returns:
        Prediction output 
    '''
    conv1_0 = F.relu(
        PF.convolution(
            x, 32, kernel=(
                3, 3), stride=(
                1, 1), pad=(
                    1, 1), name='conv1_1.0'))
    conv1_1 = conv_rl(conv1_0, 64, name='conv1_1.2')
    conv1_2 = conv_rl(conv1_1, 64, name='conv1_2')
    conv1_2norm = ins_norm(conv1_2)
    conv1_2norm_ss = conv_norm_ss(
        conv1_2norm, 64, group=64, name='conv1_2norm_ss')
    conv2_1 = conv_rl(conv1_2norm_ss, 128, name='conv2_1')
    conv2_2 = conv_rl(conv2_1, 128, name='conv2_2')
    conv2_2norm = ins_norm(conv2_2)
    conv2_2norm_ss = conv_norm_ss(
        conv2_2norm, 128, group=128, name='conv2_2norm_ss')
    conv3_1 = conv_rl(conv2_2norm_ss, 256, name='conv3_1')
    conv3_2 = conv_rl(conv3_1, 256, name='conv3_2')
    conv3_3 = conv_rl(conv3_2, 256, name='conv3_3')
    conv3_3norm = ins_norm(conv3_3)
    conv3_3norm_ss = conv_norm_ss(
        conv3_3norm, 256, group=256, name='conv3_3norm_ss')
    conv4_1 = conv_rl(conv3_3norm_ss, 512, name='conv4_1')
    conv4_2 = conv_rl(conv4_1, 512, name='conv4_2')
    conv4_3 = conv_rl(conv4_2, 512, name='conv4_3')
    conv4_3norm = ins_norm(conv4_3)
    conv5_1 = conv_rl(
        conv4_3norm, 512, pad=(
            2, 2), dilation=(
            2, 2), name='conv5_1')
    conv5_2 = conv_rl(
        conv5_1, 512, pad=(
            2, 2), dilation=(
            2, 2), name='conv5_2')
    conv5_3 = conv_rl(
        conv5_2, 512, pad=(
            2, 2), dilation=(
            2, 2), name='conv5_3')
    conv5_3norm = ins_norm(conv5_3)
    conv6_1 = conv_rl(
        conv5_3norm, 512, pad=(
            2, 2), dilation=(
            2, 2), name='conv6_1')
    conv6_2 = conv_rl(
        conv6_1, 512, pad=(
            2, 2), dilation=(
            2, 2), name='conv6_2')
    conv6_3 = conv_rl(
        conv6_2, 512, pad=(
            2, 2), dilation=(
            2, 2), name='conv6_3')
    conv6_3norm = ins_norm(conv6_3)
    conv7_1 = conv_rl(conv6_3norm, 512, name='conv7_1')
    conv7_2 = conv_rl(conv7_1, 512, name='conv7_2')
    conv7_3 = conv_rl(conv7_2, 512, name='conv7_3')
    conv7_3norm = ins_norm(conv7_3)
    conv8_1 = conv_up(conv7_3norm, 256, name='conv8_1.1')
    conv3_3_short = conv_rl(conv3_3norm, 256, name='conv3_3_short', act=False)
    conv8_1_comb = F.relu(conv8_1 + conv3_3_short)
    conv8_2 = conv_rl(conv8_1_comb, 256, name='conv8_2')
    conv8_3 = conv_rl(conv8_2, 256, name='conv8_3')
    conv8_3norm = ins_norm(conv8_3)
    conv9_1 = conv_up(conv8_3norm, 128, name='conv9_1.1')
    conv2_2_short = conv_rl(conv2_2norm, 128, name='conv2_2_short', act=False)
    conv9_1_comb = F.relu(conv9_1 + conv2_2_short)
    conv9_2 = conv_rl(conv9_1_comb, 128, name='conv9_2')
    conv9_2norm = ins_norm(conv9_2)
    conv10_1 = conv_up(conv9_2norm, 128, name='conv10_1.1')
    conv1_2_short = conv_rl(conv1_2norm, 128, name='conv1_2_short', act=False)
    conv10_1_comb = F.relu(conv10_1 + conv1_2_short)
    conv10_2 = F.leaky_relu(
        conv_rl(
            conv10_1_comb,
            128,
            name='conv10_2',
            act=False),
        alpha=0.2,
        inplace=True)
    conv10_ab = conv_rl(
        conv10_2, 2, kernel=(
            1, 1), stride=(
            1, 1), pad=None, name='conv10_ab', act=False)
    pred_ab = F.tanh(conv10_ab) * 128
    return pred_ab


# video-colorization/deep-exemplar-based-video-colorization/inference.py

# Copyright 2021 Sony Group Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import time
import cv2

import numpy as np
import nnabla as nn
import nnabla.functions as F
from nnabla.ext_utils import get_extension_context
from PIL import Image

from args import get_config
import preprocess
from vggnet import vgg_net
from nonlocal_net import frame_colorization


def transform(image, image_size):
    '''
    Preprocess the given image
    Args:
        image(cv2/numpy) : image to be preprocessed
        image_size: Output image size
    Returns:
        numpy/cv2 image of shape image_size

    '''
    image = preprocess.centerpad(image, image_size)
    image = preprocess.center_crop(image, image_size)
    image = preprocess.rgb2lab(image)
    image = np.transpose(image, [2, 0, 1])
    image = preprocess.normalize(image)
    return np.expand_dims(image, 0)


def interpolate(image, scale):
    '''
    Resize the given input with linear interpolation
    Args:
        image: numpy image of ndim 4 (B,C,H,W)
        scale: scale factor
    Returns:
        batched resized image(numpy)
    '''
    # Change order  [B, C, H, W] -> [H, W, B*C]
    image = np.transpose(image.squeeze(), (1, 2, 0))
    image = cv2.resize(image, None, fx=scale, fy=scale,
                       interpolation=cv2.INTER_LINEAR)
    # Change order  [H, W, C] -> [B, C, H, W]
    image = np.expand_dims(np.transpose(image, (2, 0, 1)), 0)
    return image


def interpolate_nn(image, frame, scale):
    '''
    Linear Interpolation on Variable image and frame
    Args:
        image  : image (Variable) 
        frame  : frame (Variable) 
    Returns
        linear interpolated image and frame 
    '''
    image = F.interpolate(image, scale)
    frame = F.interpolate(frame, scale)
    return image, frame


def get_rgb_frame(ia_lab_large, i_current_ab_predict_nn, conf):

    curr_bs_l = ia_lab_large[:, 0:1, :, :]
    curr_predict = interpolate(
        i_current_ab_predict_nn.data.data,
        scale=2) * 1.25
    if conf.wls_filter_on:
        guide_image = preprocess.uncenter_l(curr_bs_l, conf) * 255 / 100
        wls_filter = cv2.ximgproc.createFastGlobalSmootherFilter(
            guide_image[0, 0, :, :].astype(np.uint8), conf.lambda_value, conf.sigma_color)
        curr_predict_a = wls_filter.filter(
            curr_predict[0, 0, :, :])
        curr_predict_b = wls_filter.filter(
            curr_predict[0, 1, :, :])
        curr_predict_a = np.reshape(
            curr_predict_a, (1, 1, curr_predict_a.shape[0], curr_predict_a.shape[1]))
        curr_predict_b = np.reshape(
            curr_predict_b, (1, 1, curr_predict_b.shape[0], curr_predict_b.shape[1]))
        curr_predict_filter = np.concatenate(
            (curr_predict_a, curr_predict_b), axis=1)
        ia_predict_rgb = preprocess.batch_lab2rgb_transpose(
            conf, curr_bs_l[:32], curr_predict_filter[:32, ...])
    else:
        ia_predict_rgb = preprocess.batch_lab2rgb_transpose(
            conf, curr_bs_l[:32], curr_predict[:32, ...])
    return ia_predict_rgb


def colorize_video(conf, ref):
    '''
    Colorize the input frames and save the output as colorized frames and video
    Args:
        conf: conf object
        ref: refrence image
    '''
    def load_weights():
        nn.load_parameters(f'{conf.checkpoint.path}/{conf.checkpoint.vgg19}')
        nn.load_parameters(
            f'{conf.checkpoint.path}/{conf.checkpoint.non_local}')
        nn.load_parameters(
            f'{conf.checkpoint.path}/{conf.checkpoint.colornet}')

    reference_file = os.path.join(conf.data.ref_path, ref)
    output_path = os.path.join(
        conf.data.output_path, 'out_' + ref.split(".")[0])
    if not os.path.exists(output_path):
        os.makedirs(output_path)
    filenames = [f for f in os.listdir(conf.data.input_path)
                 if os.path.isfile(os.path.join(conf.data.input_path, f))]
    print(f"processing the folder: {conf.data.input_path}")
    # sort the frames in order as in video
    filenames.sort(key=lambda f: int("".join(filter(str.isdigit, f) or -1)))
    # read reference name from reference input else first frame assuming it's
    # colorized
    ref_name = conf.data.input_path + \
        filenames[0] if conf.data.frame_propagation else reference_file
    i_last_lab_predict = None
    # Load the Weights
    nn.clear_parameters()
    load_weights()
    print(f"reference = {ref_name}")
    # Preprocess reference image
    frame_ref = np.array(Image.open(ref_name))
    ib_lab_large = nn.Variable.from_numpy_array(
        transform(frame_ref, conf.data.image_size))
    for iter_num, frame_name in enumerate(filenames):
        print("input =", frame_name)
        frame = Image.open(os.path.join(conf.data.input_path, frame_name))
        ia_lab_large = nn.Variable.from_numpy_array(
            transform(np.array(frame), conf.data.image_size))
        ia_lab, ib_lab = interpolate_nn(
            ia_lab_large, ib_lab_large, scale=(0.5, 0.5))
        ia_l = ia_lab[:, 0:1, :, :]
        if i_last_lab_predict is None:
            if conf.data.frame_propagation:
                i_last_lab_predict = ib_lab
            else:
                i_last_lab_predict = nn.Variable(ia_lab.shape)
        i_reference_l = ib_lab[:, 0:1, :, :]
        i_reference_ab = ib_lab[:, 1:3, :, :]
        i_reference_rgb = preprocess.lab2rgb(
            F.concatenate(
                preprocess.uncenter_l(i_reference_l, conf),
                i_reference_ab,
                axis=1))
        if type(i_last_lab_predict).__module__ == "numpy":
            i_last_lab_predict_nn = nn.Variable.from_numpy_array(
                i_last_lab_predict)
        else:
            i_last_lab_predict_nn = i_last_lab_predict
        t_start = time.time()
        features_b_nn = vgg_net(i_reference_rgb, pre_process=True, fix=True)
        i_current_ab_predict, _i_current_nonlocal_lab, _features_gray = frame_colorization(
            ia_lab, ib_lab, i_last_lab_predict_nn,
            features_b_nn, feature_noise=0, temperature=1e-10)
        # forward the network
        nn.forward_all([i_current_ab_predict])
        i_last_lab_predict = np.concatenate(
            (ia_l.data.data, i_current_ab_predict.data.data), axis=1)
        print(f"Runtime: {time.time() - t_start:.2g} second")
        rgb_frame = get_rgb_frame(ia_lab_large.d, i_current_ab_predict, conf)
        preprocess.save_frames(rgb_frame, output_path, iter_num)
        iter_num = iter_num + 1
    # save the video
    preprocess.frames2vid(
        frame_folder=output_path,
        frame_shape=conf.data.image_size,
        output_dir=output_path,
        filename=conf.data.output_video)


def main():
    conf = get_config()
    ctx = get_extension_context(
        conf.nnabla_context.context, device_id=conf.nnabla_context.device_id)
    nn.set_default_context(ctx)
    refs = sorted(os.listdir(conf.data.ref_path))
    # sort the reference images in order
    # Inference the input frames taking each reference image
    for ref in refs:
        colorize_video(conf, ref)


if __name__ == "__main__":
    main()


# video-colorization/deep-exemplar-based-video-colorization/nonlocal_net.py


# Copyright 2021 Sony Group Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import sys

import nnabla as nn
import nnabla.functions as F
import nnabla.parametric_functions as PF

from vggnet import vgg_net
from colornet import colorvidnet
from utils_nn import gray2rgb_batch, feature_normalize


def res_block(x, out_ch, name):
    with nn.parameter_scope(name):
        residual = x
        out = F.pad(x, (1, 1, 1, 1), 'reflect')
        out = PF.convolution(
            out, out_ch, kernel=(
                3, 3), stride=(
                1, 1), name='conv1')
        out = F.instance_normalization(
            out, gamma=None, beta=None, channel_axis=1)
        out = PF.prelu(out)
        out = F.pad(out, (1, 1, 1, 1), 'reflect')
        out = PF.convolution(
            out, out_ch, kernel=(
                3, 3), stride=(
                1, 1), name='conv2')
        out = F.instance_normalization(
            out, gamma=None, beta=None, channel_axis=1)
        out += residual
        out = PF.prelu(out)
    return out


def layer2_1(x):
    pad2_1 = F.pad(x, (1, 1, 1, 1), 'reflect')
    conv2_1 = PF.convolution(
        pad2_1, 128, kernel=(
            3, 3), stride=(
            1, 1), name='layer2_1.1')
    conv2_1 = F.instance_normalization(
        conv2_1, gamma=None, beta=None, channel_axis=1)
    conv2_1 = PF.prelu(conv2_1, name='layer2_1.3')
    pad2_2 = F.pad(conv2_1, (1, 1, 1, 1), 'reflect')
    conv2_2 = PF.convolution(
        pad2_2, 64, kernel=(
            3, 3), stride=(
            2, 2), name='layer2_1.5')
    conv2_2 = F.instance_normalization(
        conv2_2, gamma=None, beta=None, channel_axis=1)
    conv2_2 = PF.prelu(conv2_2, name='layer2_1.7')
    return conv2_2


def layer3_1(x):
    pad3_1 = F.pad(x, (1, 1, 1, 1), 'reflect')
    conv3_1 = PF.convolution(
        pad3_1, 128, kernel=(
            3, 3), stride=(
            1, 1), name='layer3_1.1')
    conv3_1 = F.instance_normalization(
        conv3_1, gamma=None, beta=None, channel_axis=1)
    conv3_1 = PF.prelu(conv3_1, name='layer3_1.3')
    pad3_2 = F.pad(conv3_1, (1, 1, 1, 1), 'reflect')
    conv3_2 = PF.convolution(
        pad3_2, 64, kernel=(
            3, 3), stride=(
            1, 1), name='layer3_1.5')
    conv3_2 = F.instance_normalization(
        conv3_2, gamma=None, beta=None, channel_axis=1)
    conv3_2 = PF.prelu(conv3_2, name='layer3_1.7')
    return conv3_2


def layer4_1(x):
    pad4_1 = F.pad(x, (1, 1, 1, 1), 'reflect')
    conv4_1 = PF.convolution(
        pad4_1, 256, kernel=(
            3, 3), stride=(
            1, 1), name='layer4_1.1')
    conv4_1 = F.instance_normalization(
        conv4_1, gamma=None, beta=None, channel_axis=1)
    conv4_1 = PF.prelu(conv4_1, name='layer4_1.3')
    pad4_2 = F.pad(conv4_1, (1, 1, 1, 1), 'reflect')
    conv4_2 = PF.convolution(
        pad4_2, 64, kernel=(
            3, 3), stride=(
            1, 1), name='layer4_1.5')
    conv4_2 = F.instance_normalization(
        conv4_2, gamma=None, beta=None, channel_axis=1)
    conv4_2 = PF.prelu(conv4_2, name='layer4_1.7')
    up4_1 = F.interpolate(
        conv4_2,
        scale=(
            2,
            2),
        mode='nearest',
        align_corners=False)
    return up4_1


def layer5_1(x):
    pad5_1 = F.pad(x, (1, 1, 1, 1), 'reflect')
    conv5_1 = PF.convolution(
        pad5_1, 256, kernel=(
            3, 3), stride=(
            1, 1), name='layer5_1.1')
    conv5_1 = F.instance_normalization(
        conv5_1, gamma=None, beta=None, channel_axis=1)
    conv5_1 = PF.prelu(conv5_1, name='layer5_1.3')
    up5_1 = F.interpolate(
        conv5_1,
        scale=(
            2,
            2),
        mode='nearest',
        align_corners=False)
    pad5_2 = F.pad(up5_1, (1, 1, 1, 1), 'reflect')
    conv5_2 = PF.convolution(
        pad5_2, 64, kernel=(
            3, 3), stride=(
            1, 1), name='layer5_1.6')
    conv5_2 = F.instance_normalization(
        conv5_2, gamma=None, beta=None, channel_axis=1)
    conv5_2 = PF.prelu(conv5_2, name='layer5_1.8')
    up5_2 = F.interpolate(
        conv5_2,
        scale=(
            2,
            2),
        mode='nearest',
        align_corners=False)
    return up5_2


def layer(x, out_channel):
    res_0 = res_block(x, out_channel, name='layer.0')
    res_1 = res_block(res_0, out_channel, name='layer.1')
    res_2 = res_block(res_1, out_channel, name='layer.2')
    return res_2


def pad_replicate(x):
    start = x[:, :, 0, :]
    end = x[:, :, -1, :]
    new = F.pad(x, (1, 1, 0, 0), 'reflect')
    new[:, :, 0, :] = start
    new[:, :, -1, :] = end
    return new


def nonlocal_net(B_lab_map,
                 relu_layers,
                 temperature=0.001 * 5,
                 detach_flag=False,
                 WTA_scale_weight=1,
                 feature_noise=0):

    batch_size = B_lab_map.shape[0]
    channel = B_lab_map.shape[1]
    image_height = B_lab_map.shape[2]
    image_width = B_lab_map.shape[3]
    feature_height = int(image_height / 4)
    feature_width = int(image_width / 4)

    feature_channel = 64
    in_channels = feature_channel * 4
    inter_channels = 256

    # layer2_1
    A_feature2_1 = layer2_1(relu_layers[0])
    B_feature2_1 = layer2_1(relu_layers[4])
    # layer3_1
    A_feature3_1 = layer3_1(relu_layers[1])
    B_feature3_1 = layer3_1(relu_layers[5])
    # layer4_1
    A_feature4_1 = layer4_1(relu_layers[2])
    B_feature4_1 = layer4_1(relu_layers[6])
    # layer5_1
    A_feature5_1 = layer5_1(relu_layers[3])
    B_feature5_1 = layer5_1(relu_layers[7])

    if A_feature5_1.shape[2] != A_feature2_1.shape[2] or A_feature5_1.shape[3] != A_feature2_1.shape[3]:
        A_feature5_1 = pad_replicate(A_feature5_1)
        B_feature5_1 = pad_replicate(B_feature5_1)
    A_features = layer(
        F.concatenate(
            A_feature2_1,
            A_feature3_1,
            A_feature4_1,
            A_feature5_1,
            axis=1),
        feature_channel * 4)
    B_features = layer(
        F.concatenate(
            B_feature2_1,
            B_feature3_1,
            B_feature4_1,
            B_feature5_1,
            axis=1),
        feature_channel * 4)
    # pairwise cosine similarity
    theta = PF.convolution(
        A_features, inter_channels, kernel=(
            1, 1), stride=(
            1, 1), name='theta')
    theta_re = F.reshape(theta, (batch_size, inter_channels, -1))
    theta_re = theta_re - F.mean(theta_re, axis=2,
                                 keepdims=True)  # center the feature
    theta_norm = F.norm(
        theta_re,
        p=2,
        axis=1,
        keepdims=True) + sys.float_info.epsilon
    theta_re = F.div2(theta_re, theta_norm)
    # 2*(feature_height*feature_width)*256
    theta_permute = F.transpose(theta_re, (0, 2, 1))
    phi = PF.convolution(
        B_features, inter_channels, kernel=(
            1, 1), stride=(
            1, 1), name='phi')
    phi_re = F.reshape(phi, (batch_size, inter_channels, -1))
    # center the feature
    phi_re = phi_re - F.mean(phi_re, axis=2, keepdims=True)
    phi_norm = F.norm(phi_re, p=2, axis=1, keepdims=True) + \
        sys.float_info.epsilon
    phi_re = F.div2(phi_re, phi_norm)
    # 2*(feature_height*feature_width)*(feature_height*feature_width)
    f = F.batch_matmul(theta_permute, phi_re)

    f_shape = f.shape
    f = F.reshape(f, (1,) + f_shape)
    f_similarity = F.reshape(f, (1,) + f_shape)
    similarity_map = F.max(f_similarity, axis=3, keepdims=True)
    similarity_map = F.reshape(
        similarity_map, (batch_size, 1, feature_height, feature_width))

    # f can be negative
    # if WTA_scale_weight == 1:
    f_WTA = f

    f_WTA = f_WTA / temperature

    f_WTA_sp = f_WTA.shape
    f_WTA = F.reshape(f_WTA, (f_WTA_sp[1], f_WTA_sp[2], f_WTA_sp[3]))
    # 2*1936*1936; softmax along the horizontal line (dim=-1)
    f_div_C = F.softmax(f_WTA, axis=2)

    # downsample the reference color
    B_lab = F.average_pooling(B_lab_map, (4, 4))
    B_lab = F.reshape(B_lab, (batch_size, channel, -1))
    B_lab = F.transpose(B_lab, (0, 2, 1))  # 2*1936*channel

    # multiply the corr map with color
    y = F.batch_matmul(f_div_C, B_lab)  # 2*1936*channel
    y = F.transpose(y, (0, 2, 1))
    y = F.reshape(
        y,
        (batch_size,
         channel,
         feature_height,
         feature_width))  # 2*3*44*44
    y = F.interpolate(y, scale=(4, 4), mode='nearest', align_corners=False)
    similarity_map = F.interpolate(
        similarity_map, scale=(
            4, 4), mode='nearest', align_corners=False)

    return y, similarity_map


def warp_color(IA_l, IB_lab, features_B, feature_noise=0, temperature=0.01):
    '''
        joint_train=True, enable grad; otherwise disable grad
    '''
    # change to rgb for feature extraction
    IA_rgb_from_gray = gray2rgb_batch(IA_l)

    A_relu1_1, A_relu2_1, A_relu3_1, A_relu4_1, A_relu5_1 = vgg_net(
        IA_rgb_from_gray, pre_process=True, fix=True)
    B_relu1_1, B_relu2_1, B_relu3_1, B_relu4_1, B_relu5_1 = features_B

    # NOTE: output the feature before normalization
    features_A = [A_relu1_1, A_relu2_1, A_relu3_1, A_relu4_1, A_relu5_1]
    layers = [A_relu2_1, A_relu3_1, A_relu4_1, A_relu5_1,
              B_relu2_1, B_relu3_1, B_relu4_1, B_relu5_1]

    def feature_norm_relu(layers):
        f_norms = []
        for relu_l in layers:
            f_norm = feature_normalize(relu_l)
            f_norms.append(f_norm)
        return f_norms
    layers = feature_norm_relu(layers)
    nonlocal_BA_lab, similarity_map = nonlocal_net(
        IB_lab, layers, temperature=0.01)

    return nonlocal_BA_lab, similarity_map, features_A


def frame_colorization(IA_lab,
                       IB_lab,
                       IA_last_lab,
                       features_B,
                       joint_training=True,
                       feature_noise=0,
                       luminance_noise=0,
                       temperature=0.01):
    # change to rgb for feature extraction
    IA_l = IA_lab[:, 0:1, :, :]

    # if luminance_noise:
    nonlocal_BA_lab, similarity_map, features_A_gray = warp_color(
        IA_l, IB_lab, features_B, feature_noise, temperature=temperature)
    nonlocal_BA_ab = nonlocal_BA_lab[:, 1:3, :, :]
    color_input = F.concatenate(
        IA_l,
        nonlocal_BA_ab,
        similarity_map,
        IA_last_lab,
        axis=1)
    with nn.parameter_scope('colornet'):
        IA_ab_predict = colorvidnet(color_input)

    return IA_ab_predict, nonlocal_BA_lab, features_A_gray


# video-colorization/deep-exemplar-based-video-colorization/preprocess.py


# Copyright 2021 Sony Group Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import numbers
import collections
import cv2
import numpy as np
from skimage.transform import resize
from PIL import Image, ImageOps
import nnabla.functions as F
from skimage import color, io


def rgb2lab(img):
    return color.rgb2lab(img)


def centerpad(img, newsize):
    height_old = img.shape[0]
    width_old = img.shape[1]
    height = newsize[0]
    width = newsize[1]
    img_pad = np.zeros((height, width, img.shape[2]))
    ratio = height / width
    old_size = [height_old, width_old]
    if height_old / width_old == ratio:
        if height_old == height:
            return img
        new_size = [int(x * height / height_old) for x in old_size]
        img_resize = resize(
            img,
            new_size,
            mode="reflect",
            preserve_range=True,
            clip=False,
            anti_aliasing=True)
        return img_resize
    if height_old / width_old > ratio:  # pad the width and crop
        new_size = [int(x * width / width_old) for x in old_size]
        img_resize = resize(
            img,
            new_size,
            mode="reflect",
            preserve_range=True,
            clip=False,
            anti_aliasing=True)
        height_resize = img_resize.shape[0]
        start_height = (height_resize - height) // 2
        img_pad[:, :, :] = img_resize[start_height: (
            start_height + height), :, :]
    else:
        new_size = [int(x * height / height_old) for x in old_size]
        img_resize = resize(
            img,
            new_size,
            mode="reflect",
            preserve_range=True,
            clip=False,
            anti_aliasing=True)
        width_resize = img_resize.shape[1]
        start_width = (width_resize - width) // 2
        img_pad[:, :, :] = img_resize[:, start_width: (start_width + width), :]
    return img_pad


def crop_func(img, i, j, h, w):
    """Crop the given PIL Image.

    Args:
        img (PIL Image): Image to be cropped.
        i: Upper pixel coordinate.
        j: Left pixel coordinate.
        h: Height of the cropped image.
        w: Width of the cropped image.

    Returns:
        PIL Image: Cropped image.
    """
    # if not _is_pil_image(img):
    #    raise TypeError('img should be PIL Image. Got {}'.format(type(img)))

    return img.crop((j, i, j + w, i + h))


def center_crop(inputs, size, padding=0):
    """
    Crop the given PIL Image at a random location.

    Args:
        img (PIL Image): Image to be cropped.
        size (sequence or int): Desired output size of the crop. If size is an
            int instead of sequence like (h, w), a square crop (size, size) is
            made.
        padding (int or sequence, optional): Optional padding on each border
            of the image. Default is 0, i.e no padding. If a sequence of length
            4 is provided, it is used to pad left, top, right, bottom borders
            respectively.
        Returns:
            PIL Image: Cropped image.
    """
    def get_params(img, output_size):
        """
        Get parameters for 'crop' for a random crop.

        Args:
            img (PIL Image): Image to be cropped.
            output_size (tuple): Expected output size of the crop.

        Returns:
            tuple: params (i, j, h, w) to be passed to 'crop' for random crop.
        """
        # w, h = img.shape[:-1]
        w, h = img.size
        th, tw = output_size
        if w == tw and h == th:
            return 0, 0, h, w

        i = (h - th) // 2
        j = (w - tw) // 2
        return i, j, th, tw
    inputs = Image.fromarray(inputs.astype(np.uint8))
    if isinstance(size, numbers.Number):
        size = (int(size), int(size))

    if isinstance(inputs, list):
        i, j, h, w = get_params(inputs[0], size)
    else:
        i, j, h, w = get_params(inputs, size)
    image = crop_func(inputs, i, j, h, w)
    return np.array(image)


def standardize(arr, mean, std):
    """
    Standardize the input with mean and standard deviation.

    Args:
        arr (nd array): arr/image of size (C, H, W) to be normalized.
        mean (sequence): Sequence of means for each channel.
        std (sequence): Sequence of standard deviations for each channely.

    Returns:
        norm_arr: Normalized Tensor image.
    """
    if (arr.ndim < 3):
        arr = np.expand_dims(arr, axis=0)
    if arr.shape[0] == 1:
        norm_arr = arr.copy()
        norm_arr[0, :, :] = (arr[0, :, :] - mean) / std
    else:
        for i in range(arr.shape[0]):
            norm_arr = arr.copy()
            norm_arr[i, :, :] = (arr[i, :, :] - mean[i]) / std[i]
    return norm_arr


def normalize(image):
    """
    Normalize the input

    Args:
        image :  Input image (numpy)
    Returns:
        Normalized image in dtype float32

    """
    image[0, :, :] = standardize(image[0:1, :, :], 50, 1)
    image[1:3, :, :] = standardize(image[1:3, :, :], (0, 0), (1, 1))
    return image.astype(np.float32)


def lab2rgb(input):
    input_trans = F.split(input, axis=1)
    L, a, b = F.split(input, axis=1)
    y = (L + 16.0) / 116.0
    x = (a / 500.0) + y
    z = y - (b / 200.0)
    neg_mask = F.less_scalar(z, 0).apply(need_grad=False)
    z = z * F.logical_not(neg_mask)
    mask_Y = F.greater_scalar(y, 0.2068966).apply(need_grad=False)
    mask_X = F.greater_scalar(x, 0.2068966).apply(need_grad=False)
    mask_Z = F.greater_scalar(z, 0.2068966).apply(need_grad=False)
    Y_1 = (y ** 3) * mask_Y
    Y_2 = L / (116. * 7.787) * F.logical_not(mask_Y)
    var_Y = Y_1 + Y_2

    X_1 = (x ** 3) * mask_X
    X_2 = (x - 16. / 116.) / 7.787 * F.logical_not(mask_X)
    var_X = X_1 + X_2

    Z_1 = (z ** 3) * mask_Z
    Z_2 = (z - 16. / 116.) / 7.787 * F.logical_not(mask_Z)
    var_Z = Z_1 + Z_2

    X = 0.95047 * var_X
    Y = 1.00000 * var_Y
    Z = 1.08883 * var_Z

    var_R = X * 3.2406 + Y * -1.5372 + Z * -0.4986
    var_G = X * -0.9689 + Y * 1.8758 + Z * 0.0415
    var_B = X * 0.0557 + Y * -0.2040 + Z * 1.0570

    mask_R = F.greater_scalar(var_R, 0.0031308).apply(need_grad=False)
    n_mask_R = F.logical_not(mask_R)
    R_1 = (1.055 * (F.maximum2(var_R, n_mask_R) ** (1 / 2.4)) - 0.055) * mask_R
    R_2 = (12.92 * var_R) * n_mask_R
    var_R = R_1 + R_2

    mask_G = F.greater_scalar(var_G, 0.0031308).apply(need_grad=False)
    n_mask_G = F.logical_not(mask_G)
    G_1 = (1.055 * (F.maximum2(var_G, n_mask_G) ** (1 / 2.4)) - 0.055) * mask_G
    G_2 = (12.92 * var_G) * n_mask_G
    var_G = G_1 + G_2

    mask_B = F.greater_scalar(var_B, 0.0031308).apply(need_grad=False)
    n_mask_B = F.logical_not(mask_B)
    B_1 = (1.055 * (F.maximum2(var_B, n_mask_B) ** (1 / 2.4)) - 0.055) * mask_B
    B_2 = (12.92 * var_B) * n_mask_B
    var_B = B_1 + B_2
    return F.stack(var_R, var_G, var_B, axis=1)


def uncenter_l(l, conf):
    return l * conf.l_norm + conf.l_mean


def batch_lab2rgb_transpose(conf, img_l_mc, img_ab_mc, nrow=8):
    img_l_mc = img_l_mc
    img_ab_mc = img_ab_mc

    assert img_l_mc.ndim == 4 and img_ab_mc.ndim == 4, "only for batch input"

    img_l = img_l_mc * conf.l_norm + conf.l_mean
    img_ab = img_ab_mc * conf.ab_norm + conf.ab_mean
    pred_lab = np.concatenate((img_l, img_ab), axis=1)
    grid_lab = pred_lab.squeeze().astype(np.float64)
    return (
        np.clip(
            color.lab2rgb(
                grid_lab.transpose(
                    (1, 2, 0))), 0, 1) * 255).astype("uint8")


def save_frames(image, image_folder, index=None, frame_name=None):
    if image is not None:
        image = np.clip(image, 0, 255).astype(np.uint8)
        if frame_name:
            io.imsave(os.path.join(image_folder, frame_name), image)
        else:
            io.imsave(
                os.path.join(
                    image_folder,
                    str(index).zfill(5) +
                    '.jpg'),
                image)


def frames2vid(frame_folder, frame_shape, output_dir, filename):
    frames = sorted([img for img in os.listdir(
        frame_folder) if img.endswith(".jpg")])
    # get the height and width
    height, width = frame_shape
    print(f"writing to video file: {os.path.join(output_dir, filename)}")
    video = cv2.VideoWriter(
        os.path.join(
            output_dir, filename), cv2.VideoWriter_fourcc(
            'D', 'I', 'V', 'X'), 24, (width, height))
    for frame in frames:
        video.write(cv2.imread(os.path.join(frame_folder, frame)))
    video.release()


# video-colorization/deep-exemplar-based-video-colorization/utils_nn.py


# Copyright 2021 Sony Group Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import sys

import nnabla as nn
import nnabla.functions as F
import numpy as np

from preprocess import uncenter_l
from args import get_config

conf = get_config()


def vgg_pre_process(x):
    x_bgr = F.concatenate(
        x[:, 2:3, :, :], x[:, 1:2, :, :], x[:, 0:1, :, :], axis=1)
    # tensor_bgr = tensor[:, [2, 1, 0], ...]
    x_sub = F.reshape(nn.Variable.from_numpy_array(
        np.array([0.40760392, 0.45795686, 0.48501961])), (1, 3, 1, 1))
    x_bgr_ml = x_bgr - x_sub
    x_rst = x_bgr_ml * 255
    return x_rst


def feature_normalize(feature_in):
    feature_in_norm = F.norm(
        feature_in,
        p=2,
        axis=1,
        keepdims=True) + sys.float_info.epsilon
    feature_in_norm = F.div2(feature_in, feature_in_norm)
    return feature_in_norm


def gray2rgb_batch(l):
    # gray image tensor to rgb image tensor
    l_uncenter = uncenter_l(l, conf)
    l_uncenter = l_uncenter / (2 * conf.l_mean)
    return F.concatenate(l_uncenter, l_uncenter, l_uncenter, axis=1)


# video-colorization/deep-exemplar-based-video-colorization/vggnet.py

# Copyright 2021 Sony Group Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import nnabla.functions as F
import nnabla.parametric_functions as PF
from utils_nn import vgg_pre_process


def vgg_net(input_var, pre_process=True, fix=True):
    if pre_process:
        input_var = vgg_pre_process(input_var)

    act1_1 = F.relu(PF.convolution(input_var, outmaps=64, kernel=(
        3, 3), pad=(1, 1), fix_parameters=fix, name="conv1_1"))
    act1_2 = F.relu(PF.convolution(act1_1, outmaps=64, kernel=(
        3, 3), pad=(1, 1), fix_parameters=fix, name="conv1_2"))

    p1 = F.max_pooling(act1_2, kernel=(2, 2), stride=(2, 2))
    act2_1 = F.relu(PF.convolution(p1, outmaps=128, kernel=(
        3, 3), pad=(1, 1), fix_parameters=fix, name="conv2_1"))
    act2_2 = F.relu(PF.convolution(act2_1, outmaps=128, kernel=(
        3, 3), pad=(1, 1), fix_parameters=fix, name="conv2_2"))

    p2 = F.max_pooling(act2_2, kernel=(2, 2), stride=(2, 2))
    act3_1 = F.relu(PF.convolution(p2, outmaps=256, kernel=(
        3, 3), pad=(1, 1), fix_parameters=fix, name="conv3_1"))
    act3_2 = F.relu(PF.convolution(act3_1, outmaps=256, kernel=(
        3, 3), pad=(1, 1), fix_parameters=fix, name="conv3_2"))
    act3_3 = F.relu(PF.convolution(act3_2, outmaps=256, kernel=(
        3, 3), pad=(1, 1), fix_parameters=fix, name="conv3_3"))
    act3_4 = F.relu(PF.convolution(act3_3, outmaps=256, kernel=(
        3, 3), pad=(1, 1), fix_parameters=fix, name="conv3_4"))

    p3 = F.max_pooling(act3_4, kernel=(2, 2), stride=(2, 2))
    act4_1 = F.relu(PF.convolution(p3, outmaps=512, kernel=(
        3, 3), pad=(1, 1), fix_parameters=fix, name="conv4_1"))
    act4_2 = F.relu(PF.convolution(act4_1, outmaps=512, kernel=(
        3, 3), pad=(1, 1), fix_parameters=fix, name="conv4_2"))
    act4_3 = F.relu(PF.convolution(act4_2, outmaps=512, kernel=(
        3, 3), pad=(1, 1), fix_parameters=fix, name="conv4_3"))
    act4_4 = F.relu(PF.convolution(act4_3, outmaps=512, kernel=(
        3, 3), pad=(1, 1), fix_parameters=fix, name="conv4_4"))

    p4 = F.max_pooling(act4_4, kernel=(2, 2), stride=(2, 2))
    act5_1 = F.relu(PF.convolution(p4, outmaps=512, kernel=(
        3, 3), pad=(1, 1), fix_parameters=fix, name="conv5_1"))
    act5_2 = F.relu(PF.convolution(act5_1, outmaps=512, kernel=(
        3, 3), pad=(1, 1), fix_parameters=fix, name="conv5_2"))
    act5_3 = F.relu(PF.convolution(act5_2, outmaps=512, kernel=(
        3, 3), pad=(1, 1), fix_parameters=fix, name="conv5_3"))
    act5_4 = F.relu(PF.convolution(act5_3, outmaps=512, kernel=(
        3, 3), pad=(1, 1), fix_parameters=fix, name="conv5_4"))

    return [act1_2, act2_2, act3_2, act4_2, act5_2]


# video-superresolution/jsigan/args.py

# Copyright 2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
from utils import read_yaml


def get_config():
    """
    Get command line arguments.
    Arguments set the default values of command line arguments
    """
    parser = argparse.ArgumentParser(description='JSIGAN')
    parser.add_argument('--cfg', default="./config.yaml")
    args, _ = parser.parse_known_args()

    conf = read_yaml(args.cfg)

    parser.add_argument('--lr-sdr-test', type=str, default="datasets/",
                        help='The directory of the input data, for testing')
    parser.add_argument('--hr-hdr-test', type=str, default="datasets/",
                        help='The directory of the input data, for testing')
    parser.add_argument('--lr-sdr-train', type=str, default="datasets/",
                        help='The directory of the input data, for training')
    parser.add_argument('--hr-hdr-train', type=str, default="datasets/",
                        help='The directory of the input data, for training')
    parser.add_argument('--scaling-factor', type=int, default=4,
                        help='LR to HR scaling factor')
    parser.add_argument('--pre-trained-model', type=str, default="models/",
                        help='Path of the pre trained weights')
    parser.add_argument('--output-dir', type=str, default="models/",
                        help='Path to save weight files during training')
    parser.add_argument('--jsigan', action='store_true', default=False,
                        help='If True, GAN network will be trained, False otherwise')
    parser.add_argument('--save-images', action='store_true', default=False,
                        help='If True, images will be saved during inference; False to save time')

    args = parser.parse_args()

    conf.data.lr_sdr_test = args.lr_sdr_test
    conf.data.hr_hdr_test = args.hr_hdr_test
    conf.data.lr_sdr_train = args.lr_sdr_train
    conf.data.hr_hdr_train = args.hr_hdr_train
    conf.scaling_factor = args.scaling_factor
    conf.pre_trained_model = args.pre_trained_model
    conf.output_dir = args.output_dir
    conf.jsigan = args.jsigan
    conf.save_images = args.save_images

    return conf


# video-superresolution/jsigan/data_loader.py

# Copyright 2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from nnabla.utils.data_source import DataSource
from nnabla.utils.data_iterator import data_iterator
import h5py
import numpy as np
import nnabla as nn


def read_mat_file(data_fname, label_fname, data_name, label_name, train):
    """
    Read training data from .mat file
    """
    data_file = h5py.File(data_fname, 'r')
    label_file = h5py.File(label_fname, 'r')
    data = data_file[data_name][()]
    label = label_file[label_name][()]

    # change type & reorder
    data = np.array(data, dtype=np.float32) / 255.
    label = np.array(label, dtype=np.float32) / 1023.
    data = np.swapaxes(data, 1, 3)
    label = np.swapaxes(label, 1, 3)
    if not train:
        data = nn.Variable.from_numpy_array(data)
        label = nn.Variable.from_numpy_array(label)
    print('.mat files read successfully.')
    return data, label


class JSIData(DataSource):
    def _get_data(self, position):
        if self.train:
            data_train = self.data_train[self._indexes[position]]
            label_train = self.label_train[self._indexes[position]]

            return data_train, label_train

    def __init__(self, conf, train, shuffle=True, rng=None):
        super(JSIData, self).__init__(shuffle=shuffle, rng=rng)
        self.train = train
        self.data_fname = conf.data.lr_sdr_train
        self.label_fname = conf.data.hr_hdr_train
        self.data_name = conf.data.d_name_train
        self.label_name = conf.data.l_name_train
        self.val_data_size = conf.data.val_data_size
        self.data, self.label = read_mat_file(self.data_fname, self.label_fname, self.data_name,
                                              self.label_name, train)
        self.data_train = self.data[:-self.val_data_size, :, :, :]
        self.label_train = self.label[:-self.val_data_size, :, :, :]
        self._size = len(self.data_train)
        self.rng = rng
        self._variables = ('data_train', 'label_train')

    def reset(self):
        if self._shuffle:
            self._indexes = self._rng.permutation(self._size)
        else:
            self._indexes = np.arange(self._size)
        super(JSIData, self).reset()


def jsi_iterator(batch_size, conf, train, with_memory_cache=False, with_file_cache=False, rng=None):
    return data_iterator(JSIData(conf, train, shuffle=True, rng=None), batch_size, rng,
                         with_memory_cache, with_file_cache)


# video-superresolution/jsigan/generate_lr_hdr.py

# Copyright 2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import math
import argparse
import h5py
import numpy as np


def cubic(x):
    """
    Cubic kernel
    """
    absx = np.abs(x)
    absx2 = absx**2
    absx3 = absx**3
    return (1.5 * absx3 - 2.5 * absx2 + 1) * (
        (absx <= 1).astype(type(absx))) + (-0.5 * absx3 + 2.5 * absx2 - 4 * absx + 2) * ((
            (absx > 1) * (absx <= 2)).astype(type(absx)))


def calculate_weights_indices(in_length, out_length, scale, kernel_width, antialiasing):
    """
    Get weights and indices
    """
    if (scale < 1) and (antialiasing):
        # Use a modified kernel to simultaneously interpolate and antialias- larger kernel width
        kernel_width = kernel_width / scale

    # Output-space coordinates
    x = np.linspace(1, out_length, out_length)

    # Input-space coordinates. Calculate the inverse mapping such that 0.5
    # in output space maps to 0.5 in input space, and 0.5+scale in output
    # space maps to 1.5 in input space.
    u = x / scale + 0.5 * (1 - 1 / scale)

    # What is the left-most pixel that can be involved in the computation?
    left = np.floor(u - kernel_width / 2)

    # What is the maximum number of pixels that can be involved in the
    # computation?  Note: it's OK to use an extra pixel here; if the
    # corresponding weights are all zero, it will be eliminated at the end
    # of this function.
    P = math.ceil(kernel_width) + 2

    # The indices of the input pixels involved in computing the k-th output
    # pixel are in row k of the indices matrix.
    indices = np.repeat(left.reshape(out_length, 1), P).reshape(out_length, P) +  \
        np.broadcast_to(np.linspace(
            0, P - 1, P).reshape(1, P), (out_length, P))

    # The weights used to compute the k-th output pixel are in row k of the
    # weights matrix.
    distance_to_center = np.repeat(
        u.reshape(out_length, 1), P).reshape(out_length, P) - indices

    # apply cubic kernel
    if (scale < 1) and (antialiasing):
        weights = scale * cubic(distance_to_center * scale)
    else:
        weights = cubic(distance_to_center)
    # Normalize the weights matrix so that each row sums to 1.
    weights_sum = np.sum(weights, 1).reshape(out_length, 1)
    weights = weights / np.repeat(weights_sum, P).reshape(out_length, P)

    # If a column in weights is all zero, get rid of it. only consider the first and last column.
    weights_zero_tmp = np.sum((weights == 0), 0)
    if not math.isclose(weights_zero_tmp[0], 0, rel_tol=1e-6):
        # indices = indices.narrow(1, 1, P - 2)
        indices = indices[:, 1:P-1]
        # weights = weights.narrow(1, 1, P - 2)
        weights = weights[:, 1:P-1]
    if not math.isclose(weights_zero_tmp[-1], 0, rel_tol=1e-6):
        # indices = indices.narrow(1, 0, P - 2)
        indices = indices[:, 0:P-1]
        # weights = weights.narrow(1, 0, P - 2)
        weights = weights[:, 0:P-1]

    weights = np.ascontiguousarray(weights)
    indices = np.ascontiguousarray(indices)
    sym_len_s = -indices.min() + 1
    sym_len_e = indices.max() - in_length
    indices = indices + sym_len_s - 1
    return weights, indices, int(sym_len_s), int(sym_len_e)


def imresize_np(img, scale, antialiasing=True):
    """
    The scale should be the same for H and W
    Input: img: Numpy, HWC BGR [0,1]
    Output: HWC BGR [0,1] w/o round
    """
    in_H, in_W, in_C = img.shape
    _, out_H, out_W = in_C, math.ceil(in_H * scale), math.ceil(in_W * scale)
    kernel_width = 4

    # Return the desired dimension order for performing the resize.  The
    # strategy is to perform the resize first along the dimension with the
    # smallest scale factor.
    # Now we do not support this.

    # get weights and indices
    weights_H, indices_H, sym_len_Hs, sym_len_He = calculate_weights_indices(
        in_H, out_H, scale, kernel_width, antialiasing)
    weights_W, indices_W, sym_len_Ws, sym_len_We = calculate_weights_indices(
        in_W, out_W, scale, kernel_width, antialiasing)
    # process H dimension
    # symmetric copying
    img_aug = np.ndarray((in_H + sym_len_Hs + sym_len_He,
                          in_W, in_C), dtype='float32')
    # img = img_aug[sym_len_Hs:sym_len_Hs+in_H,:,:]
    img_aug[sym_len_Hs:sym_len_Hs+in_H, :, :] = img

    sym_patch = img[:sym_len_Hs, :, :]
    inv_idx = np.arange(sym_patch.shape[0] - 1, -1, -1).astype('int64')
    sym_patch_inv = np.take(sym_patch, inv_idx, 0)
    # sym_patch_inv = img_aug[0:sym_len_Hs:,:,:]
    img_aug[0:sym_len_Hs:, :, :] = sym_patch_inv

    sym_patch = img[-sym_len_He:, :, :]
    inv_idx = np.arange(sym_patch.shape[0] - 1, -1, -1).astype('int64')
    sym_patch_inv = np.take(sym_patch, inv_idx, 0)
    # sym_patch_inv = img_aug[sym_len_Hs + in_H:sym_len_Hs + in_H + sym_len_He,:,:]
    img_aug[sym_len_Hs + in_H:sym_len_Hs +
            in_H + sym_len_He, :, :] = sym_patch_inv

    out_1 = np.ndarray((out_H, in_W, in_C), dtype='float32')
    kernel_width = weights_H.shape[1]

    for i in range(out_H):
        idx = int(indices_H[i][0])
        out_1[i, :, 0] = np.dot(
            img_aug[idx:idx + kernel_width, :, 0].transpose(1, 0), weights_H[i])
        out_1[i, :, 1] = np.dot(
            img_aug[idx:idx + kernel_width, :, 1].transpose(1, 0), weights_H[i])
        out_1[i, :, 2] = np.dot(
            img_aug[idx:idx + kernel_width, :, 2].transpose(1, 0), weights_H[i])

    # process W dimension
    # symmetric copying
    out_1_aug = np.ndarray(
        (out_H, in_W + sym_len_Ws + sym_len_We, in_C), dtype='float32')
    # out_1 = out_1_aug[:,sym_len_Ws:sym_len_Ws + in_W]
    out_1_aug[:, sym_len_Ws:sym_len_Ws + in_W] = out_1

    sym_patch = out_1[:, :sym_len_Ws, :]
    inv_idx = np.arange(sym_patch.shape[1] - 1, -1, -1).astype('int64')
    sym_patch_inv = np.take(sym_patch, inv_idx, 1)
    # sym_patch_inv = out_1_aug[:,0:0+sym_len_Ws]
    out_1_aug[:, 0:0+sym_len_Ws] = sym_patch_inv

    sym_patch = out_1[:, -sym_len_We:, :]
    inv_idx = np.arange(sym_patch.shape[1] - 1, -1, -1).astype('int64')
    sym_patch_inv = np.take(sym_patch, inv_idx, 1)
    # sym_patch_inv = out_1_aug[:,sym_len_Ws + in_W:sym_len_Ws + in_W + sym_len_We]
    out_1_aug[:, sym_len_Ws + in_W:sym_len_Ws +
              in_W + sym_len_We] = sym_patch_inv

    out_2 = np.ndarray((out_H, out_W, in_C), dtype='float32')
    kernel_width = weights_W.shape[1]
    for i in range(out_W):
        idx = int(indices_W[i][0])
        out_2[:, i, 0] = np.dot(
            out_1_aug[:, idx:idx + kernel_width, 0], weights_W[i])
        out_2[:, i, 1] = np.dot(
            out_1_aug[:, idx:idx + kernel_width, 1], weights_W[i])
        out_2[:, i, 2] = np.dot(
            out_1_aug[:, idx:idx + kernel_width, 2], weights_W[i])

    return out_2


# Command line arguments
data_parser = argparse.ArgumentParser(description='dataset path')

data_parser.add_argument('--hr-hdr-path', type=str, default="HDR_youtube_80.mat",
                         help='add path where HR-HDR dataset resides')

data_parser.add_argument('--lr-hdr-path', type=str, default="LR_HDR_youtube_80.mat",
                         help='where the LR-HDR images will be saved')

data_parser.add_argument('--reduction-factor', type=int, default=2,
                         help='by how much factor you want to reduce the size of the images')

data_parser.add_argument('--create-test-dataset', action='store_true', default=False,
                         help='If True, HDR_YUV is name of group in hdf5 dataset')

args = data_parser.parse_args()

if __name__ == "__main__":
    f_hr_hdr = h5py.File((args.hr_hdr_path))
    if args.create_test_dataset:
        imgs = f_hr_hdr['HDR_YUV']
    else:
        imgs = f_hr_hdr['HDR_data']
    img_lr_hdr_data = []

    for idx in range(len(imgs)):
        print("Image-", idx)
        img = imgs[idx]
        img_s = np.swapaxes(img, 0, 2)
        # Create LR images
        img_lr_hdr = imresize_np(img_s, 1 / args.reduction_factor, True)
        img_lr_hdr = np.swapaxes(img_lr_hdr, 0, 2)
        img_lr_hdr_data.append(img_lr_hdr)

    img_lr_hdr_data = np.array(img_lr_hdr_data)

    if args.create_test_dataset:
        with h5py.File(args.lr_hdr_path, "w") as f_lr_hdr:
            dset = f_lr_hdr.create_dataset('HDR_YUV', data=img_lr_hdr_data)
    else:
        with h5py.File(args.lr_hdr_path, "w") as f_lr_hdr:
            dset = f_lr_hdr.create_dataset('HDR_data', data=img_lr_hdr_data)

    f_hr_hdr.close()


# video-superresolution/jsigan/inference.py

# Copyright 2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import time
import os
import nnabla as nn
import nnabla.functions as F
from nnabla.ext_utils import get_extension_context
import numpy as np
from utils import get_hw_boundary, trim_patch_boundary, compute_psnr, save_results_yuv
from args import get_config
from ops import model
from data_loader import read_mat_file


def inference():
    """
    Inference function to generate high resolution hdr images
    """
    conf = get_config()
    ctx = get_extension_context(
        conf.nnabla_context.context, device_id=conf.nnabla_context.device_id)
    nn.set_default_context(ctx)

    data, target = read_mat_file(conf.data.lr_sdr_test, conf.data.hr_hdr_test, conf.data.d_name_test,
                                 conf.data.l_name_test, train=False)

    if not os.path.exists(conf.test_img_dir):
        os.makedirs(conf.test_img_dir)

    data_sz = data.shape
    target_sz = target.shape
    PATCH_BOUNDARY = 10  # set patch boundary to reduce edge effect around patch edges
    test_loss_PSNR_list_for_epoch = []
    inf_time = []
    start_time = time.time()

    test_pred_full = np.zeros((target_sz[1], target_sz[2], target_sz[3]))

    print("Loading pre trained model.........", conf.pre_trained_model)
    nn.load_parameters(conf.pre_trained_model)

    for index in range(data_sz[0]):
        ### ======== Divide Into Patches ======== ###
        for p in range(conf.test_patch ** 2):
            pH = p // conf.test_patch
            pW = p % conf.test_patch
            sH = data_sz[1] // conf.test_patch
            sW = data_sz[2] // conf.test_patch
            H_low_ind, H_high_ind, W_low_ind, W_high_ind = \
                get_hw_boundary(
                    PATCH_BOUNDARY, data_sz[1], data_sz[2], pH, sH, pW, sW)
            data_test_p = nn.Variable.from_numpy_array(
                data.d[index, H_low_ind: H_high_ind, W_low_ind: W_high_ind, :])
            data_test_sz = data_test_p.shape
            data_test_p = F.reshape(
                data_test_p, (1, data_test_sz[0], data_test_sz[1], data_test_sz[2]))
            st = time.time()
            net = model(data_test_p, conf.scaling_factor)
            net.pred.forward()
            test_pred_temp = net.pred.d
            inf_time.append(time.time() - st)
            test_pred_t = trim_patch_boundary(test_pred_temp, PATCH_BOUNDARY, data_sz[1], data_sz[2],
                                              pH, sH, pW, sW, conf.scaling_factor)
            # pred_sz = test_pred_t.shape
            test_pred_t = np.squeeze(test_pred_t)
            test_pred_full[pH * sH * conf.scaling_factor: (pH + 1) * sH * conf.scaling_factor,
                           pW * sW * conf.scaling_factor: (pW + 1) * sW * conf.scaling_factor, :] = test_pred_t

        ### ======== Compute PSNR & Print Results======== ###
        test_GT = np.squeeze(target.d[index, :, :, :])
        test_PSNR = compute_psnr(test_pred_full, test_GT, 1.)
        test_loss_PSNR_list_for_epoch.append(test_PSNR)
        print(" <Test> [%4d/%4d]-th images, time: %4.4f(minutes), test_PSNR: %.8f[dB]  "
              % (int(index), int(data_sz[0]), (time.time() - start_time) / 60, test_PSNR))
        if conf.save_images:
            # comment for faster testing
            save_results_yuv(test_pred_full, index, conf.test_img_dir)
    test_PSNR_per_epoch = np.mean(test_loss_PSNR_list_for_epoch)

    print("######### Average Test PSNR: %.8f[dB]  #########" % (
        test_PSNR_per_epoch))
    print("######### Estimated Inference Time (per 4K frame): %.8f[s]  #########" %
          (np.mean(inf_time) * conf.test_patch * conf.test_patch))


if __name__ == '__main__':
    inference()


# video-superresolution/jsigan/ops.py

# Copyright 2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections import namedtuple
import nnabla as nn
import nnabla.functions as F
import nnabla.parametric_functions as PF
import nnabla.initializer as I
import numpy as np
from utils import depth_to_space


def box_filter(x, szf):
    """
    Box filter
    """
    y = F.identity(x)
    szy = list(y.shape)
    b_filt = nn.Variable((szf, szf, 1, 1))
    b_filt.data.fill(1.)
    b_filt = b_filt / (szf ** 2)
    # 5,5,1,1
    b_filt = F.tile(b_filt, [1, 1, szy[3], 1])
    b_filt = F.transpose(b_filt, (3, 2, 0, 1))
    b_filt = F.reshape(b_filt, (6, 5, 5))
    pp = int((szf - 1) / 2)
    y = F.pad(y, (0, 0, pp, pp, pp, pp, 0, 0), mode='reflect')
    y_chw = F.transpose(y, (0, 3, 1, 2))
    y_chw = F.depthwise_convolution(y_chw, b_filt, multiplier=1, stride=(1, 1))
    y_hwc = F.transpose(y_chw, (0, 2, 3, 1))
    return y_hwc


def guided_filter(img, r, eps):
    """
    Edge preserving filter
    """
    img2 = F.concatenate(img, img * img, axis=3)
    img2 = box_filter(img2, r)
    mean = F.split(img2, axis=3)
    mean_i = F.stack(mean[0], mean[1], mean[2], axis=3)
    mean_ii = F.stack(mean[3], mean[4], mean[5], axis=3)
    var_i = mean_ii - mean_i * mean_i
    a = var_i / (var_i + eps)
    b = mean_i - a * mean_i
    ab = F.concatenate(a, b, axis=3)
    ab = box_filter(ab, r)
    mean_ab = F.split(ab, axis=3)
    mean_a = F.stack(mean_ab[0], mean_ab[1], mean_ab[2], axis=3)
    mean_b = F.stack(mean_ab[3], mean_ab[4], mean_ab[5], axis=3)
    q = mean_a * img + mean_b
    return q


def conv_2d(x, o_ch, kernel, name=None):
    """
    Convolution for JSInet
    """
    b = I.ConstantInitializer(0.)
    h = PF.convolution(x, o_ch, kernel=kernel, stride=(1, 1), pad=(1, 1), channel_last=True,
                       b_init=b, name=name)
    return h


def res_block(x, out_ch, name):
    """
    Create residual block
    """
    with nn.parameter_scope(name):
        h = conv_2d(F.relu(x), out_ch, kernel=(3, 3), name='conv/0')
        h = conv_2d(F.relu(h), out_ch, kernel=(3, 3), name='conv/1')
        h = x + h
    return h


def dyn_2d_filter(x, lf_2d, k_sz):
    """
    Dynamic 2d filtering
    """
    with nn.parameter_scope('Dynamic_2D_Filtering'):
        f_localexpand = nn.Variable.from_numpy_array(
            np.eye(k_sz[0] * k_sz[1], k_sz[0] * k_sz[1]))
        f_localexpand = F.reshape(f_localexpand,
                                  (k_sz[0], k_sz[1], 1, k_sz[0] * k_sz[1]))  # (9,9,1,81))
        f_localexpand = F.transpose(f_localexpand, (3, 0, 1, 2))  # (81,9,9,1))
        x_sz = x.shape
        x = F.reshape(x, (x_sz[0], x_sz[1], x_sz[2], 1))  # (1,100,170,1)
        x_localexpand = F.convolution(x, f_localexpand, stride=(1, 1), pad=(4, 4),
                                      channel_last=True)  # (1,100,170,81)
        x_le_sz = x_localexpand.shape
        x_localexpand = F.reshape(x_localexpand,
                                  (x_le_sz[0], x_le_sz[1], x_le_sz[2], 1, x_le_sz[3]))
        y = F.batch_matmul(x_localexpand, lf_2d)
        y_sz = y.shape
        y = F.reshape(y, (y_sz[0], y_sz[1], y_sz[2], y_sz[4]))
    return y


def dyn_2d_up_operation(x, lf_2d, k_sz, sf=2):
    """
    Dynamic 2d upsampling
    """
    with nn.parameter_scope("Dynamic_2D_Upsampling"):
        y = []
        sz = lf_2d.shape
        lf_2d_new = F.reshape(
            lf_2d, (sz[0], sz[1], sz[2], k_sz[0] * k_sz[0], sf ** 2))
        lf_2d_new = F.softmax(lf_2d_new, axis=3)
        for ch in range(3):  # loop over YUV channels
            # apply dynamic filtering operation
            temp = dyn_2d_filter(x[:, :, :, ch], lf_2d_new, k_sz)
            temp = depth_to_space(temp, sf)
            y += [temp]
        y = F.concatenate(*y, axis=3)
    return y


def dyn_sep_up_operation(x, dr_k_v, dr_k_h, k_sz, sf):
    """
    Dynamic separable upsampling operation with 1D separable local kernels.
    x: [B, H, W, C], dr_k_v: [B, H, W, 41*sf*sf], dr_k_h: [B, H, W, 41*sf*sf]
    out: [B, H*sf, W*sf, C]
    """
    sz = x.shape
    pad = k_sz // 2  # local filter pad size
    # [B, H, W, C*sf*sf]
    out_v = nn.Variable((sz[0], sz[1], sz[2], sz[3] * sf ** 2))
    out_v.data.zero()
    # [B, H, W, C*sf*sf]
    out_h = nn.Variable((sz[0], sz[1], sz[2], sz[3] * sf ** 2))
    out_h.data.zero()
    img_pad = F.pad(x, (0, 0, pad, pad, 0, 0, 0, 0))
    img_pad_y = F.reshape(img_pad[:, :, :, 0],
                          (img_pad.shape[0], img_pad.shape[1], img_pad.shape[2], 1))
    img_pad_y = F.tile(img_pad_y, [1, 1, 1, sf ** 2])
    img_pad_u = F.reshape(img_pad[:, :, :, 1],
                          (img_pad.shape[0], img_pad.shape[1], img_pad.shape[2], 1))
    img_pad_u = F.tile(img_pad_u, [1, 1, 1, sf ** 2])
    img_pad_v = F.reshape(img_pad[:, :, :, 2],
                          (img_pad.shape[0], img_pad.shape[1], img_pad.shape[2], 1))
    img_pad_v = F.tile(img_pad_v, [1, 1, 1, sf ** 2])
    img_pad = F.concatenate(img_pad_y, img_pad_u, img_pad_v, axis=3)

    # vertical 1D filter
    for i in range(k_sz):
        out_v = out_v + img_pad[:, i:i + sz[1], :, :] * F.tile(
            dr_k_v[:, :, :, i:k_sz * sf ** 2:k_sz], [1, 1, 1, 3])
    img_pad = F.pad(out_v, (0, 0, 0, 0, pad, pad, 0, 0))
    # horizontal 1D filter
    for i in range(k_sz):
        out_h = out_h + img_pad[:, :, i:i + sz[2], :] * F.tile(
            dr_k_h[:, :, :, i:k_sz * sf ** 2:k_sz], [1, 1, 1, 3])

    # depth to space upsampling (YUV)
    out = depth_to_space(out_h[:, :, :, 0:sf ** 2], sf)
    out = F.concatenate(out, depth_to_space(
        out_h[:, :, :, sf ** 2:2 * sf ** 2], sf), axis=3)
    out = F.concatenate(out, depth_to_space(
        out_h[:, :, :, 2 * sf ** 2:3 * sf ** 2], sf), axis=3)
    return out


def res_block_concat(x, out_ch, name):
    """
    Basic residual block -> [conv-relu | conv-relu] + input
    """
    with nn.parameter_scope(name):
        h = conv_2d(F.relu(x), out_ch, kernel=(3, 3), name='conv/0')
        h = conv_2d(F.relu(h), out_ch, kernel=(3, 3), name='conv/1')
        h = x[:, :, :, :out_ch] + h
    return h


def model(img, sf):
    """
    Define JSInet model
    """
    with nn.parameter_scope('Network'):
        with nn.parameter_scope('local_contrast_enhancement'):
            ## ================= Local Contrast Enhancement Subnet ============================ ##
            ch = 64
            b = guided_filter(img, 5, 0.01)
            n1 = conv_2d(b, ch, kernel=(3, 3), name='conv/0')
            for i in range(4):
                n1 = res_block(n1, ch, 'res_block/%d' % i)
            n1 = F.relu(n1)
            local_filter_2d = conv_2d(n1, (9 ** 2) * (sf ** 2), kernel=(3, 3),
                                      name='conv_k')  # [B, H, W, (9x9)*(sfxsf)]
            # dynamic 2D upsampling with 2D local filters
            pred_C = dyn_2d_up_operation(b, local_filter_2d, (9, 9), sf)
            # local contrast mask
            pred_C = 2 * F.sigmoid(pred_C)
            ## ================= Detail Restoration Subnet ============================ ##
            ch = 64
            d = F.div2(img, b + 1e-15)
        with nn.parameter_scope('detail_restoration'):
            n3 = conv_2d(d, ch, kernel=(3, 3), name='conv/0')
            for i in range(4):
                n3 = res_block(n3, ch, 'res_block/%d' % i)
                if i == 0:
                    d_feature = n3
            n3 = F.relu(n3)
            # separable 1D filters
            dr_k_h = conv_2d(n3, 41 * sf ** 2, kernel=(3, 3), name='conv_k_h')
            dr_k_v = conv_2d(n3, 41 * sf ** 2, kernel=(3, 3), name='conv_k_v')
            # dynamic separable upsampling with with separable 1D local filters
            pred_D = dyn_sep_up_operation(d, dr_k_v, dr_k_h, 41, sf)
        ## ================= Image Reconstruction Subnet ============================ ##
        with nn.parameter_scope('image_reconstruction'):
            n4 = conv_2d(img, ch, kernel=(3, 3), name='conv/0')
            for i in range(4):
                if i == 1:
                    n4 = F.concatenate(n4, d_feature, axis=3)
                    n4 = res_block_concat(n4, ch, 'res_block/%d' % i)
                else:
                    n4 = res_block(n4, ch, 'res_block/%d' % i)
            n4 = F.relu(n4)

            n4 = F.relu(conv_2d(n4, ch * sf * sf, kernel=(3, 3),
                                name='conv/1'))
            # (1,100,170,1024) -> (1,100,170,4,4,64) -> (1,100,4,170,4,64)
            # pixel shuffle
            n4 = depth_to_space(n4, sf)
            pred_I = conv_2d(n4, 3, kernel=(3, 3), name='conv/2')

    pred = F.add2(pred_I, pred_D) * pred_C
    jsinet = namedtuple('jsinet', ['pred'])
    return jsinet(pred)


def truncated_normal(w_shape, mean, std):
    """
    Numpy truncated normal
    """
    init = I.NormalInitializer()
    tmp = init(w_shape + (4,))
    valid = np.logical_and((np.less(tmp, 2)), (np.greater(tmp, -2)))
    ind = np.argmax(valid, axis=-1)
    ind1 = (np.expand_dims(ind, -1))
    trunc_norm = np.take_along_axis(tmp, ind1, axis=4).squeeze(-1)
    trunc_norm = trunc_norm * std + mean
    return trunc_norm


def conv(x, channels, kernel=4, stride=2, pad=0, pad_type='zero', use_bias=True, scope='conv_0'):
    """
    Convolution for discriminator
    """
    w_n_shape = (channels, kernel, kernel, x.shape[-1])
    w_init = truncated_normal(w_n_shape, mean=0.0, std=0.02)
    b_init = I.ConstantInitializer(0.)
    with nn.parameter_scope(scope):
        if pad > 0:
            h = x.shape[1]
            if h % stride == 0:
                pad = pad * 2
            else:
                pad = max(kernel - (h % stride), 0)
            pad_top = pad // 2
            pad_bottom = pad - pad_top
            pad_left = pad // 2
            pad_right = pad - pad_left

            if pad_type == 'zero':
                x = F.pad(x, (0, 0, pad_top, pad_bottom,
                              pad_left, pad_right, 0, 0))
            if pad_type == 'reflect':
                x = F.pad(x, (0, 0, pad_top, pad_bottom, pad_left,
                              pad_right, 0, 0), mode='reflect')

        def apply_w(w):
            return PF.spectral_norm(w, dim=0)

        x = PF.convolution(x, channels, kernel=(kernel, kernel), stride=(
            stride, stride), apply_w=apply_w, w_init=w_init, b_init=b_init, with_bias=use_bias,
                           channel_last=True)

    return x


def dis_block(n, c, i, train=True):
    """
    Discriminator conv_bn_relu block
    """
    out = conv(n, channels=c, kernel=4, stride=2, pad=1, use_bias=False,
               scope='d_conv/' + str(2 * i + 2))
    out_fm = F.leaky_relu(
        PF.batch_normalization(
            out, axes=[3], batch_stat=train, name='d_bn/' + str(2 * i + 1)),
        alpha=0.2)

    out = conv(out_fm, channels=c * 2, kernel=3, stride=1, pad=1, use_bias=False,
               scope='d_conv/' + str(2 * i + 3))
    out = F.leaky_relu(
        PF.batch_normalization(
            out, axes=[3], batch_stat=train, name='d_bn/' + str(2 * i + 2)),
        alpha=0.2)
    return out, out_fm


def discriminator_fm(x, sf, scope="Discriminator_FM"):
    """
    Feature matching discriminator
    """
    with nn.parameter_scope(scope):
        fm_list = []
        ch = 32
        n = F.leaky_relu(conv(x, ch, 3, 1, 1, scope='d_conv/1'), alpha=0.2)
        for i in range(4):
            n, out_fm = dis_block(n, ch, i, train=True)
            ch = ch * 2
            fm_list.append(out_fm)
        n = F.leaky_relu(PF.batch_normalization(
            conv(n, channels=ch, kernel=4, stride=2,
                 pad=1, use_bias=False, scope='d_conv/10'),
            axes=[3], batch_stat=True, name='d_bn/9'), alpha=0.2
            )

        if sf == 1:
            n = F.leaky_relu(PF.batch_normalization(
                conv(n, channels=ch, kernel=5, stride=1,
                     pad=1, use_bias=False, scope='d_conv/11'),
                axes=[3], batch_stat=True, name='d_bn/10'), alpha=0.2)
        else:
            n = F.leaky_relu(PF.batch_normalization(
                conv(n, channels=ch, kernel=5, stride=1,
                     use_bias=False, scope='d_conv/11'),
                axes=[3], batch_stat=True, name='d_bn/10'), alpha=0.2)

        n = PF.batch_normalization(
            conv(n, channels=1, kernel=1, stride=1,
                 use_bias=False, scope='d_conv/12'),
            axes=[3], batch_stat=True, name='d_bn/11')

        out_logit = n
        out = F.sigmoid(out_logit)  # [B,1]
        return out, out_logit, fm_list


def discriminator_loss(real, fake):
    """
    Calculate discriminator loss
    """
    real_loss = F.mean(
        F.relu(1.0 - (real - F.reshape(F.mean(fake), (1, 1, 1, 1)))))
    fake_loss = F.mean(
        F.relu(1.0 + (fake - F.reshape(F.mean(real), (1, 1, 1, 1)))))
    l_d = real_loss + fake_loss
    return l_d


def generator_loss(real, fake):
    """
    Calculate generator loss
    """
    real_loss = F.mean(
        F.relu(1.0 + (real - F.reshape(F.mean(fake), (1, 1, 1, 1)))))
    fake_loss = F.mean(
        F.relu(1.0 - (fake - F.reshape(F.mean(real), (1, 1, 1, 1)))))
    l_g = real_loss + fake_loss
    return l_g


def feature_matching_loss(x, y, num=4):
    """
    Calculate feature matching loss
    """
    fm_loss = 0.0
    for i in range(num):
        fm_loss += F.mean(F.squared_error(x[i], y[i]))
    return fm_loss


def gan_model(label_ph, pred, conf):
    """
    Define GAN model with adversarial and discriminator losses and their orchestration
    """
    # Define Discriminator
    _, d_real_logits, d_real_fm_list = discriminator_fm(
        label_ph, conf.scaling_factor, scope="Discriminator_FM")
    # output of D for fake images
    _, d_fake_logits, d_fake_fm_list = discriminator_fm(
        pred, conf.scaling_factor, scope="Discriminator_FM")

    # Define Detail Discriminator
    # compute the detail layers for the dicriminator (reuse)
    base_gt = guided_filter(label_ph, 5, 0.01)
    detail_gt = F.div2(label_ph, base_gt + 1e-15)
    base_pred = guided_filter(pred, 5, 0.01)
    detail_pred = F.div2(pred, base_pred + 1e-15)

    # detail layer output of D for real images
    _, d_detail_real_logits, d_detail_real_fm_list = \
        discriminator_fm(detail_gt, conf.scaling_factor,
                         scope="Discriminator_Detail")

    # detail layer output of D for fake images
    _, d_detail_fake_logits, d_detail_fake_fm_list = \
        discriminator_fm(detail_pred, conf.scaling_factor,
                         scope="Discriminator_Detail")

    # Loss
    # original GAN (hinge GAN)
    d_adv_loss = discriminator_loss(d_real_logits, d_fake_logits)
    d_adv_loss.persistent = True
    g_adv_loss = generator_loss(d_real_logits, d_fake_logits)
    g_adv_loss.persistent = True

    # detail GAN (hinge GAN)
    d_detail_adv_loss = conf.detail_lambda * \
        discriminator_loss(d_detail_real_logits, d_detail_fake_logits)
    d_detail_adv_loss.persistent = True
    g_detail_adv_loss = conf.detail_lambda * \
        generator_loss(d_detail_real_logits, d_detail_fake_logits)
    g_detail_adv_loss.persistent = True

    # feature matching (FM) loss
    fm_loss = feature_matching_loss(d_real_fm_list, d_fake_fm_list, 4)
    fm_loss.persistent = True
    fm_detail_loss = conf.detail_lambda * feature_matching_loss(d_detail_real_fm_list,
                                                                d_detail_fake_fm_list, 4)
    fm_detail_loss.persistent = True

    jsigan = namedtuple('jsigan',
                        ['d_adv_loss', 'd_detail_adv_loss', 'g_adv_loss', 'g_detail_adv_loss',
                         'fm_loss', 'fm_detail_loss'])
    return jsigan(d_adv_loss, d_detail_adv_loss, g_adv_loss, g_detail_adv_loss, fm_loss,
                  fm_detail_loss)


# video-superresolution/jsigan/train.py

# Copyright 2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import datetime
import os
import nnabla as nn
import nnabla.functions as F
import nnabla.solvers as S
from nnabla.monitor import Monitor, MonitorSeries, MonitorTimeElapsed
from nnabla.ext_utils import get_extension_context
from args import get_config
from ops import model, gan_model
from data_loader import jsi_iterator
from utils import compute_psnr, get_learning_rate, CommunicatorWrapper


def setup_monitor(conf, monitor):
    """
    Setup monitor to keep track of losses and times to log them
    """
    jsi_monitor = {'rec_loss': MonitorSeries('rec_loss', monitor, interval=conf.monitor_interval),
                   'psnr': MonitorSeries('psnr', monitor, interval=conf.monitor_interval),
                   'lr': MonitorSeries('learning_rate', monitor,
                                       interval=conf.monitor_interval),
                   'g_final_loss': MonitorSeries('g_final_loss', monitor,
                                                 interval=conf.monitor_interval),
                   'd_final_fm_loss': MonitorSeries('d_final_fm_loss', monitor,
                                                    interval=conf.monitor_interval),
                   'd_final_detail_loss': MonitorSeries('d_final_detail_loss', monitor,
                                                        interval=conf.monitor_interval),
                   'g_adv_loss': MonitorSeries('g_adv_loss', monitor,
                                               interval=conf.monitor_interval),
                   'g_detail_adv_loss': MonitorSeries('g_detail_adv_loss', monitor,
                                                      interval=conf.monitor_interval),
                   'fm_loss': MonitorSeries('fm_loss', monitor,
                                            interval=conf.monitor_interval),
                   'fm_detail_loss': MonitorSeries('fm_detail_loss', monitor,
                                                   interval=conf.monitor_interval),
                   'time': MonitorTimeElapsed("Training time per epoch", monitor,
                                              interval=conf.monitor_interval)}
    return jsi_monitor


def main():
    conf = get_config()
    extension_module = conf.nnabla_context.context
    ctx = get_extension_context(
        extension_module, device_id=conf.nnabla_context.device_id)
    comm = CommunicatorWrapper(ctx)
    nn.set_default_context(comm.ctx)
    print("#GPU Count: ", comm.n_procs)

    data_iterator_train = jsi_iterator(conf.batch_size, conf, train=True)
    if conf.scaling_factor == 1:
        d_t = nn.Variable((conf.batch_size, 80, 80, 3),
                          need_grad=True)
        l_t = nn.Variable((conf.batch_size, 80, 80, 3), need_grad=True)

    else:
        d_t = nn.Variable((conf.batch_size, 160 / conf.scaling_factor, 160 / conf.scaling_factor, 3),
                          need_grad=True)
        l_t = nn.Variable((conf.batch_size, 160, 160, 3), need_grad=True)

    if comm.n_procs > 1:
        data_iterator_train = data_iterator_train.slice(
            rng=None, num_of_slices=comm.n_procs, slice_pos=comm.rank)

    monitor_path = './nnmonitor' + \
        str(datetime.datetime.now().strftime("%Y%m%d%H%M%S"))

    monitor = Monitor(monitor_path)
    jsi_monitor = setup_monitor(conf, monitor)

    with nn.parameter_scope("jsinet"):
        nn.load_parameters(conf.pre_trained_model)
        net = model(d_t, conf.scaling_factor)
        net.pred.persistent = True
    rec_loss = F.mean(F.squared_error(net.pred, l_t))
    rec_loss.persistent = True
    g_final_loss = rec_loss

    if conf.jsigan:
        net_gan = gan_model(l_t, net.pred, conf)
        d_final_fm_loss = net_gan.d_adv_loss
        d_final_fm_loss.persistent = True
        d_final_detail_loss = net_gan.d_detail_adv_loss
        d_final_detail_loss.persistent = True
        g_final_loss = conf.rec_lambda * rec_loss + conf.adv_lambda * (
                net_gan.g_adv_loss + net_gan.g_detail_adv_loss) + conf.fm_lambda * (
                               net_gan.fm_loss + net_gan.fm_detail_loss)
        g_final_loss.persistent = True

    max_iter = data_iterator_train._size // (conf.batch_size)
    if comm.rank == 0:
        print("max_iter", data_iterator_train._size, max_iter)

    iteration = 0
    if not conf.jsigan:
        start_epoch = 0
        end_epoch = conf.adv_weight_point
        lr = conf.learning_rate * comm.n_procs
    else:
        start_epoch = conf.adv_weight_point
        end_epoch = conf.epoch
        lr = conf.learning_rate * comm.n_procs
        w_d = conf.weight_decay * comm.n_procs

    # Set generator parameters
    with nn.parameter_scope("jsinet"):
        solver_jsinet = S.Adam(alpha=lr, beta1=0.9, beta2=0.999, eps=1e-08)
        solver_jsinet.set_parameters(nn.get_parameters())

    if conf.jsigan:
        solver_disc_fm = S.Adam(alpha=lr, beta1=0.9, beta2=0.999, eps=1e-08)
        solver_disc_detail = S.Adam(
            alpha=lr, beta1=0.9, beta2=0.999, eps=1e-08)
        with nn.parameter_scope("Discriminator_FM"):
            solver_disc_fm.set_parameters(nn.get_parameters())
        with nn.parameter_scope("Discriminator_Detail"):
            solver_disc_detail.set_parameters(nn.get_parameters())

    for epoch in range(start_epoch, end_epoch):
        for index in range(max_iter):
            d_t.d, l_t.d = data_iterator_train.next()

            if not conf.jsigan:
                # JSI-net -> Generator
                lr_stair_decay_points = [200, 225]
                lr_net = get_learning_rate(lr, iteration, lr_stair_decay_points,
                                           conf.lr_decreasing_factor)
                g_final_loss.forward(clear_no_need_grad=True)
                solver_jsinet.zero_grad()
                if comm.n_procs > 1:
                    all_reduce_callback = comm.get_all_reduce_callback()
                    g_final_loss.backward(clear_buffer=True,
                                          communicator_callbacks=all_reduce_callback)
                else:
                    g_final_loss.backward(clear_buffer=True)
                solver_jsinet.set_learning_rate(lr_net)
                solver_jsinet.update()
            else:
                # GAN part (discriminator + generator)
                lr_gan = lr if epoch < conf.gan_lr_linear_decay_point \
                    else lr * (end_epoch - epoch) / (end_epoch - conf.gan_lr_linear_decay_point)
                lr_gan = lr_gan * conf.gan_ratio

                net.pred.need_grad = False

                # Discriminator_FM
                solver_disc_fm.zero_grad()
                d_final_fm_loss.forward(clear_no_need_grad=True)
                if comm.n_procs > 1:
                    all_reduce_callback = comm.get_all_reduce_callback()
                    d_final_fm_loss.backward(clear_buffer=True,
                                             communicator_callbacks=all_reduce_callback)
                else:
                    d_final_fm_loss.backward(clear_buffer=True)
                solver_disc_fm.set_learning_rate(lr_gan)
                solver_disc_fm.weight_decay(w_d)
                solver_disc_fm.update()

                # Discriminator_Detail
                solver_disc_detail.zero_grad()
                d_final_detail_loss.forward(clear_no_need_grad=True)
                if comm.n_procs > 1:
                    all_reduce_callback = comm.get_all_reduce_callback()
                    d_final_detail_loss.backward(clear_buffer=True,
                                                 communicator_callbacks=all_reduce_callback)
                else:
                    d_final_detail_loss.backward(clear_buffer=True)
                solver_disc_detail.set_learning_rate(lr_gan)
                solver_disc_detail.weight_decay(w_d)
                solver_disc_detail.update()

                # Generator
                net.pred.need_grad = True
                solver_jsinet.zero_grad()
                g_final_loss.forward(clear_no_need_grad=True)
                if comm.n_procs > 1:
                    all_reduce_callback = comm.get_all_reduce_callback()
                    g_final_loss.backward(clear_buffer=True,
                                          communicator_callbacks=all_reduce_callback)
                else:
                    g_final_loss.backward(clear_buffer=True)
                solver_jsinet.set_learning_rate(lr_gan)
                solver_jsinet.update()

            iteration += 1
            if comm.rank == 0:
                train_psnr = compute_psnr(net.pred.d, l_t.d, 1.)
                jsi_monitor['psnr'].add(iteration, train_psnr)
                jsi_monitor['rec_loss'].add(iteration, rec_loss.d.copy())
                jsi_monitor['time'].add(iteration)

            if comm.rank == 0:
                if conf.jsigan:
                    jsi_monitor['g_final_loss'].add(
                        iteration, g_final_loss.d.copy())
                    jsi_monitor['g_adv_loss'].add(
                        iteration, net_gan.g_adv_loss.d.copy())
                    jsi_monitor['g_detail_adv_loss'].add(iteration,
                                                         net_gan.g_detail_adv_loss.d.copy())
                    jsi_monitor['d_final_fm_loss'].add(
                        iteration, d_final_fm_loss.d.copy())
                    jsi_monitor['d_final_detail_loss'].add(
                        iteration, d_final_detail_loss.d.copy())
                    jsi_monitor['fm_loss'].add(
                        iteration, net_gan.fm_loss.d.copy())
                    jsi_monitor['fm_detail_loss'].add(
                        iteration, net_gan.fm_detail_loss.d.copy())
                    jsi_monitor['lr'].add(iteration, lr_gan)

        if comm.rank == 0:
            if not os.path.exists(conf.output_dir):
                os.makedirs(conf.output_dir)
            with nn.parameter_scope("jsinet"):
                nn.save_parameters(os.path.join(
                    conf.output_dir, "model_param_%04d.h5" % epoch))


if __name__ == "__main__":
    main()


# video-superresolution/jsigan/utils.py

# Copyright 2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys
import imageio
import numpy as np
import nnabla.functions as F

common_utils_path = os.path.abspath(
    os.path.join(os.path.dirname(__file__), '..', '..', 'utils'))
sys.path.append(common_utils_path)

from neu.yaml_wrapper import read_yaml
from neu.misc import AttrDict
from neu.comm import CommunicatorWrapper
from neu.variable_utils import set_persistent_all


def get_hw_boundary(patch_boundary, h, w, pH, sH, pW, sW):
    """
    Calculate height and width of patch
    """
    h_low_ind = max(pH * sH - patch_boundary, 0)
    h_high_ind = min((pH + 1) * sH + patch_boundary, h)
    w_low_ind = max(pW * sW - patch_boundary, 0)
    w_high_ind = min((pW + 1) * sW + patch_boundary, w)

    return h_low_ind, h_high_ind, w_low_ind, w_high_ind


def depth_to_space(x, sf):
    """
    x: (B, H, W, sf*sf)
    reshape x to (B, H, W, sf,sf,1) then transpose to (B, H, sf, W, sf, 1) & finally reshape (B, H*sf, W*sf, 1)
    return: (B, H*sf, W*sf, 1)
    """
    x_sz = x.shape
    if x_sz[3] == sf*sf:
        x = F.reshape(x, (x_sz[0], x_sz[1], x_sz[2], sf, sf, 1))
        x = F.reshape(F.transpose(x, (0, 1, 3, 2, 4, 5)),
                      (x_sz[0], x_sz[1] * sf, x_sz[2] * sf, 1))
    else:
        ch = 64
        x = F.reshape(x, (x_sz[0], x_sz[1], x_sz[2], sf, sf, ch))
        x = F.reshape(F.transpose(x, (0, 1, 3, 2, 4, 5)),
                      (x_sz[0], x_sz[1] * sf, x_sz[2] * sf, ch))
    return x


def trim_patch_boundary(img, patch_boundary, h, w, pH, sH, pW, sW, sf):
    """
    Remove both rows and columns to reduce edge effect 
    around patch edges.
    """
    # trim rows
    if pH * sH >= patch_boundary:
        img = img[:, patch_boundary * sf:, :, :]
    if (pH + 1) * sH + patch_boundary <= h:
        img = img[:, :-patch_boundary * sf, :, :]

    # trim columns
    if pW * sW >= patch_boundary:
        img = img[:, :, patch_boundary * sf:, :]
    if (pW + 1) * sW + patch_boundary <= w:
        img = img[:, :, :-patch_boundary * sf, :]

    return img


def compute_psnr(img_orig, img_out, peak):
    """
    Calculate PSNR value
    """
    mse = np.mean((img_orig - img_out) ** 2)
    psnr = 10 * np.log10(1.*1. / mse)
    return psnr


def get_learning_rate(init_lr, iteration, lr_stair_decay_points, lr_decreasing_factor):
    """
    Calculate learning rate decay
    """
    epoch_lr_to_be_decayed_boundaries = [
        y * (iteration) for y in lr_stair_decay_points]
    epoch_lr_to_be_decayed_value = [init_lr * (lr_decreasing_factor ** y) for y in
                                    range(len(lr_stair_decay_points) + 1)]
    lr = init_lr
    if epoch_lr_to_be_decayed_boundaries[0] < iteration <= epoch_lr_to_be_decayed_boundaries[1]:
        lr = epoch_lr_to_be_decayed_value[1]
    if iteration > epoch_lr_to_be_decayed_boundaries[1]:
        lr = epoch_lr_to_be_decayed_value[2]

    return lr


def save_results_yuv(pred, index, test_img_dir):
    """
    Saves generated YUV image seperately by channel  
    """
    test_pred = np.squeeze(pred)
    test_pred = np.clip(test_pred, 0, 1) * 1023
    test_pred = np.uint16(test_pred)

    # split image
    pred_y = test_pred[:, :, 0]
    pred_u = test_pred[:, :, 1]
    pred_v = test_pred[:, :, 2]

    # save prediction - must be saved in separate channels due to 16-bit pixel depth
    imageio.imwrite(os.path.join(test_img_dir, "{}-y_pred.png".format(str(int(index) + 1).zfill(2))),
                    pred_y)
    imageio.imwrite(os.path.join(test_img_dir, "{}-u_pred.png".format(str(int(index) + 1).zfill(2))),
                    pred_u)
    imageio.imwrite(os.path.join(test_img_dir, "{}-v_pred.png".format(str(int(index) + 1).zfill(2))),
                    pred_v)


# video-superresolution/tecogan/args.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse

from utils import read_yaml


def get_config():
    """
    Get command line arguments.
    Arguments set the default values of command line arguments.
    """

    parser = argparse.ArgumentParser(description='TecoGAN')

    parser.add_argument('--cfg', default="./config.yaml")

    args, _ = parser.parse_known_args()

    conf = read_yaml(args.cfg)

    parser.add_argument('--input_video_dir', default=conf.data.input_video_dir,
                        help='The directory of the video input data, for training')
    parser.add_argument('--output_dir', default=conf.data.output_dir,
                        help='The output directory of the checkpoint')
    parser.add_argument('--num_resblock', type=int, default=conf.train.num_resblock,
                        help='How many residual blocks are there in the generator')
    parser.add_argument('--max_iter', type=int, default=conf.train.max_iter,
                        help='max iteration for training')
    parser.add_argument('--pre_trained_frvsr_weights', type=str, default=conf.train.pre_trained_frvsr_weights,
                        help='the weight of frvsr generator will be loaded as an initial point')
    parser.add_argument('--vgg_pre_trained_weights', type=str,
                        default=conf.train.vgg_pre_trained_weights,
                        help='path to pre-trained weights for the vgg19')
    parser.add_argument('--tecogan', action='store_true',
                        default=conf.train.tecogan,
                        help='True for Tecogan training False for FRVSR training')
    parser.add_argument('--checkpoint', type=str,
                        default=conf.train.checkpoint, help='path to checkpoint file')
    args = parser.parse_args()

    # refine config
    conf.data.input_video_dir = args.input_video_dir
    conf.train.max_iter = args.max_iter
    conf.data.output_dir = args.output_dir
    conf.train.num_resblock = args.num_resblock
    conf.train.pre_trained_frvsr_weights = args.pre_trained_frvsr_weights
    conf.train.vgg_pre_trained_weights = args.vgg_pre_trained_weights
    conf.train.tecogan = args.tecogan
    conf.train.checkpoint = args.checkpoint

    return conf


# video-superresolution/tecogan/convert_tf_vgg19_weights.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse

import numpy as np

import nnabla as nn
import nnabla.parametric_functions as PF

from tensorflow.python import pywrap_tensorflow

parser = argparse.ArgumentParser(
    description='Convert TF VGG19 weights to NNabla format')
parser.add_argument('--pre_trained_model', default='./vgg_19.ckpt',
                    help='path to tensorflow pretrained model')
parser.add_argument('--save-path', default='./VGG19_NNabla_model.h5',
                    help='Path to save h5 file')
args = parser.parse_args()


def tf_to_nn_param_map(affine):
    """
    Map from tensorflow default param names to NNabla default param names
    """
    d_1 = {'weights': 'conv/W', 'biases': 'conv/b',
           "conv1_1": "conv1", "conv1_2": "conv2",
           "conv2_1": "conv3", "conv2_2": "conv4",
           "conv3_1": "conv5", "conv3_2": "conv6", "conv3_3": "conv7", "conv3_4": "conv8",
           "conv4_1": "conv9", "conv4_2": "conv10", "conv4_3": "conv11", "conv4_4": "conv12",
           "conv5_1": "conv13", "conv5_2": "conv14", "conv5_3": "conv15", "conv5_4": "conv16",
           "vgg19": ""}
    d_2 = {"fc6": "classifier/0", "weights": "affine/W", "biases": "affine/b",
           "fc7": "classifier/3",
           "fc8": "classifier/6", "vgg19": ""}
    if not affine:
        return d_1
    return d_2


def rename_params(param_name, affine):
    """
    Rename the tensorflow param names to corresponding NNabla param names
    """
    py_to_nn_dict = tf_to_nn_param_map(affine)
    params = param_name.split("/")
    new_param = []
    for param in params:
        if param in py_to_nn_dict.keys():
            param = py_to_nn_dict[param]
            new_param.append(param)
    return "/".join(new_param)


def convert_ckpt_to_h5(input_file, h5_file):
    """
    Convert the input checkpoint file to output hdf5 file
    """
    # Get tensorflow checkpoint reader
    reader = pywrap_tensorflow.NewCheckpointReader(input_file)
    var_to_shape_map = reader.get_variable_to_shape_map()

    # Loop through each tensor name from the variable to shape map
    for key in var_to_shape_map:
        # Read tensor values for each tensor name
        weight = reader.get_tensor(key)
        if not str(key).startswith("vgg_19/mean_rgb") and not str(key).startswith("global_step"):
            s = key.split('/')
            if str(s[-2]).startswith("fc"):
                k = ('/'.join(s))
                key = rename_params(str(k), affine=True)
            else:
                s.remove(s[1])
                k = ('/'.join(s))
                key = rename_params(str(k), affine=False)
            if weight.ndim == 4:
                # transpose TF weight to NNabla weight format
                weight = np.transpose(weight, (3, 0, 1, 2))

                # Create parameter with the same tensor name and shape
            params = PF.get_parameter_or_create(key, shape=weight.shape)
            params.d = weight

        # Save to a h5 file
    nn.parameter.save_parameters(h5_file)


def main():
    """
    Conversion of VGG19 weights from Tensorflow to NNabla
    """
    convert_ckpt_to_h5(args.pre_trained_model, args.save_path)


if __name__ == '__main__':
    main()


# video-superresolution/tecogan/data_loader.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import collections
import numpy as np
import cv2 as cv
from scipy import signal
import nnabla as nn
import nnabla.functions as F
from nnabla.utils.data_iterator import data_iterator_simple


def inference_data_loader(filedir):
    """
    Read and prepare the inference data from specified directory
    filedir: inference data directory name
    return: list of inference image names and images to be enhanced
    """

    image_list_lr_temp = os.listdir(filedir)
    image_list_lr_temp = [_ for _ in image_list_lr_temp if _.endswith(".png")]
    # first sort according to abc, then sort according to 123
    image_list_lr_temp = sorted(image_list_lr_temp)
    image_list_lr_temp.sort(key=lambda f: int(
        ''.join(list(filter(str.isdigit, f))) or -1))
    image_list_lr = [os.path.join(filedir, _) for _ in image_list_lr_temp]

    # Read in and preprocess the images
    def preprocess_test(name):
        ip_img = cv.imread(name, 3).astype(np.float32)[:, :, ::-1]
        max_divided_img = ip_img / 255.0  # equivalent to np.max(ip_img)
        return max_divided_img

    image_lr = [preprocess_test(_) for _ in image_list_lr]

    # a hard-coded symmetric padding
    image_list_lr = image_list_lr[5:0:-1] + image_list_lr
    image_lr = image_lr[5:0:-1] + image_lr

    Data = collections.namedtuple('Data', 'paths_lr, inputs')
    return Data(
        paths_lr=image_list_lr,
        inputs=image_lr
    )


def gaussian_2dkernel(size=5, sig=1.0):
    """
    Returns a 2D Gaussian kernel array with side length size and a sigma of sig
    """
    g_kern1d = signal.gaussian(size, std=sig).reshape(size, 1)
    g_kern2d = np.outer(g_kern1d, g_kern1d)
    return g_kern2d/g_kern2d.sum()


def nn_data_gauss_down_quad(hr_data, sigma=1.5):
    """
    2D down-scaling by 4 with Gaussian blur
    sigma: the sigma used for Gaussian blur
    return: down-scaled data
    """

    k_w = 1 + 2 * int(sigma * 3.0)
    gau_k = gaussian_2dkernel(k_w, sigma)
    gau_0 = np.zeros_like(gau_k)
    gau_wei = np.float32([
        [gau_k, gau_0, gau_0],
        [gau_0, gau_k, gau_0],
        [gau_0, gau_0, gau_k]])  # only works for RGB images!
    gau_wei = np.transpose(gau_wei, [0, 2, 3, 1])
    gau_wei = nn.Variable.from_numpy_array(gau_wei)
    down_sampled_data = F.convolution(
        hr_data, weight=gau_wei, stride=(4, 4), channel_last=True)
    return down_sampled_data


def get_sample_name_grid(conf):
    """
    Return a conf.train.rnn_n number of image names list
    for ex., for 120 frame video, return a list of 10, each with 110 * sequences file names
    conf: configuration information
    return: grid of file names for TecoGAN training
    """

    # Check the input directory
    if conf.data.input_video_dir == '':
        raise ValueError(
            'Video input directory input_video_dir is not provided')

    if not os.path.exists(conf.data.input_video_dir):
        raise ValueError('Video input directory not found')

    image_grid_hr_r = [[] for _ in range(conf.train.rnn_n)]  # all empty lists

    for dir_i in range(conf.data.str_dir, conf.data.end_dir+1):
        ip_dir = os.path.join(conf.data.input_video_dir,
                              '%s_%04d' % (conf.data.input_video_pre, dir_i))
        if os.path.exists(ip_dir):  # the following names are hard coded
            if not os.path.exists(os.path.join(ip_dir, 'col_high_%04d.png' % conf.data.max_frm)):
                print("Skip %s, since folder doesn't contain enough frames!" % ip_dir)
                continue
            for f_i in range(conf.train.rnn_n):
                image_grid_hr_r[f_i] += [os.path.join(ip_dir, 'col_high_%04d.png' % frame_i)
                                         for frame_i in
                                         range(f_i, conf.data.max_frm - conf.train.rnn_n + f_i + 1)]

    return image_grid_hr_r


# moving decision
def moving_decision(conf):
    """
    Decide based on random sampling whether to randomly move first frame for data augmentation
    conf: configuration information
    return: whether to randomly move first frame, also movement offset parameters
    """

    offset_xy = np.floor(
        np.random.uniform(-3.5, 4.5, [conf.train.rnn_n, 2])).astype(int)

    # [FLAGS.RNN_N , 2], shifts
    pos_xy_tmp = np.cumsum(offset_xy, axis=0)
    pos_xy = np.zeros([conf.train.rnn_n, 2])
    pos_xy[1:, :] = pos_xy_tmp[:-1, :]

    min_pos = np.min(pos_xy, axis=0)
    range_pos = np.max(pos_xy, axis=0) - min_pos  # [ shrink x, shrink y ]
    lefttop_pos = pos_xy - min_pos  # crop point
    is_move = np.random.uniform(0, 1, []).astype(float)

    return (lefttop_pos).astype(int), (range_pos).astype(int), is_move


def preprocess(image):
    """
    modify image range from [0,1] to [-1, 1]
    return: range modified image
    """
    return image * 2 - 1


def data_iterator_sr(conf, num_samples, sample_names, tar_size, shuffle, rng=None):
    """
    Data iterator for TecoGAN training
    return: makes provision for low res & high res frames in RNN segments for specified batch_size
    """

    def populate_hr_data(i):

        hr_data = []  # high res rgb, in range 0-1, shape any

        # moving first frame -> data augmentation
        # our data augmentation, moving first frame to mimic camera motion
        if conf.train.movingFirstFrame:
            lefttop_pos, range_pos, is_move = moving_decision(conf)

        for f_i in range(conf.train.rnn_n):
            img_name = sample_names[f_i][i]
            img_data = cv.imread(img_name, 3).astype(np.float32)
            img_data = img_data/255

            if conf.train.movingFirstFrame:
                if f_i == 0:
                    img_data_0 = img_data
                    target_size = img_data.shape

                # random data augmentation -> move first frame only with 30% probability
                if not is_move < 0.7:
                    img_data = img_data_0[lefttop_pos[f_i][1]:target_size[0]
                                          - range_pos[1] + lefttop_pos[f_i][1],
                                          lefttop_pos[f_i][0]:target_size[1]
                                          - range_pos[0] + lefttop_pos[f_i][0], :]
            hr_data.append(img_data)

        return hr_data

    def dataset_load_func(i):

        hr_data = populate_hr_data(i)

        # random crop each batch entry separately
        # Check whether perform crop
        if conf.train.random_crop is True:
            cur_size = hr_data[0].shape
            offset_h = np.floor(np.random.uniform(
                0, cur_size[0] - tar_size, [])).astype(int)
            offset_w = np.floor(np.random.uniform(
                0, cur_size[1] - tar_size, [])).astype(int)
            for frame_t in range(conf.train.rnn_n):
                hr_data[frame_t] = hr_data[frame_t][offset_h:offset_h +
                                                    tar_size, offset_w:offset_w + tar_size, :]

        # random flip:
        if conf.train.flip is True:
            # Produce the decision of random flip
            flip_decision = np.random.uniform(0, 1, []).astype(float)
            for frame_t in range(conf.train.rnn_n):
                if flip_decision < 0.5:
                    np.fliplr(hr_data[frame_t])

        hr_frames = hr_data
        target_frames = []

        k_w_border = int(1.5 * 3.0)
        for rnn_inst in range(conf.train.rnn_n):
            # crop out desired data
            cropped_data = hr_data[rnn_inst][k_w_border:k_w_border+conf.train.crop_size*4,
                                             k_w_border:k_w_border+conf.train.crop_size*4, :]
            pre_processed_data = preprocess(cropped_data)
            target_frames.append(pre_processed_data)

        return hr_frames, target_frames

    return data_iterator_simple(dataset_load_func, num_samples, conf.train.batch_size,
                                shuffle=shuffle, rng=rng, with_file_cache=False,
                                with_memory_cache=False)


# video-superresolution/tecogan/generate.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import argparse
import numpy as np
import nnabla as nn
import nnabla.functions as F
from nnabla.ext_utils import get_extension_context
from utils.utils import save_img, space_to_depth, upscale_four, warp_by_flow
from data_loader import inference_data_loader
from models import generator, flow_estimator

parser = argparse.ArgumentParser(description='TecoGAN')
parser.add_argument('--model', default='./TecoGAN_NNabla_model.h5',
                    help='path to converted NNabla weight file')
parser.add_argument('--input-dir-lr', default='./LR/calendar',
                    help='Path to input LR directory')
parser.add_argument('--output-dir', default='./results/calendar',
                    help='Path to save the output HR frames')
parser.add_argument('--output-name', default='output',
                    help='The pre name of the outputs')
parser.add_argument('--output-ext', default='png',
                    help='The format of the output')
parser.add_argument('--num_resblock', default=16, type=int,
                    help='No of residual blocks in generator. (16 for TecoGAN and 10 FRVSR)')
args = parser.parse_args()

ctx = get_extension_context('cudnn')
nn.set_default_context(ctx)


def main():
    """
        Inference function to generate SR images.
    """
    nn.load_parameters(args.model)
    # Inference data loader
    inference_data = inference_data_loader(args.input_dir_lr)
    input_shape = [1, ] + list(inference_data.inputs[0].shape)
    output_shape = [1, input_shape[1]*4, input_shape[2]*4, 3]
    oh = input_shape[1] - input_shape[1]//8 * 8
    ow = input_shape[2] - input_shape[2]//8 * 8

    # Build the computation graph
    inputs_raw = nn.Variable(input_shape)
    pre_inputs = nn.Variable(input_shape)
    pre_gen = nn.Variable(output_shape)
    pre_warp = nn.Variable(output_shape)

    transposed_pre_warp = space_to_depth(pre_warp)
    inputs_all = F.concatenate(inputs_raw, transposed_pre_warp)
    with nn.parameter_scope("generator"):
        gen_output = generator(inputs_all, 3, args.num_resblock)
    outputs = (gen_output + 1) / 2
    inputs_frames = F.concatenate(pre_inputs, inputs_raw)
    with nn.parameter_scope("fnet"):
        flow_lr = flow_estimator(inputs_frames)
    flow_lr = F.pad(flow_lr, (0, 0, 0, oh, 0, ow, 0, 0), "reflect")
    flow_hr = upscale_four(flow_lr*4.0)
    pre_gen_warp = warp_by_flow(pre_gen, flow_hr)

    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    max_iter = len(inference_data.inputs)
    print('Frame evaluation starts!!')
    pre_inputs.d, pre_gen.d, pre_warp.d = 0, 0, 0
    for i in range(max_iter):
        inputs_raw.d = np.array([inference_data.inputs[i]]).astype(np.float32)
        if i != 0:
            pre_gen_warp.forward(clear_buffer=True)
            pre_warp.data.copy_from(pre_gen_warp.data)
        outputs.forward(clear_buffer=True)
        output_frame = outputs.d

        if i >= 5:
            name, _ = os.path.splitext(
                os.path.basename(str(inference_data.paths_lr[i])))
            filename = args.output_name+'_'+name
            print('saving image %s' % filename)
            out_path = os.path.join(args.output_dir, "%s.%s" %
                                    (filename, args.output_ext))
            save_img(out_path, output_frame[0])
        else:  # First 5 is a hard-coded symmetric frame padding, ignored but time added!
            print("Warming up %d" % (5-i))

        pre_inputs.data.copy_from(inputs_raw.data)
        pre_gen.data.copy_from(outputs.data)


if __name__ == '__main__':
    main()


# video-superresolution/tecogan/models.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import nnabla as nn
import nnabla.functions as F
import nnabla.parametric_functions as PF
from utils.utils import bicubic_four


def conv2d(input, output_channels, kernel, stride, name='', pad=(1, 1), bias=True):
    """
    Convolution layer:
    """
    return PF.convolution(input, output_channels, kernel=kernel, stride=stride,
                          pad=pad, with_bias=bias, name=name, channel_last=True)


def deconv2d(input, output_channels, kernel, stride, name='', pad=(1, 1), output_padding=(1, 1), bias=True):
    """
    Deconvolution layer
    """
    return PF.deconvolution(input, output_channels, kernel=kernel, stride=stride, output_padding=output_padding,
                            pad=pad, with_bias=bias, name=name, channel_last=True)


# Definition of the flow_estimation : Flow-Estimation network
def flow_estimator(fnet_input):
    """
    Flow Estimator network
        Returns the flow estmation for given batch of consecutive frames.
    """
    def down_block(input, output_channels=64, stride=1, scope='down_block'):
        with nn.parameter_scope(scope):
            net = conv2d(input, output_channels, (3, 3),
                         (stride, stride), name='conv_1')
            net = F.leaky_relu(net, 0.2)
            net = conv2d(net, output_channels, (3, 3),
                         (stride, stride), name='conv_2')
            net = F.leaky_relu(net, 0.2)
            net = F.max_pooling(net, (2, 2), channel_last=True)
        return net

    def up_block(input, output_channels=64, stride=1, scope='up_block'):
        with nn.parameter_scope(scope):
            net = conv2d(input, output_channels, (3, 3),
                         (stride, stride), name='conv_1')
            net = F.leaky_relu(net, 0.2)
            net = conv2d(net, output_channels, (3, 3),
                         (stride, stride), name='conv_2')
            net = F.leaky_relu(net, 0.2)
            net = F.interpolate(net, scale=(2, 2), channel_last=True)
        return net

    with nn.parameter_scope('autoencode_unit'):
        net = down_block(fnet_input, 32, scope='encoder_1')
        net = down_block(net, 64, scope='encoder_2')
        net = down_block(net, 128, scope='encoder_3')

        net = up_block(net, 256, scope='decoder_1')
        net = up_block(net, 128, scope='decoder_2')
        net = up_block(net, 64, scope='decoder_3')
        with nn.parameter_scope('output_stage'):
            net = conv2d(net, 32, (3, 3), (1, 1), name='conv1')
            net = F.leaky_relu(net, 0.2)
            net = conv2d(net, 2, (3, 3), (1, 1), name='conv2')
            # the 24.0 is the max Velocity, details can be found in TecoGAN paper
            net = F.tanh(net) * 24.0
    return net


def generator(gen_input, gen_output_channels, num_resblock=16):
    """
    Generator network
        Returns the HR frames for given batch of LR frames
    """
    # The residual blocks
    def residual_block(input, output_channels=64, stride=1, scope='res_block'):
        with nn.parameter_scope(scope):
            net = conv2d(input, output_channels, (3, 3),
                         (stride, stride), name='conv_1')
            net = F.relu(net)
            net = conv2d(net, output_channels, (3, 3),
                         (stride, stride), name='conv_2')
            net = net + input
        return net

    with nn.parameter_scope('generator_unit'):
        # The input layer
        with nn.parameter_scope('input_stage'):
            net = conv2d(gen_input, 64, (3, 3), (1, 1), name='conv')
            net = F.relu(net)

        # The residual block parts
        for i in range(0, num_resblock, 1):  # should be 16 for TecoGAN, and 10 for TecoGANmini
            net = residual_block(net, 64, 1, scope='resblock_%d' % (i+1))

        with nn.parameter_scope('conv_tran2highres'):
            net = deconv2d(net, 64, (3, 3), (2, 2), name='conv_tran1')
            net = F.relu(net)
            net = deconv2d(net, 64, (3, 3), (2, 2), name='conv_tran2')
            net = F.relu(net)

        with nn.parameter_scope('output_stage'):
            net = conv2d(net, gen_output_channels, (3, 3), (1, 1), name='conv')
            low_res_in = gen_input[:, :, :, 0:3]  # ignore warped pre high res
            bicubic_hi = bicubic_four(low_res_in)
            net = net + bicubic_hi
            net = net * 2 - 1
    return net


def discriminator_block(input, out_channels, scope, kernel=(4, 4), stride=(2, 2)):
    with nn.parameter_scope(scope):
        net = conv2d(input, out_channels, kernel, stride, bias=False)
        net = PF.batch_normalization(net, eps=0.001, axes=[3], no_scale=True)
        net = F.leaky_relu(net, alpha=0.2)
    return net


def discriminator(dis_input):
    """
    Discriminator network
        Returns output from dicsriminator block 1 to 7
    """
    layer_list = []
    net = conv2d(dis_input, 64, (3, 3), (1, 1), name='conv')
    net = F.leaky_relu(net, alpha=0.2)        # (b,h,w,64)
    net = discriminator_block(net, 64, 'disblock_1')        # (b,h/2,w/2,64)
    layer_list += [net]
    net = discriminator_block(net, 64, 'disblock_3')        # (b,h/4,w/4,64)
    layer_list += [net]
    net = discriminator_block(net, 128, 'disblock_5')       # (b,h/8,w/8,64)
    layer_list += [net]
    net = discriminator_block(net, 256, 'disblock_7')       # (b,h/16,w/16,64)
    layer_list += [net]
    # channel wise affine
    net = PF.affine(net, 1, base_axis=3, name='affine0')
    net = F.sigmoid(net)                                    # (b,h/16,w/16,1)
    return net, layer_list


# video-superresolution/tecogan/tecogan_model.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import collections
import nnabla as nn
import nnabla.functions as F
from nnabla.monitor import Monitor, MonitorSeries, MonitorTimeElapsed
from models import flow_estimator, generator, discriminator
from utils import set_persistent_all
from utils.utils import deprocess, warp_by_flow, space_to_depth, space_to_depth_disc, upscale_four
from vgg19 import VGG19


def get_common_monitors(monitor):
    """
    Create monitors for displaying and storing losses.
    """
    monitor_content_loss = MonitorSeries(
        'content loss', monitor, interval=20)
    monitor_gen_loss = MonitorSeries(
        'generator loss', monitor, interval=20)
    monitor_warp_loss = MonitorSeries(
        'warp loss', monitor, interval=20)
    monitor_lr = MonitorSeries(
        'learning rate', monitor, interval=20)
    monitor_time = MonitorTimeElapsed(
        "training time per iteration", monitor, interval=20)
    Monitor_common = collections.namedtuple('Monitor_common',
                                            ['monitor_content_loss', 'monitor_gen_loss', 'monitor_warp_loss',
                                             'monitor_lr', 'monitor_time'])
    return Monitor_common(monitor_content_loss, monitor_gen_loss, monitor_warp_loss, monitor_lr, monitor_time)


def get_tecogan_monitors(monitor):
    """
    Create monitors for displaying and storing TECOGAN losses.
    """
    monitor_vgg_loss = MonitorSeries(
        'vgg loss', monitor, interval=20)
    monitor_pp_loss = MonitorSeries(
        'ping pong', monitor, interval=20)
    monitor_sum_layer_loss = MonitorSeries(
        'd layer loss', monitor, interval=20)
    monitor_adv_loss = MonitorSeries(
        'adversarial loss', monitor, interval=20)
    monitor_disc_loss = MonitorSeries(
        'discriminator loss', monitor, interval=20)
    monitor_tb = MonitorSeries(
        'tb', monitor, interval=20)
    Monitor_tecogan = collections.namedtuple('Monitor_tecogan',
                                             ['monitor_vgg_loss', 'monitor_pp_loss', 'monitor_sum_layer_loss',
                                              'monitor_adv_loss', 'monitor_disc_loss', 'monitor_tb'])
    return Monitor_tecogan(monitor_vgg_loss, monitor_pp_loss, monitor_sum_layer_loss, monitor_adv_loss, monitor_disc_loss, monitor_tb)


def get_tecogan_inputs(r_inputs, r_targets):
    """
    Generate and return the ping-pong sequence (forward and backward) from given inputs and targets
    """
    r_inputs = F.concatenate(r_inputs, r_inputs[:, -2::-1, :, :, :], axis=1)
    r_targets = F.concatenate(r_targets, r_targets[:, -2::-1, :, :, :], axis=1)
    return r_inputs, r_targets


def get_fnet_output(conf, rnn_length, frame_t_pre, frame_t, scope_name):
    """
    Return the flow estimations for LR and HR from flow-estimator network
    """
    fnet_input = F.concatenate(frame_t_pre, frame_t)
    fnet_input = F.reshape(fnet_input, (conf.train.batch_size*(rnn_length-1),
                                        conf.train.crop_size, conf.train.crop_size, 2*3))
    with nn.parameter_scope(scope_name + "fnet"):
        flow_lr = flow_estimator(fnet_input)
    flow_hr = upscale_four(flow_lr*4.0)  # a linear up-sampling
    flow_hr = F.reshape(flow_hr, (conf.train.batch_size, (rnn_length-1), conf.train.crop_size*4,
                                  conf.train.crop_size*4, 2), inplace=False)
    fnet_output = collections.namedtuple('fnet_output', 'flow_lr, flow_hr')
    return fnet_output(
        flow_lr=flow_lr,
        flow_hr=flow_hr,
                )


def get_generator_output(conf, rnn_length, r_inputs, flow_hr, scope_name):
    """
    Return the generated HR frames
    """
    # list for all outputs
    gen_outputs = []

    # for the first frame, concat with zeros
    input0 = F.concatenate(r_inputs[:, 0, :, :, :], F.constant(
        0, (conf.train.batch_size, conf.train.crop_size, conf.train.crop_size, 3*4*4)))
    with nn.parameter_scope(scope_name + "generator"):
        gen_pre_output = generator(input0, 3, conf.train.num_resblock)
    gen_outputs.append(gen_pre_output)  # append generated HR frame-0

    for frame_i in range(rnn_length - 1):
        cur_flow = flow_hr[:, frame_i, :, :, :]
        # warp the previously generated frame
        gen_pre_output_warp = warp_by_flow(gen_pre_output, cur_flow)
        gen_pre_output_warp = F.identity(deprocess(gen_pre_output_warp))
        # apply space-to-depth transform
        gen_pre_output_warp = space_to_depth(gen_pre_output_warp)
        # pack it as the recurrent input
        inputs = F.concatenate(
            r_inputs[:, frame_i+1, :, :, :], gen_pre_output_warp)
        # super-resolution part
        with nn.parameter_scope(scope_name + "generator"):
            gen_output = generator(inputs, 3, conf.train.num_resblock)
        gen_outputs.append(gen_output)
        gen_pre_output = gen_output

    # gen_outputs, a list, len = frame, shape = (batch, FLAGS.crop_size*4, FLAGS.crop_size*4, 3)
    gen_outputs = F.stack(*gen_outputs, axis=1)
    # gen_outputs, nn.Variable with shape = (batch, frame, FLAGS.crop_size*4, FLAGS.crop_size*4, 3)

    return gen_outputs


def get_warp_loss(conf, rnn_length, frame_t, frame_t_pre, flow_lr):
    """
    Warp loss
    """
    input_frames = F.reshape(frame_t, (conf.train.batch_size*(rnn_length-1),
                                       conf.train.crop_size, conf.train.crop_size, 3))
    frame_t_pre_reshaped = F.reshape(frame_t_pre, (conf.train.batch_size*(rnn_length-1),
                                                   conf.train.crop_size, conf.train.crop_size, 3))
    s_input_warp = warp_by_flow(frame_t_pre_reshaped, flow_lr)

    warp_loss = F.mean(F.sum(F.squared_error(
        input_frames, s_input_warp), axis=[3]))

    return warp_loss


def get_d_data(conf, flow_hr, gen_outputs, r_targets, rnn_length):
    """
    prepare data for temporal Discriminators
    """
    # 3 frames are used as one entry, the last input images%3 frames are abandoned
    t_size = int(3 * (rnn_length // 3))
    t_gen_output = F.reshape(gen_outputs[:, :t_size, :, :, :], (conf.train.batch_size *
                                                                t_size, conf.train.crop_size*4, conf.train.crop_size*4, 3),
                             inplace=False)
    t_targets = F.reshape(r_targets[:, :t_size, :, :, :], (conf.train.batch_size *
                                                           t_size, conf.train.crop_size*4, conf.train.crop_size*4, 3),
                          inplace=False)
    t_batch = conf.train.batch_size * t_size // 3
    t_inputs_v_pre_batch = F.identity(
                flow_hr[:, 0:t_size:3, :, :, :])  # forward motion reused,
    t_inputs_v_batch = nn.Variable(t_inputs_v_pre_batch.shape)
    # no motion for middle frames
    t_inputs_v_batch.data.zero()
    t_inputs_v_nxt_batch = F.identity(
                flow_hr[:, -2:-1-t_size:-3, :, :, :])  # backward motion

    t_vel = F.stack(
                *[t_inputs_v_pre_batch, t_inputs_v_batch, t_inputs_v_nxt_batch], axis=2)
    # batch, t_size/3, 3, FLAGS.crop_size*4, FLAGS.crop_size*4, 2
    t_vel = F.reshape(t_vel, (conf.train.batch_size*t_size,
                              conf.train.crop_size*4, conf.train.crop_size*4, 2), inplace=False)
    # Stop gradient to fnet from discriminator, details in TecoGAN supplemental paper
    t_vel.need_grad = False

    disc_data = collections.namedtuple('disc_data',
                                       't_vel, t_gen_output, t_batch, t_targets, t_size')
    return disc_data(
        t_vel=t_vel,
        t_gen_output=t_gen_output,
        t_batch=t_batch,
        t_targets=t_targets,
        t_size=t_size
                )


def get_t_d(conf, r_inputs, d_data):
    """
    Create Real and fake temoral discriminators
    """
    # to crop out unstable part for temporal discriminator, details in TecoGAN supplemental paper
    crop_size_dt = int(conf.train.crop_size * 4 * conf.gan.crop_dt)
    offset_dt = (conf.train.crop_size * 4 - crop_size_dt) // 2
    crop_size_dt = conf.train.crop_size * 4 - offset_dt*2
    paddings = (0, 0, offset_dt, offset_dt, offset_dt, offset_dt, 0, 0)

    with nn.parameter_scope("discriminator"):
        real_warp = warp_by_flow(d_data.t_targets, d_data.t_vel)
        real_warp = space_to_depth_disc(real_warp, d_data.t_batch)

        # equivalent to tf.image.crop_to_bounding_box
        real_warp = real_warp[:, offset_dt:offset_dt +
                              crop_size_dt, offset_dt:offset_dt+crop_size_dt, :]
        real_warp = F.pad(real_warp, paddings)
        before_warp = space_to_depth_disc(d_data.t_targets, d_data.t_batch)
        t_input = space_to_depth_disc(
            r_inputs[:, :d_data.t_size, :, :, :], d_data.t_batch)
        # resizing using bilinear interpolation
        input_hi = F.interpolate(t_input, scale=(4, 4), mode='linear',
                                 channel_last=True)
        real_warp = F.concatenate(before_warp, real_warp, input_hi)

        tdiscrim_real_output, real_layers = discriminator(real_warp)

        fake_warp = warp_by_flow(d_data.t_gen_output, d_data.t_vel)
        fake_warp = space_to_depth_disc(fake_warp, d_data.t_batch)
        fake_warp = fake_warp[:, offset_dt:offset_dt +
                              crop_size_dt, offset_dt:offset_dt+crop_size_dt, :]
        fake_warp = F.pad(fake_warp, paddings)
        before_warp = space_to_depth_disc(
            d_data.t_gen_output, d_data.t_batch, inplace=False)
        fake_warp = F.concatenate(before_warp, fake_warp, input_hi)
        tdiscrim_fake_output, fake_layers = discriminator(fake_warp)

    temporal_disc = collections.namedtuple('temporal_disc', 'tdiscrim_real_output,'
                                           'real_layers, tdiscrim_fake_output, fake_layers')
    return temporal_disc(
        tdiscrim_real_output=tdiscrim_real_output,
        real_layers=real_layers,
        tdiscrim_fake_output=tdiscrim_fake_output,
        fake_layers=fake_layers
                )


def get_d_layer(real_layers, fake_layers):
    """
    discriminator layer loss
    """
    fix_range = 0.02  # hard coded, all layers are roughly scaled to this value
    sum_layer_loss = 0  # adds-on for generator
    layer_loss_list = []
    layer_n = len(real_layers)

    # hard coded, an overall average of all layers
    layer_norm = [12.0, 14.0, 24.0, 100.0]

    for layer_i in range(layer_n):
        real_layer = real_layers[layer_i]
        false_layer = fake_layers[layer_i]

        layer_diff = real_layer - false_layer
        layer_loss = F.mean(
                        F.sum(F.abs(layer_diff), axis=[3]))  # an l1 loss
        layer_loss_list += [layer_loss]
        scaled_layer_loss = fix_range * \
            layer_loss / layer_norm[layer_i]
        sum_layer_loss += scaled_layer_loss

    return sum_layer_loss


def get_vgg_loss(s_gen_output, s_targets):
    """
    VGG loss
    """
    vgg19 = VGG19()
    vgg_gen_fea = vgg19(s_gen_output)
    vgg_target_fea = vgg19(s_targets)
    vgg_loss_list = []
    # we use 4 VGG layers
    vgg_wei_list = [1.0, 1.0, 1.0, 1.0]
    vgg_loss = 0
    vgg_layer_n = 4

    for layer_i in range(vgg_layer_n):
        # cosine similarity, -1~1, 1 best
        curvgg_diff = F.sum(
            vgg_gen_fea[layer_i]*vgg_target_fea[layer_i], axis=[3])
        curvgg_diff = 1.0 - F.mean(curvgg_diff)  # 0 ~ 2, 0 best
        scaled_layer_loss = vgg_wei_list[layer_i] * curvgg_diff
        vgg_loss_list += [curvgg_diff]
        vgg_loss += scaled_layer_loss

    return vgg_loss


def get_tecogan_model(conf, r_inputs, r_targets, scope_name, tecogan=True):
    """
    Create computation graph and variables for TecoGAN.
    """
    # r_inputs, r_targets : shape (batch, conf.train.rnn_n, h, w, c)
    rnn_length = conf.train.rnn_n
    if tecogan:
        r_inputs, r_targets = get_tecogan_inputs(r_inputs, r_targets)
        rnn_length = rnn_length * 2 - 1

    # get the consecutive frame sequences from the input sequence
    frame_t_pre, frame_t = r_inputs[:, 0:-1, :, :, :], r_inputs[:, 1:, :, :, :]

    # Get flow estimations
    fnet_output = get_fnet_output(
        conf, rnn_length, frame_t_pre, frame_t, scope_name)

    # Get the generated HR output frames
    gen_outputs = get_generator_output(
        conf, rnn_length, r_inputs, fnet_output.flow_hr, scope_name)

    s_gen_output = F.reshape(gen_outputs, (conf.train.batch_size * rnn_length,
                                           conf.train.crop_size * 4,
                                           conf.train.crop_size*4, 3), inplace=False)
    s_targets = F.reshape(r_targets, (conf.train.batch_size * rnn_length, conf.train.crop_size * 4,
                                      conf.train.crop_size * 4, 3), inplace=False)

    # Content loss (l2 loss)
    content_loss = F.mean(
        F.sum(F.squared_error(s_gen_output, s_targets), axis=[3]))
    # Warp loss (l2 loss)
    warp_loss = get_warp_loss(
        conf, rnn_length, frame_t, frame_t_pre, fnet_output.flow_lr)

    if tecogan:
        d_data = get_d_data(conf, fnet_output.flow_hr,
                            gen_outputs, r_targets, rnn_length)
        # Build the tempo discriminator for the real part and fake part
        t_d = get_t_d(conf, r_inputs, d_data)

        # Discriminator layer loss:
        d_layer_loss = get_d_layer(t_d.real_layers, t_d.fake_layers)
        # vgg loss (cosine similarity)
        loss_vgg = get_vgg_loss(s_gen_output, s_targets)
        # ping pong loss (an l1 loss)
        gen_out_first = gen_outputs[:, 0:conf.train.rnn_n-1, :, :, :]
        gen_out_last_rev = gen_outputs[:, -1:-conf.train.rnn_n:-1, :, :, :]
        pp_loss = F.mean(F.abs(gen_out_first - gen_out_last_rev))
        # adversarial loss
        t_adversarial_loss = F.mean(
            -F.log(t_d.tdiscrim_fake_output + conf.train.eps))

        # Overall generator loss
        gen_loss = content_loss + pp_loss * conf.gan.pp_scaling + conf.gan.ratio * \
            t_adversarial_loss + conf.gan.vgg_scaling * loss_vgg + \
            conf.gan.dt_ratio_0 * d_layer_loss

        # Discriminator loss
        t_discrim_fake_loss = F.log(
            1 - t_d.tdiscrim_fake_output + conf.train.eps)
        t_discrim_real_loss = F.log(t_d.tdiscrim_real_output + conf.train.eps)
        t_discrim_loss = F.mean(-(t_discrim_fake_loss + t_discrim_real_loss))

        fnet_loss = gen_loss + warp_loss

        set_persistent_all(r_targets, r_inputs, loss_vgg, gen_out_first, gen_out_last_rev, pp_loss,
                           d_layer_loss, content_loss, warp_loss, gen_loss, t_adversarial_loss,
                           t_discrim_loss, t_discrim_real_loss, d_data.t_vel, d_data.t_gen_output,
                           s_gen_output, s_targets)

        Network = collections.namedtuple('Network', 'content_loss,  warp_loss, fnet_loss, vgg_loss,'
                                         'gen_loss, pp_loss, sum_layer_loss,t_adversarial_loss,'
                                         't_discrim_loss,t_gen_output,t_discrim_real_loss')
        return Network(
            content_loss=content_loss,
            warp_loss=warp_loss,
            fnet_loss=fnet_loss,
            vgg_loss=loss_vgg,
            gen_loss=gen_loss,
            pp_loss=pp_loss,
            sum_layer_loss=d_layer_loss,
            t_adversarial_loss=t_adversarial_loss,
            t_discrim_loss=t_discrim_loss,
            t_gen_output=d_data.t_gen_output,
            t_discrim_real_loss=t_discrim_real_loss
        )

    gen_loss = content_loss
    fnet_loss = gen_loss + warp_loss
    set_persistent_all(content_loss, s_gen_output,
                       warp_loss, gen_loss, fnet_loss)

    Network = collections.namedtuple(
        'Network', 'content_loss, warp_loss, fnet_loss, gen_loss')
    return Network(
        content_loss=content_loss,
        warp_loss=warp_loss,
        fnet_loss=fnet_loss,
        gen_loss=gen_loss,
    )


def get_frvsr_model(conf, r_inputs, r_targets, scope_name):
    """
    Create computation graph and variables for FRVSR.
    """
    return get_tecogan_model(conf, r_inputs, r_targets, scope_name, False)


# video-superresolution/tecogan/train.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import datetime
import numpy as np
import nnabla as nn
import nnabla.functions as F
import nnabla.solvers as S
from nnabla.ext_utils import get_extension_context
from nnabla.monitor import Monitor
from tecogan_model import get_tecogan_model, get_frvsr_model, get_common_monitors, get_tecogan_monitors
from utils import CommunicatorWrapper, save_checkpoint, load_checkpoint
from utils.utils import ExponentialMovingAverage
from args import get_config


def main():
    conf = get_config()
    extension_module = conf.nnabla_context.context
    ctx = get_extension_context(
        extension_module, device_id=conf.nnabla_context.device_id)
    comm = CommunicatorWrapper(ctx)
    nn.set_default_context(comm.ctx)
    print("comm rank", comm.rank)

    # data iterators for train and val data
    from data_loader import data_iterator_sr, get_sample_name_grid, nn_data_gauss_down_quad

    sample_names = get_sample_name_grid(conf)
    num_samples = len(sample_names[0])
    print("No of training samples :", num_samples)

    tar_size = conf.train.crop_size
    tar_size = (conf.train.crop_size * 4) + int(1.5 * 3.0) * \
        2  # crop_size * 4, and Gaussian blur margin

    data_iterator_train = data_iterator_sr(
        conf, num_samples, sample_names, tar_size, shuffle=True)

    if comm.n_procs > 1:
        data_iterator_train = data_iterator_train.slice(
            rng=None, num_of_slices=comm.n_procs, slice_pos=comm.rank)

    train_hr = nn.Variable(
        (conf.train.batch_size, conf.train.rnn_n, conf.train.crop_size*4, conf.train.crop_size*4, 3))
    data_hr = nn.Variable(
        (conf.train.batch_size, conf.train.rnn_n, tar_size, tar_size, 3))
    train_lr = nn_data_gauss_down_quad(data_hr.reshape(
        (conf.train.batch_size * conf.train.rnn_n, tar_size, tar_size, 3)))
    train_lr = F.reshape(
        train_lr, (conf.train.batch_size, conf.train.rnn_n, conf.train.crop_size, conf.train.crop_size, 3))

    # setting up monitors for logging
    monitor_path = './nnmonitor' + \
        str(datetime.datetime.now().strftime("%Y%m%d%H%M%S"))
    monitor = Monitor(monitor_path)
    common_monitor = get_common_monitors(monitor)

    # Change max_iter and learning_rate when batch size or no. of gpu devices change.
    div_factor = conf.train.batch_size * comm.n_procs
    max_iter = (conf.train.max_iter * 4) // div_factor
    learning_rate = conf.train.learning_rate * \
        (conf.train.batch_size / 4) * comm.n_procs

    if comm.rank == 0:
        print("maximum iterations", max_iter)

    scope_name = 'frvsr/'
    if conf.train.tecogan:
        scope_name = 'tecogan/'
        if not conf.train.checkpoint:
            print('loading pretrained FRVSR model',
                  conf.train.pre_trained_frvsr_weights)
            with nn.parameter_scope(scope_name):
                nn.load_parameters(conf.train.pre_trained_frvsr_weights)
                params_from_pre_trained_model = []
                for key, val in nn.get_parameters().items():
                    params_from_pre_trained_model.append(scope_name + key)

            network = get_tecogan_model(conf, train_lr, train_hr, scope_name)
            params_from_graph = nn.get_parameters()

            # Set the Generator parameters which are not in FRVSR to zero,
            # as done in orig implementation.
            for key, val in params_from_graph.items():
                if key in params_from_pre_trained_model or key.startswith('vgg') or key.startswith('disc'):
                    continue
                print(key)
                val.data.zero()  # fill with zero

        else:
            network = get_tecogan_model(conf, train_lr, train_hr, scope_name)

        # Define discriminator optimizer/solver
        solver_disc = S.Adam(alpha=learning_rate,
                             beta1=conf.train.beta, eps=conf.train.adameps)
        # Set discriminator Parameters
        with nn.parameter_scope("discriminator"):
            solver_disc.set_parameters(nn.get_parameters())

        # setting up monitors for TecoGAN
        tecogan_monitor = get_tecogan_monitors(monitor)

    else:
        network = get_frvsr_model(conf, train_lr, train_hr, scope_name)

    # Define generator and fnet optimizer/solver
    solver_gen = S.Adam(alpha=learning_rate,
                        beta1=conf.train.beta, eps=conf.train.adameps)
    solver_fnet = S.Adam(alpha=learning_rate,
                         beta1=conf.train.beta, eps=conf.train.adameps)

    # Set generator and fnet Parameters
    with nn.parameter_scope(scope_name + "generator"):
        solver_gen.set_parameters(nn.get_parameters())
    with nn.parameter_scope(scope_name + "fnet"):
        solver_fnet.set_parameters(nn.get_parameters())

    if conf.train.tecogan:
        solver_dict = {"gen": solver_gen,
                       "fnet": solver_fnet, "disc": solver_disc}
    else:
        solver_dict = {"gen": solver_gen, "fnet": solver_fnet}

    start_point = 0
    if conf.train.checkpoint:
        # Load optimizer/solver information and model weights from checkpoint
        start_point = load_checkpoint(conf.train.checkpoint, solver_dict)

    # Exponential Moving Average Calculation for tb
    ema = ExponentialMovingAverage(conf.train.decay)
    tb = 0

    # Create output directory if it doesn't exist
    if not os.path.exists(conf.data.output_dir):
        os.makedirs(conf.data.output_dir)

    # Training loop.
    for i in range(start_point, max_iter):
        # Get Training Data
        data_hr.d, train_hr.d = data_iterator_train.next()

        if conf.train.tecogan:
            network.t_discrim_loss.forward(clear_no_need_grad=True)
            if np.less(tb, 0.4):  # train gen with d
                # Compute grads for discriminator and update
                solver_disc.zero_grad()
                # Stop back-propagation from t_discrim_loss to generator
                network.t_gen_output.need_grad = False
                if comm.n_procs > 1:
                    all_reduce_callback = comm.get_all_reduce_callback()
                    network.t_discrim_loss.backward(clear_buffer=True,
                                                    communicator_callbacks=all_reduce_callback)
                else:
                    network.t_discrim_loss.backward(clear_buffer=True)
                solver_disc.update()  # Update grads
                # Enable back propagation from fnet_loss to Generator
                network.t_gen_output.need_grad = True

        # Compute grads for fnet and generator together using fnet_loss
        solver_fnet.zero_grad()
        solver_gen.zero_grad()
        # Apply forward and backward propagation on fnet_loss
        network.fnet_loss.forward(clear_no_need_grad=True)
        if comm.n_procs > 1:
            all_reduce_callback = comm.get_all_reduce_callback()
            network.fnet_loss.backward(clear_buffer=True,
                                       communicator_callbacks=all_reduce_callback)
        else:
            network.fnet_loss.backward(clear_buffer=True)
        # Update grads for fnet and generator
        solver_gen.update()
        solver_fnet.update()

        if conf.train.tecogan:
            if comm.n_procs > 1:
                comm.all_reduce([network.t_discrim_real_loss.data,
                                 network.t_adversarial_loss.data], division=True, inplace=True)
            t_balance = F.mean(network.t_discrim_real_loss.data) + \
                network.t_adversarial_loss.data
            if i == 0:
                ema.register(t_balance)
            else:
                tb = ema(t_balance)
            if comm.rank == 0:
                tecogan_monitor.monitor_pp_loss.add(
                    i, network.pp_loss.d.copy())
                tecogan_monitor.monitor_vgg_loss.add(
                    i, network.vgg_loss.d.copy())
                tecogan_monitor.monitor_sum_layer_loss.add(
                    i, network.sum_layer_loss.d.copy())
                tecogan_monitor.monitor_adv_loss.add(
                    i, network.t_adversarial_loss.d.copy())
                tecogan_monitor.monitor_disc_loss.add(
                    i, network.t_discrim_loss.d.copy())
                tecogan_monitor.monitor_tb.add(i, tb)

        if comm.rank == 0:
            common_monitor.monitor_content_loss.add(
                i, network.content_loss.d.copy())
            common_monitor.monitor_gen_loss.add(i, network.gen_loss.d.copy())
            common_monitor.monitor_warp_loss.add(i, network.warp_loss.d.copy())
            common_monitor.monitor_lr.add(i, learning_rate)
            common_monitor.monitor_time.add(i)
            if (i % conf.train.save_freq) == 0:
                # Save intermediate model parameters
                with nn.parameter_scope(scope_name):
                    nn.save_parameters(os.path.join(
                        conf.data.output_dir, "model_param_%08d.h5" % i))

                # Save intermediate check_points
                save_checkpoint(conf.data.output_dir, i, solver_dict)

    # save final Generator and Fnet network parameters
    if comm.rank == 0:
        with nn.parameter_scope(scope_name):
            nn.save_parameters(os.path.join(
                conf.data.output_dir, "model_param_%08d.h5" % i))


if __name__ == "__main__":
    main()


# video-superresolution/tecogan/vgg19.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import nnabla as nn
import nnabla.functions as F
import nnabla.parametric_functions as PF
from utils.utils import deprocess
from args import get_config


def get_vgg_feat(input_var):
    """
        Exactly the same architecture used for LPIPS.
                Returns output features from conv1 to relu5_4, classification layers are dropped.
    """
    # assert input_var.shape[1] == 3
    act1 = F.relu(PF.convolution(input_var, outmaps=64,
                                 kernel=(3, 3), pad=(1, 1), channel_last=True, fix_parameters=True, name="conv1"), True)
    act1 = F.relu(PF.convolution(act1, outmaps=64, kernel=(
        3, 3), pad=(1, 1), channel_last=True, fix_parameters=True, name="conv2"), True)

    act2 = F.max_pooling(act1, kernel=(2, 2), stride=(2, 2), channel_last=True)
    act2 = F.relu(PF.convolution(act2, outmaps=128, kernel=(
        3, 3), pad=(1, 1), channel_last=True, fix_parameters=True, name="conv3"), True)
    act2 = F.relu(PF.convolution(act2, outmaps=128, kernel=(
        3, 3), pad=(1, 1), channel_last=True, fix_parameters=True, name="conv4"), True)

    act3 = F.max_pooling(act2, kernel=(2, 2), stride=(2, 2), channel_last=True)
    act3 = F.relu(PF.convolution(act3, outmaps=256, kernel=(
        3, 3), pad=(1, 1), channel_last=True, fix_parameters=True, name="conv5"), True)
    act3 = F.relu(PF.convolution(act3, outmaps=256, kernel=(
        3, 3), pad=(1, 1), channel_last=True, fix_parameters=True, name="conv6"), True)
    act3 = F.relu(PF.convolution(act3, outmaps=256, kernel=(
        3, 3), pad=(1, 1), channel_last=True, fix_parameters=True, name="conv7"), True)
    act3 = F.relu(PF.convolution(act3, outmaps=256, kernel=(
        3, 3), pad=(1, 1), channel_last=True, fix_parameters=True, name="conv8"), True)

    act4 = F.max_pooling(act3, kernel=(2, 2), stride=(2, 2), channel_last=True)
    act4 = F.relu(PF.convolution(act4, outmaps=512, kernel=(
        3, 3), pad=(1, 1), channel_last=True, fix_parameters=True, name="conv9"), True)
    act4 = F.relu(PF.convolution(act4, outmaps=512, kernel=(
        3, 3), pad=(1, 1), channel_last=True, fix_parameters=True, name="conv10"), True)
    act4 = F.relu(PF.convolution(act4, outmaps=512, kernel=(
        3, 3), pad=(1, 1), channel_last=True, fix_parameters=True, name="conv11"), True)
    act4 = F.relu(PF.convolution(act4, outmaps=512, kernel=(
        3, 3), pad=(1, 1), channel_last=True, fix_parameters=True, name="conv12"), True)

    act5 = F.max_pooling(act4, kernel=(2, 2), stride=(2, 2), channel_last=True)
    act5 = F.relu(PF.convolution(act5, outmaps=512, kernel=(
        3, 3), pad=(1, 1), channel_last=True, fix_parameters=True, name="conv13"), True)
    act5 = F.relu(PF.convolution(act5, outmaps=512, kernel=(
        3, 3), pad=(1, 1), channel_last=True, fix_parameters=True, name="conv14"), True)
    act5 = F.relu(PF.convolution(act5, outmaps=512, kernel=(
        3, 3), pad=(1, 1), channel_last=True, fix_parameters=True, name="conv15"), True)
    act5 = F.relu(PF.convolution(act5, outmaps=512, kernel=(
        3, 3), pad=(1, 1), channel_last=True, fix_parameters=True, name="conv16"), True)

    return [act2, act3, act4, act5]


class VGG19(object):
    """
        Load VGG19 pretrained weights and
        return output features from conv1 to relu5_4.
        """

    def __init__(self):
        conf = get_config()

        with nn.parameter_scope("vgg19"):
            if not conf.train.checkpoint:
                print("Loading pre-trained vgg19 weights from ",
                      conf.train.vgg_pre_trained_weights)
                nn.load_parameters(conf.train.vgg_pre_trained_weights)

                # drop all the affine layers from pre-trained model for finetuning.
                drop_layers = ['classifier/0/affine',
                               'classifier/3/affine', 'classifier/6/affine']
                for layers in drop_layers:
                    nn.parameter.pop_parameter((layers + '/W'))
                    nn.parameter.pop_parameter((layers + '/b'))
            self.mean = nn.Variable.from_numpy_array(np.asarray(
                [123.68, 116.78, 103.94]).reshape(1, 1, 1, 3))

    def __call__(self, x):
        with nn.parameter_scope("vgg19"):
            results = list()
            input = deprocess(x)
            x_in = F.sub2(input * 255.0, self.mean)
            features = get_vgg_feat(x_in)
            for i in range(len(features)):
                orig_deep_feature = features[i]
                orig_len = (F.sum((orig_deep_feature ** 2),
                                  axis=[3], keepdims=True)+1e-12) ** 0.5
                results.append(orig_deep_feature / orig_len)
            return results


# view-synthesis/nerf/main.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys
from argparse import ArgumentParser
import yaml
import numpy as np

import nnabla as nn
import nnabla.communicators as C

from train.nerf import train_nerf

common_utils_path = os.path.abspath(
    os.path.join(os.path.dirname(__file__), '..', '..', 'utils'))
sys.path.append(common_utils_path)
from neu.yaml_wrapper import read_yaml
from neu.misc import init_nnabla


def get_args():
    parser = ArgumentParser(description='NeRF: Nnabla implementation')

    parser.add_argument('--config-path', '-c', type=str, default='configs/tiny_data.yaml',
                        required=True,
                        help='model and training configuration file')
    parser.add_argument('--model', type=str, choices=['wild', 'uncertainty', 'appearance', 'vanilla'],
                        required=True,
                        help='Select the model to train')

    parser.add_argument('--data-perturb', type=str, choices=['none', 'color', 'occ', 'both'],
                        default='none',
                        help='Select the model to train')
    parser.add_argument('--dataset', type=str, choices=['blender', 'phototourism'],
                        default='blender',
                        help='Select the dataset')
    parser.add_argument('--save-results-dir', type=str, default='',
                        help='Path to save trained model parameters')

    return parser.parse_args()


def main():

    args = get_args()

    config = read_yaml(args.config_path)

    nn.set_auto_forward(True)
    comm = init_nnabla(ext_name="cuda", device_id='0', type_config='float')

    if args.save_results_dir != '':
        config.log.save_results_dir = args.save_results_dir

    config.data.color_perturb = True if (
        args.data_perturb == 'color' or args.data_perturb == 'both') else False
    config.data.occ_perturb = True if (
        args.data_perturb == 'occ' or args.data_perturb == 'both') else False

    if comm is None or comm.rank == 0:
        if config.data.color_perturb:
            print('Applying color perturbation to the dataset')
        if config.data.occ_perturb:
            print('Applying occlusion perturbation to the dataset')
        if not config.data.color_perturb and not config.data.occ_perturb:
            print('No perturbation will be applied to the data')

    train_nerf(config, comm, args.model, args.dataset)


if __name__ == '__main__':
    main()


# view-synthesis/nerf/phototourism_eval.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import nnabla as nn
import nnabla.functions as F
import nnabla.parametric_functions as PF
from nnabla.utils.data_iterator import data_iterator


import numpy as np
import imageio
from tqdm import tqdm, trange
import argparse
import os
import sys

from train.nerf import forward_pass
from train.common import *
from data_iterator import get_photo_tourism_dataiterator


common_utils_path = os.path.abspath(
    os.path.join(os.path.dirname(__file__), '..', '..', 'utils'))
sys.path.append(common_utils_path)
from neu.yaml_wrapper import read_yaml
from neu.misc import init_nnabla


def main():

    parser = argparse.ArgumentParser()

    parser.add_argument('--output-filename', '-o', type=str, default='video.gif',
                        help="name of an output file.")
    parser.add_argument('--output-static-filename', '-os', type=str, default='video_static.gif',
                        help="name of an output file.")
    parser.add_argument('--config-path', '-c', type=str, default='configs/llff.yaml',
                        required=True,
                        help='model and training configuration file')
    parser.add_argument('--weight-path', '-w', type=str, default='configs/llff.yaml',
                        required=True,
                        help='path to pretrained NeRF parameters')
    parser.add_argument('--model', type=str, choices=['wild', 'uncertainty', 'appearance', 'vanilla'],
                        required=True,
                        help='Select the model to train')

    parser.add_argument('--visualization-type', '-v', type=str,
                        choices=['zoom', '360-rotation', 'default'],
                        default='default-render-poses',
                        help='type of visualization')

    parser.add_argument('--downscale', '-d', default=1, type=float,
                        help="downsampling factor for the rendered images for faster inference")

    parser.add_argument('--num-images', '-n', default=120, type=int,
                        help="Number of images to generate for the output video/gif")

    args = parser.parse_args()

    nn.set_auto_forward(True)
    comm = init_nnabla(ext_name="cudnn")

    use_transient = False
    use_embedding = False

    if args.model == 'wild':
        use_transient = True
        use_embedding = True
    elif args.model == 'uncertainty':
        use_transient = True
    elif args.model == 'appearance':
        use_embedding = True

    args = parser.parse_args()
    config = read_yaml(args.config_path)

    config.data.downscale = args.downscale
    nn.load_parameters(args.weight_path)

    data_source = get_photo_tourism_dataiterator(config, 'test', comm)

    # Pose, Appearance index for generating novel views
    # as well as camera trajectory is hard-coded here.
    data_source.test_appearance_idx = 125
    pose_idx = 125
    dx = np.linspace(-0.2, 0.15, args.num_images//3)
    dy = -0.15
    dz = np.linspace(0.1, 0.22, args.num_images//3)

    embed_idx_list = list(data_source.poses_dict.keys())

    data_source.poses_test = np.tile(
        data_source.poses_dict[pose_idx], (args.num_images, 1, 1))
    for i in range(0, args.num_images//3):
        data_source.poses_test[i, 0, 3] += dx[i]
        data_source.poses_test[i, 1, 3] += dy
    for i in range(args.num_images//3, args.num_images//2):
        data_source.poses_test[i, 0, 3] += dx[len(dx)-1-i]
        data_source.poses_test[i, 1, 3] += dy

    for i in range(args.num_images//2, 5*args.num_images//6):
        data_source.poses_test[i, 2, 3] += dz[i-args.num_images//2]
        data_source.poses_test[i, 1, 3] += dy
        data_source.poses_test[i, 0, 3] += dx[len(dx)//2]

    for i in range(5*args.num_images//6, args.num_images):
        data_source.poses_test[i, 2, 3] += dz[args.num_images-1 - i]
        data_source.poses_test[i, 1, 3] += dy
        data_source.poses_test[i, 0, 3] += dx[len(dx)//2]

    # mapping_net = MLP
    encode_position_function = get_encoding_function(
        config.train.num_encodings_position, True, True)
    if config.train.use_view_directions:
        encode_direction_function = get_encoding_function(
            config.train.num_encodings_direction, True, True)
    else:
        encode_direction_function = None

    frames = []
    if use_transient:
        static_frames = []

    pbar = tqdm(np.arange(0, data_source.poses_test.shape[0]))
    data_source._size = data_source.poses_test.shape[0]
    data_source.test_img_w = 400
    data_source.test_img_h = 400
    data_source.test_focal = data_source.test_img_w/2/np.tan(np.pi/6)
    data_source.test_K = np.array([[data_source.test_focal, 0, data_source.test_img_w/2],
                                   [0, data_source.test_focal,
                                       data_source.test_img_h/2],
                                   [0, 0, 1]])

    data_source._indexes = np.arange(0, data_source._size)

    di = data_iterator(data_source, batch_size=1)

    print(f'Rendering {args.num_images} poses...')

    a = [1, 128]
    alpha = np.linspace(0, 1, args.num_images)

    for th in pbar:

        rays, embed_inp = di.next()
        ray_origins = nn.NdArray.from_numpy_array(rays[0, :, :3])
        ray_directions = nn.NdArray.from_numpy_array(rays[0, :, 3:6])
        near_plane_ = nn.NdArray.from_numpy_array(rays[0, :, 6])
        far_plane_ = nn.NdArray.from_numpy_array(rays[0, :, 7])

        embed_inp = nn.NdArray.from_numpy_array(
            embed_inp[0, :config.train.chunksize_fine])
        image_shape = (data_source.test_img_w, data_source.test_img_h, 3)

        ray_directions = F.reshape(ray_directions, (-1, 3))
        ray_origins = F.reshape(ray_origins, (-1, 3))

        num_ray_batches = (
            ray_directions.shape[0] + config.train.ray_batch_size-1)//config.train.ray_batch_size

        app_emb, trans_emb = None, None
        if use_embedding:
            with nn.parameter_scope('embedding_a'):
                embed_inp_app = nn.NdArray.from_numpy_array(
                    np.full((config.train.chunksize_fine,), a[0], dtype=int))
                app_emb = PF.embed(
                    embed_inp_app, config.train.n_vocab, config.train.n_app)

                embed_inp_app = nn.NdArray.from_numpy_array(
                    np.full((config.train.chunksize_fine,), a[1], dtype=int))
                app_emb_2 = PF.embed(
                    embed_inp_app, config.train.n_vocab, config.train.n_app)

                app_emb = app_emb * alpha[th] + app_emb_2 * (1-alpha[th])

        if use_transient:
            with nn.parameter_scope('embedding_t'):
                trans_emb = PF.embed(
                    embed_inp, config.train.n_vocab, config.train.n_trans)

            static_rgb_map_fine_list, transient_rgb_map_fine_list = [], []

        rgb_map_fine_list = []

        for i in trange(num_ray_batches):
            ray_d, ray_o = ray_directions[i*config.train.ray_batch_size:(
                i+1)*config.train.ray_batch_size], ray_origins[i*config.train.ray_batch_size:(i+1)*config.train.ray_batch_size]

            near_plane = near_plane_[i*config.train.ray_batch_size:(
                            i+1)*config.train.ray_batch_size]
            far_plane = far_plane_[i*config.train.ray_batch_size:(
                i+1)*config.train.ray_batch_size]

            if use_transient:
                _, rgb_map_fine, static_rgb_map_fine, transient_rgb_map_fine, _, _, _ = forward_pass(ray_d, ray_o, near_plane, far_plane, app_emb,
                                                                                                     trans_emb, encode_position_function, encode_direction_function, config, use_transient)

                static_rgb_map_fine_list.append(static_rgb_map_fine)
                transient_rgb_map_fine_list.append(transient_rgb_map_fine)

            else:
                _, _, _, _, rgb_map_fine, _, _, _ = \
                    forward_pass(ray_d, ray_o, near_plane, far_plane, app_emb, trans_emb,
                                 encode_position_function, encode_direction_function, config, use_transient)

            rgb_map_fine_list.append(rgb_map_fine)

        rgb_map_fine = F.concatenate(*rgb_map_fine_list, axis=0)
        rgb_map_fine = F.reshape(rgb_map_fine, image_shape)

        if use_transient:
            static_rgb_map_fine = F.concatenate(
                *static_rgb_map_fine_list, axis=0)
            static_rgb_map_fine = F.reshape(static_rgb_map_fine, image_shape)

        frames.append((255*np.clip(rgb_map_fine.data, 0, 1)).astype(np.uint8))
        if use_transient:
            static_frames.append(
                (255*np.clip(static_rgb_map_fine.data, 0, 1)).astype(np.uint8))

    imageio.mimwrite(args.output_filename, frames, fps=30)
    imageio.mimwrite(args.output_static_filename, static_frames, fps=30)


if __name__ == '__main__':
    main()


# view-synthesis/nerf/visualize.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import nnabla as nn
import nnabla.functions as F
import nnabla.parametric_functions as PF
from nnabla.ext_utils import get_extension_context

import numpy as np
import imageio
from tqdm import tqdm, trange
import argparse
import os
import sys

from train.nerf import forward_pass
from train.common import *
from data_iterator.get_data import get_data


common_utils_path = os.path.abspath(
    os.path.join(os.path.dirname(__file__), '..', '..', 'utils'))
sys.path.append(common_utils_path)
from neu.yaml_wrapper import read_yaml, write_yaml


def trans_t(t): return np.array([
    [1, 0, 0, 0],
    [0, 1, 0, 0],
    [0, 0, 1, t],
    [0, 0, 0, 1],
], dtype=np.float32)


def rot_phi(phi): return np.array([
    [1, 0, 0, 0],
    [0, np.cos(phi), -np.sin(phi), 0],
    [0, np.sin(phi), np.cos(phi), 0],
    [0, 0, 0, 1],
], dtype=np.float32)


def rot_theta(th): return np.array([
    [np.cos(th), 0, -np.sin(th), 0],
    [0, 1, 0, 0],
    [np.sin(th), 0, np.cos(th), 0],
    [0, 0, 0, 1],
], dtype=np.float32)


def pose_spherical(theta, phi, radius):
    c2w = trans_t(radius)
    c2w = rot_phi(phi/180.*np.pi) @ c2w
    c2w = rot_theta(theta/180.*np.pi) @ c2w
    c2w = np.array([[-1, 0, 0, 0], [0, 0, 1, 0],
                    [0, 1, 0, 0], [0, 0, 0, 1]]) @ c2w
    return c2w


def main():

    parser = argparse.ArgumentParser()

    parser.add_argument('--output-filename', '-o', type=str, default='video.gif',
                        help="name of an output file.")
    parser.add_argument('--output-static-filename', '-os', type=str, default='video_static.gif',
                        help="name of an output file.")
    parser.add_argument('--config-path', '-c', type=str, default='configs/llff.yaml',
                        required=True,
                        help='model and training configuration file')
    parser.add_argument('--weight-path', '-w', type=str, default='configs/llff.yaml',
                        required=True,
                        help='path to pretrained NeRF parameters')
    parser.add_argument('--model', type=str, choices=['wild', 'uncertainty', 'appearance', 'vanilla'],
                        required=True,
                        help='Select the model to train')

    parser.add_argument('--visualization-type', '-v', type=str,
                        choices=['zoom', '360-rotation', 'default'],
                        default='default-render-poses',
                        help='type of visualization')

    parser.add_argument('--downscale', '-d', default=1, type=float,
                        help="downsampling factor for the rendered images for faster inference")

    parser.add_argument('--num-images', '-n', default=120, type=int,
                        help="Number of images to generate for the output video/gif")

    parser.add_argument("--fast", help="Use Fast NeRF architecture",
                        action="store_true")

    args = parser.parse_args()

    use_transient = False
    use_embedding = False

    if args.model == 'wild':
        use_transient = True
        use_embedding = True
    elif args.model == 'uncertainty':
        use_transient = True
    elif args.model == 'appearance':
        use_embedding = True

    args = parser.parse_args()
    config = read_yaml(args.config_path)

    config.data.downscale = args.downscale

    nn.set_auto_forward(True)
    ctx = get_extension_context('cuda')
    nn.set_default_context(ctx)
    nn.load_parameters(args.weight_path)

    _, _, render_poses, hwf, _, _, near_plane, far_plane = get_data(config)
    height, width, focal_length = hwf
    print(
        f'Rendering with Height {height}, Width {width}, Focal Length: {focal_length}')

    # mapping_net = MLP
    encode_position_function = get_encoding_function(
        config.train.num_encodings_position, True, True)
    if config.train.use_view_directions:
        encode_direction_function = get_encoding_function(
            config.train.num_encodings_direction, True, True)
    else:
        encode_direction_function = None

    frames = []
    if use_transient:
        static_frames = []

    if args.visualization_type == '360-rotation':
        print('The 360 degree roation result will not work with LLFF data!')
        pbar = tqdm(np.linspace(0, 360, args.num_images, endpoint=False))
    elif args.visualization_type == 'zoom':
        pbar = tqdm(np.linspace(near_plane, far_plane,
                                args.num_images, endpoint=False))
    else:
        args.num_images = min(args.num_images, render_poses.shape[0])
        pbar = tqdm(
            np.arange(0, render_poses.shape[0], render_poses.shape[0]//args.num_images))

    print(f'Rendering {args.num_images} poses...')

    for th in pbar:

        if args.visualization_type == '360-rotation':
            pose = nn.NdArray.from_numpy_array(pose_spherical(th, -30., 4.))
        elif args.visualization_type == 'zoom':
            pose = nn.NdArray.from_numpy_array(trans_t(th))
        else:
            pose = nn.NdArray.from_numpy_array(render_poses[th][:3, :4])
            # pose = nn.NdArray.from_numpy_array(render_poses[0][:3, :4])

        ray_directions, ray_origins = get_ray_bundle(
            height, width, focal_length, pose)

        ray_directions = F.reshape(ray_directions, (-1, 3))
        ray_origins = F.reshape(ray_origins, (-1, 3))

        num_ray_batches = ray_directions.shape[0]//config.train.ray_batch_size+1

        app_emb, trans_emb = None, None
        if use_embedding:
            with nn.parameter_scope('embedding_a'):
                embed_inp = nn.NdArray.from_numpy_array(
                    np.full((config.train.chunksize_fine,), 1, dtype=int))
                app_emb = PF.embed(
                    embed_inp, config.train.n_vocab, config.train.n_app)

        if use_transient:
            with nn.parameter_scope('embedding_t'):
                embed_inp = nn.NdArray.from_numpy_array(
                    np.full((config.train.chunksize_fine,), th, dtype=int))
                trans_emb = PF.embed(
                    embed_inp, config.train.n_vocab, config.train.n_trans)

            static_rgb_map_fine_list, transient_rgb_map_fine_list = [], []

        rgb_map_fine_list = []

        for i in trange(num_ray_batches):
            if i != num_ray_batches-1:
                ray_d, ray_o = ray_directions[i*config.train.ray_batch_size:(
                    i+1)*config.train.ray_batch_size], ray_origins[i*config.train.ray_batch_size:(i+1)*config.train.ray_batch_size]
            else:
                ray_d, ray_o = ray_directions[i*config.train.ray_batch_size:,
                                              :], ray_origins[i*config.train.ray_batch_size:, :]

            if use_transient:
                _, rgb_map_fine, static_rgb_map_fine, transient_rgb_map_fine, _, _, _ = forward_pass(ray_d, ray_o, near_plane, far_plane,
                                                                                                     app_emb, trans_emb, encode_position_function, encode_direction_function, config, use_transient, hwf=hwf, fast=args.fast)

                static_rgb_map_fine_list.append(static_rgb_map_fine)
                transient_rgb_map_fine_list.append(transient_rgb_map_fine)

            else:
                _, _, _, _, rgb_map_fine, _, _, _ = \
                    forward_pass(ray_d, ray_o, near_plane, far_plane, app_emb, trans_emb, encode_position_function,
                                 encode_direction_function, config, use_transient, hwf=hwf, fast=args.fast)
            rgb_map_fine_list.append(rgb_map_fine)

        rgb_map_fine = F.concatenate(*rgb_map_fine_list, axis=0)
        rgb_map_fine = F.reshape(rgb_map_fine, (height, width, 3))

        if use_transient:
            static_rgb_map_fine = F.concatenate(
                *static_rgb_map_fine_list, axis=0)
            static_rgb_map_fine = F.reshape(
                static_rgb_map_fine, (height, width, 3))

        frames.append((255*np.clip(rgb_map_fine.data, 0, 1)).astype(np.uint8))
        if use_transient:
            static_frames.append(
                (255*np.clip(static_rgb_map_fine.data, 0, 1)).astype(np.uint8))

    imageio.mimwrite(args.output_filename, frames, fps=30)
    if use_transient:
        imageio.mimwrite(args.output_static_filename, static_frames, fps=30)


if __name__ == '__main__':
    main()


# vision-and-language/clip/args.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


def get_args(monitor_path="tmp.monitor"):
    import argparse
    import os
    description = ("NNabla implementation of CLIP")
    parser = argparse.ArgumentParser(description)
    parser.add_argument('--context', '-c', type=str,
                        default=None, help="Extension modules. ex) 'cpu', 'cudnn'.")
    parser.add_argument("--type-config", "-t", type=str, default='float',
                        help='Type configuration (float or half)')
    parser.add_argument("--batch-size", "-b", type=int, default=1)
    parser.add_argument("--image-size", type=int, default=224)
    parser.add_argument("--monitor-path", "-m",
                        type=str, default=monitor_path,
                        help='Path monitoring logs saved.')
    parser.add_argument("--epochs", "-ep", type=int, default=-1,
                        help='Number of training epochs.')
    parser.add_argument("--max-iter", "-i", type=int, default=-1,
                        help='Max iterations.')
    parser.add_argument("--warmup-iter", "-wi", type=int, default=10000,
                        help='iterations for warmup.')
    parser.add_argument("--device-id", "-d", type=int, default=0,
                        help='Device ID the training run on. This is only valid if you specify `-c cudnn`.')
    parser.add_argument("--model-save-interval",
                        type=float, default=5000,
                        help='Model save interval')
    parser.add_argument("--model-save-path", "-o",
                        type=str, default=monitor_path,
                        help='Path where model parameters are saved.')
    parser.add_argument("--model-load-path",
                        type=str, default="asset/ViT-B-32.h5",
                        help='Path where model parameters are loaded.')
    parser.add_argument("--solver",
                        type=str,
                        help='Solvertype. [Adam, AdamW]')
    parser.add_argument('--lr', default=5e-4, type=float,
                        help="Learning rate")
    parser.add_argument('--beta1', default=0.9, type=float,
                        help="Beta1 of Adam")
    parser.add_argument('--beta2', default=0.98, type=float,
                        help="Beta2 of Adam")
    parser.add_argument('--eps', default=1.0e-6, type=float,
                        help="epsiron decay of Adam")
    parser.add_argument('--wd', default=1.0e-4, type=float,
                        help="weight decay of AdamW")
    parser.add_argument("--seed", default=412, type=int,
                        help='Use style encodoing in the other domain.')
    parser.add_argument("--finetuning", "-ft", action="store_true",
                        help='Finetune the pretrained model')
    parser.add_argument("--context_length", default=77, type=int,
                        help='Set context_length')
    parser.add_argument("--train-txt-path", default="asset/minisample.txt",
                        type=str,
                        help='Path to train txt file of image-text pair')
    parser.add_argument("--val-txt-path", default="asset/minisample.txt",
                        type=str,
                        help='Path to val txt file of image-text pair')
    parser.add_argument("--aggregate", "-ag", action="store_true",
                        help='Aggregate features across multiple gpus ')

    args = parser.parse_args()
    return args


def save_args(args, mode="train"):
    from nnabla import logger
    import os
    if not os.path.exists(args.monitor_path):
        os.makedirs(args.monitor_path)

    path = "{}/Arguments-{}.txt".format(args.monitor_path, mode)
    logger.info("Arguments are saved to {}.".format(path))
    with open(path, "w") as fp:
        for k, v in sorted(vars(args).items()):
            logger.info("{}={}".format(k, v))
            fp.write("{}={}\n".format(k, v))


# vision-and-language/clip/datasets.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import numpy as np
import pandas as pd
import nnabla as nn
import nnabla.functions as F
from nnabla.logger import logger
from nnabla.utils.data_iterator import data_iterator, data_iterator_simple
from nnabla.utils.data_source import DataSource

import glob
from PIL import Image


# Data interator for CLIP
# preprocess for image and text
# return [image, text]
def clip_data_iterator(clip, txt_path, batch_size=1, image_size=224, num_samples=-1,
                       normalize_method=lambda x: (x - 127.5) / 127.5,
                       shuffle=True,
                       image_key="filepath",
                       caption_key="title",
                       rng=None):
    # Load txt file (image_path and ception)
    df = pd.read_csv(txt_path, sep="\t")
    images = df[image_key].tolist()
    captions = df[caption_key].tolist()

    if num_samples == -1:
        assert len(images) == len(
            captions), "The number of images and captions must be same"
        num_samples = len(images)
    else:
        logger.info(
            "Num. of data ({}) is used for debugging".format(num_samples))

    def load_func(i):
        text = captions[i]
        # Note: change prefix as per your environment
        img_path = "/conceptual_captions/cc_3m/" + images[i]
        image = clip.preprocess(Image.open(img_path))
        # shape [num_text_per_image, 77] but in training phase, num_text_per_image=1
        text = clip.tokenize(text, truncate=True)
        text = np.reshape(text, (77,))

        # Return numpy.ndarray, not Variable
        return image, text

    return data_iterator_simple(load_func, num_samples, batch_size,
                                shuffle=shuffle, rng=rng, with_file_cache=False)


if __name__ == '__main__':
    # Hand-made test
    import clip
    clip.load("asset/ViT-B-32.h5")
    txt_path = "asset/minisample.txt"
    loader = clip_data_iterator(clip, txt_path, batch_size=1, image_size=256)

    for _ in range(4):
        imgs, texts = loader.next()
        print(imgs[0].shape, texts[0])


# vision-and-language/clip/demo_zero_shot_classification.py

import nnabla as nn
import numpy as np

from typing import Union, Optional, List
from io import BytesIO

from clip import CLIP

PRESET_HAIR = '\n'.join(f'A person with {color} hair' for color in [
                        'white', 'black', 'brown', 'red', 'pink', 'blue', 'blond'])
PRESET_EXPRESSION = '\n'.join(f'A {expression} person' for expression in [
                              'happy', 'sad', 'calm', 'angry', 'funny', 'crying'])
IMAGE_EXAMPLES = [
    'https://images.pexels.com/photos/5682847/pexels-photo-5682847.jpeg?auto=compress&cs=tinysrgb&w=600',
    'https://images.pexels.com/photos/2741036/pexels-photo-2741036.jpeg?auto=compress&cs=tinysrgb&w=600',
    'https://images.pexels.com/photos/1852300/pexels-photo-1852300.jpeg?auto=compress&cs=tinysrgb&w=600',
    'https://images.pexels.com/photos/13198518/pexels-photo-13198518.jpeg?auto=compress&cs=tinysrgb&w=600',
]


def read_image_from_bytes(b: BytesIO) -> np.ndarray:
    from PIL import Image
    pil_img = Image.open(b)
    np_img = np.asarray(pil_img, dtype=np.uint8)
    return np_img


def download_image(url: str, return_bytes: bool = False) -> Union[np.ndarray, BytesIO]:
    import requests
    # Download and read image as numpy array
    c = requests.get(url).content
    b = BytesIO(c)
    if return_bytes:
        return b
    np_img = read_image_from_bytes(b)
    return np_img


def parse_categories(texts: str) -> List[str]:
    return list(
        filter(
            lambda x: x,
            map(lambda x: x.strip(), texts.splitlines())))


class GradioCLIP(CLIP):

    def run_by_url(self, url: str, categories: str):
        img_bytes = download_image(url, return_bytes=True)
        np_img = read_image_from_bytes(img_bytes)
        categories = parse_categories(categories)
        img_bytes.seek(0)
        probs = self.run(img_bytes, categories)
        inds = probs.argsort()[::-1]
        labels = {categories[i]: float(probs[i]) for i in inds}
        return np_img, labels


def main():

    import gradio as gr
    from nnabla.ext_utils import get_extension_context
    ctx = get_extension_context('cudnn')
    nn.set_default_context(ctx)

    # Create layout
    with gr.Blocks() as demo:
        w_md = gr.Markdown('''
        # Zero-shot classification using nnabla CLIP

        1. Provide URL of an input image.
        2. Define categories for zero-shot classification. See category definition examples below. Each line reprsents a category.
        3. Click the "Run" button to run predicted scores.
        ''')
        with gr.Row():
            with gr.Column(scale=2):
                w_url = gr.Textbox(label='Image URL', lines=1,
                                   max_lines=1, value=IMAGE_EXAMPLES[0])
                w_categories = gr.Textbox(label='Zero-shot categories', lines=10, max_lines=100,
                                          value=PRESET_HAIR)
                w_predict = gr.Button(label='Run')
            with gr.Column(scale=3):
                w_image = gr.Image(label='Image')
            with gr.Column(scale=4):
                w_label = gr.Label(label='Prediction')

        # Set callbacks
        gradio_clip = GradioCLIP()
        w_preset_categories = gr.Examples(
            [PRESET_HAIR, PRESET_EXPRESSION],
            [w_categories],
            label='Category definition examples',
        )
        w_preset_images = gr.Examples(
            [[x] for x in IMAGE_EXAMPLES],
            inputs=[w_url],
            label='Example images',
        )

        w_url.change(download_image, inputs=[w_url], outputs=w_image)
        w_predict.click(gradio_clip.run_by_url, inputs=[
                        w_url, w_categories], outputs=[w_image, w_label])

    # Launch servers
    demo.queue(concurrency_count=1)
    demo.launch(server_name='0.0.0.0', server_port=8889)


if __name__ == '__main__':
    main()


# vision-and-language/clip/losses.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import nnabla as nn
import nnabla.functions as F

import numpy as np


def get_logits(image_features, text_features, aggregate=False, checking=nn.Variable()):

    # normalized features
    image_features = image_features / \
        F.norm(image_features, axis=1, keepdims=True)
    text_features = text_features / \
        F.norm(text_features, axis=1, keepdims=True)

    # cosine similarity as logits
    logit_scale = nn.parameter.get_parameter_or_create(
        name='logit_scale', shape=())
    checking = logit_scale
    logit_scale = F.mean(F.exp(logit_scale))

    image_features = image_features.reshape(
        (1, image_features.shape[0], image_features.shape[1]))
    text_features = F.transpose(text_features, (1, 0))
    text_features = text_features.reshape(
        (1, text_features.shape[0], text_features.shape[1]))

    per_image = F.batch_matmul(image_features, text_features).reshape(
            (image_features.shape[1], image_features.shape[1])
        )  # shape = [global_batch_size, global_batch_size]
    logits_per_image = logit_scale.reshape((1, 1)) * per_image

    logits_per_text = F.transpose(logits_per_image, (1, 0))

    # shape = [global_batch_size, global_batch_size]
    return logits_per_image, logits_per_text, checking


# vision-and-language/clip/main.py

# Copyright 2021 Sony Group Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import nnabla as nn
import nnabla.functions as F
from nnabla.ext_utils import get_extension_context
from PIL import Image

import clip

ctx = get_extension_context('cudnn')
nn.set_default_context(ctx)


def main():
    with nn.auto_forward():
        clip.load('data/ViT-B-32.h5')

        image = nn.Variable.from_numpy_array(
            clip.preprocess(Image.open("CLIP.png")))
        image = F.reshape(
            image, (1, image.shape[0], image.shape[1], image.shape[2]))
        text = nn.Variable.from_numpy_array(
            clip.tokenize(["a diagram", "a dog", "a cat"]))

        image_features = clip.encode_image(image)
        text_features = clip.encode_text(text)

        logits_per_image, logits_per_text = clip.logits(image, text)
        probs = F.softmax(logits_per_image, axis=-1)

        # prints: [[0.9927937  0.00421068 0.00299572]]
        print("Label probs:", probs.d)


if __name__ == "__main__":
    main()


# vision-and-language/clip/scheduler.py

'''
Forked from https://github.com/sony/nnabla-examples/blob/master/utils/neu/learning_rate_scheduler.py
'''

import nnabla


class BaseLearningRateScheduler(object):

    def __init__(self):
        self._iter = 0
        self._epoch = 0

    def set_iter_per_epoch(self, it):
        pass

    def set_epoch(self, epoch):
        '''Set current epoch number.
        '''
        self._epoch = epoch

    def get_lr_and_update(self):
        '''
        Get current learning rate and update itereation count.
        The iteration count is calculated by how many times this method is called.
        Returns: Current learning rate
        '''
        lr = self.get_lr()
        self._iter += 1
        return lr

    def get_lr(self):
        '''
        Get current learning rate according to the schedule.
        '''
        return self._get_lr(self._epoch, self._iter)

    def _get_lr(self, current_epoch, current_iter):
        '''
        Get learning rate by current iteration.
        Args:
            current_epoch(int): Epoch count.
            current_iter(int):
                Current iteration count from the beginning of training.
        Note:
            A derived class must override this method. 
        '''
        raise NotImplementedError('')


class IterCosineLearningRateScheduler(BaseLearningRateScheduler):
    '''
    ** Modified `EpochCosineLearningRateScheduler` **

    Cosine Annealing Decay with warmup.
    The learning rate gradually increases linearly towards `base_lr` during
    `warmup_iters`, then gradually decreases with cosine decay towards 0 for
    `iters - warmup_iters`.
    Args:
        base_lr (float): Base learning rate
        iters (int): See description above.
        warmup_iters (int): It performs warmup during this period.
    '''

    def __init__(self, base_lr, iters, warmup_iters=5):

        # https://github.com/sony/nnabla/blob/master/python/src/nnabla/utils/learning_rate_scheduler.py
        from nnabla.utils.learning_rate_scheduler import CosineScheduler
        super().__init__()
        self.base_lr = base_lr
        self.iters = iters  # total iterations for training
        self.warmup_iters = warmup_iters
        self.cosine = CosineScheduler(
            self.base_lr, self.iters - self.warmup_iters)  # (init_lr, max_iter)

    def _get_lr(self, current_epoch, current_iter):

        # Warmup
        if current_iter < self.warmup_iters:
            return self.base_lr * (current_iter + 1) / self.warmup_iters

        # Cosine decay
        return self.cosine.get_learning_rate(
            current_iter - self.warmup_iters)


# vision-and-language/clip/train.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import os
import numpy as np
import argparse
import time
import datetime
import tqdm
import nnabla as nn
import nnabla.logger as logger
import nnabla.functions as F
import nnabla.parametric_functions as PF
import nnabla.solvers as S
import nnabla.communicators as C

from nnabla.monitor import Monitor, MonitorImage, MonitorSeries, MonitorTimeElapsed
from nnabla.ext_utils import get_extension_context
from nnabla.utils.profiler import GraphProfiler

from collections import OrderedDict

from args import get_args, save_args
from datasets import clip_data_iterator
from losses import get_logits
from scheduler import IterCosineLearningRateScheduler

import clip


def train(args):
    # Create Communicator and Context
    extension_module = "cudnn"
    ctx = get_extension_context(extension_module, type_config=args.type_config)
    comm = C.MultiProcessDataParalellCommunicator(ctx)
    comm.init()
    n_devices = comm.size
    mpi_rank = comm.rank
    mpi_local_rank = comm.local_rank
    device_id = mpi_local_rank
    ctx.device_id = str(device_id)
    nn.set_default_context(ctx)
    np.random.seed(args.seed)

    # Load checkpoint if file exists.
    if args.finetuning:
        # load nnabla params
        clip_model_path = args.model_load_path  # ex) 'asset/ViT-B-32.h5'
        clip.load(clip_model_path)
        print(f"Fine-tuning from {clip_model_path}.")
    else:
        # start training from initialized params
        clip_model_path = 'asset/ViT-B-32-initialized.h5'
        clip.load(clip_model_path)
        print(f"Training from {clip_model_path}.")

    # Create training graphs
    # Input
    b, c, h, w = args.batch_size, 3, args.image_size, args.image_size
    image = nn.Variable([b, c, h, w])
    text = nn.Variable([b, args.context_length])  # context_length = 77.

    # Graph
    image_features = clip.encode_image(image)
    text_features = clip.encode_text(text)
    image_features.persistent = True
    text_features.persistent = True

    if args.aggregate:
        assert comm is not None, "Set communicator"
        # gather all tensor across multiple gpu for a large matrix
        i_features_list = [nn.Variable(image_features.shape)
                           for _ in range(comm.size)]
        t_features_list = [nn.Variable(text_features.shape)
                           for _ in range(comm.size)]
        unliked_image_features = image_features.get_unlinked_variable()
        unliked_text_features = text_features.get_unlinked_variable()
        print(
            f"All Gather Executed: comm.size={comm.size}, comm.rank={comm.rank}, comm.local_rank={comm.local_rank}")
        comm.all_gather(unliked_image_features.data, [
                        i_features.data for i_features in i_features_list])
        comm.all_gather(unliked_text_features.data, [
                        t_features.data for t_features in t_features_list])

        # aggregate features_list to one features
        all_image_features = F.concatenate(
            unliked_image_features, *i_features_list[:comm.rank], *i_features_list[comm.rank+1:], axis=0)
        all_text_features = F.concatenate(
            unliked_text_features, *t_features_list[:comm.rank], *t_features_list[comm.rank+1:], axis=0)
        print(
            f"aggregated to image_features:{all_image_features.shape}, text_features:{all_text_features.shape}")
        logits_per_image, logits_per_text, logit_scale = get_logits(
            all_image_features, all_text_features, aggregate=args.aggregate)
        all_image_features.persistent = True
        all_text_features.persistent = True
    else:
        logits_per_image, logits_per_text, logit_scale = get_logits(
            image_features, text_features, aggregate=args.aggregate)
        logit_scale.persistent = True

    # Get ground_truth
    ground_truth_d = np.arange(
        logits_per_image.shape[0], dtype="int").reshape(-1, 1)
    ground_truth = nn.Variable((logits_per_image.shape[0], 1))

    # Cacl losses
    loss_image = F.mean(F.softmax_cross_entropy(
        logits_per_image, ground_truth))
    loss_image.persistent = True
    loss_text = F.mean(F.softmax_cross_entropy(logits_per_text, ground_truth))
    loss_text.persistent = True
    total_loss = (loss_image + loss_text) / 2
    total_loss.persistent = True

    # Create validation graphs
    # Input
    image_val = nn.Variable([b, c, h, w])
    # context_length = 77. Should be fixed following the official CLIP implement.
    text_val = nn.Variable([b, args.context_length])
    # Graph
    image_features_val = clip.encode_image(image_val)
    text_features_val = clip.encode_text(text_val)
    logits_per_image_val, logits_per_text_val, _ = get_logits(
        image_features_val, text_features_val, aggregate=args.aggregate)
    # Get ground_truth
    ground_truth_d_val = np.arange(
        logits_per_image_val.shape[0], dtype="int").reshape(-1, 1)
    ground_truth_val = nn.Variable((logits_per_image_val.shape[0], 1))
    # Culc losses
    loss_image_val = F.mean(F.softmax_cross_entropy(
        logits_per_image_val, ground_truth_val))
    loss_text_val = F.mean(F.softmax_cross_entropy(
        logits_per_text_val, ground_truth_val))
    total_loss_val = (loss_image_val + loss_text_val) / 2

    # Solver
    # Divide parameters into two groups mainly for AdamW, whose gain and bias are unsuitable for weight decay.
    def exclude(
        n, p): return p.ndim < 2 or "/bn" in n or "/ln" in n or "/b" in n or 'logit_scale' in n

    def include(n, p): return not exclude(n, p)

    named_parameters = nn.get_parameters()

    gain_or_bias_params = OrderedDict(
        [(n, p) for n, p in named_parameters.items() if exclude(n, p) and p.need_grad])
    rest_params = OrderedDict(
        [(n, p) for n, p in named_parameters.items() if include(n, p) and p.need_grad])

    if args.solver == "Adam":
        solver1 = S.Adam(args.lr, args.beta1, args.beta2, args.eps)
        solver2 = S.Adam(args.lr, args.beta1, args.beta2, args.eps)
    elif args.solver == "AdamW":
        solver1 = S.AdamW(args.lr, args.beta1, args.beta2, args.eps, 0.0)
        # ex) --lr 5.0e-4 --beta1 0.9 --beta2 0.98 --eps 1.0e-6 --wd 1.0e-4
        solver2 = S.AdamW(args.lr, args.beta1, args.beta2, args.eps, args.wd)
    else:
        raise NotImplementedError()
    solver1.set_parameters(gain_or_bias_params)
    solver2.set_parameters(rest_params)

    # Monitor
    monitor_interval = 100 * args.batch_size * n_devices
    monitor = Monitor(args.monitor_path)
    monitor_time = MonitorTimeElapsed(
        "Training time", monitor, interval=monitor_interval)
    monitor_loss_image = MonitorSeries(
        "Loss Image", monitor, interval=monitor_interval)
    monitor_loss_text = MonitorSeries(
        "Loss Text", monitor, interval=monitor_interval)
    monitor_total_loss = MonitorSeries(
        "Total Loss", monitor, interval=monitor_interval)
    monitor_logit_scale = MonitorSeries(
        "Logit Scale", monitor, interval=monitor_interval)
    monitor_losses = [
        (monitor_loss_image, loss_image),
        (monitor_loss_text, loss_text),
        (monitor_total_loss, total_loss),
        (monitor_logit_scale, logit_scale)
    ]
    # Monitor for eval
    monitor_loss_val = MonitorSeries(
        "Eval Loss", monitor, interval=1)  # monitoring each_epoch
    monitor_val_loss = [
        (monitor_loss_val, total_loss_val),
    ]

    # Data Iterator
    rng = np.random.RandomState(args.seed)
    train_loader = clip_data_iterator(
        clip, args.train_txt_path, args.batch_size, shuffle=True, rng=rng)
    print(f"train_loader.size: {train_loader.size}")
    if comm.size > 1:
        train_loader = train_loader.slice(
            rng=None, num_of_slices=comm.size, slice_pos=comm.rank)
        print(
            f"train_loader.size: {train_loader.size}, num_of_slices={comm.size}, slice_pos={comm.rank}")
    val_loader = clip_data_iterator(
        clip, args.val_txt_path, batch_size=args.batch_size, shuffle=False)
    iters_per_one_epoch = train_loader.size // args.batch_size

    comm.barrier()

    # set max iteration
    if args.max_iter == -1:
        args.max_iter = args.epochs * iters_per_one_epoch
        print(f"Max iteration = {args.max_iter}")

    # Learning rate scheduler
    learning_rate_scheduler = IterCosineLearningRateScheduler(
        base_lr=args.lr,
        iters=args.max_iter,
        warmup_iters=args.warmup_iter
    )

    comm.barrier()

    # Train loop
    for i in range(args.max_iter):
        ii = i * args.batch_size * n_devices
        current_epoch = i // iters_per_one_epoch

        image_data, text_data = train_loader.next()
        image.d, text.d = image_data, text_data
        ground_truth.d = ground_truth_d

        if args.aggregate:
            image_features.forward(clear_no_need_grad=True)
            text_features.forward(clear_no_need_grad=True)

            comm.all_gather(unliked_image_features.data, [
                            i_features.data for i_features in i_features_list])
            comm.all_gather(unliked_text_features.data, [
                            t_features.data for t_features in t_features_list])

            unliked_image_features.grad.zero()
            unliked_text_features.grad.zero()

        total_loss.forward(clear_no_need_grad=True)

        solver1.zero_grad()
        solver2.zero_grad()
        solver1.set_learning_rate(
            learning_rate_scheduler._get_lr(
                current_epoch=current_epoch, current_iter=i)
        )
        solver2.set_learning_rate(
            learning_rate_scheduler._get_lr(
                current_epoch=current_epoch, current_iter=i)
        )

        total_loss.backward(clear_buffer=True)

        if args.aggregate:
            image_features.backward(grad=None, clear_buffer=True)
            text_features.backward(grad=None, clear_buffer=True)
            comm.all_reduce([w.grad for w in gain_or_bias_params.values()])
            comm.all_reduce([w.grad for w in rest_params.values()])

        solver1.update()
        solver2.update()

        if mpi_local_rank == 0:
            # Monitor
            monitor_time.add(ii)
            for mon, loss in monitor_losses:
                mon.add(ii, loss.d)
            # Save Params and evaluate validation loss, at the end of each epoch
            if (i + 1) % iters_per_one_epoch == 0:
                # save
                nn.save_parameters(os.path.join(
                    args.monitor_path, "param_epoch_{:05d}.h5".format(current_epoch)))
                # val loss
                val_loss = []
                for vi in tqdm.tqdm(range(val_loader.size // args.batch_size)):
                    image_data_val, text_data_val = val_loader.next()
                    image_val.d, text_val.d = image_data_val, text_data_val
                    ground_truth_val.d = ground_truth_d_val
                    total_loss_val.forward(clear_buffer=True)
                    val_loss.append(total_loss_val.d)
                val_loss = np.array(val_loss).mean()
                for mon, loss in monitor_val_loss:
                    mon.add(current_epoch, val_loss)

        if (i + 1) % iters_per_one_epoch == 0:
            # re-shuffle for nnabla v1.27 or v1.28 by changing random seeds
            train_loader = clip_data_iterator(clip, args.train_txt_path, args.batch_size,
                                              shuffle=True, rng=np.random.RandomState(args.seed + (current_epoch+1)))
            if comm.size > 1:
                train_loader = train_loader.slice(
                    rng=None, num_of_slices=comm.size, slice_pos=comm.rank)

    if mpi_local_rank == 0:
        # Monitor
        for mon, loss in monitor_losses:
            mon.add(ii, loss.d)
        # Save
        nn.save_parameters(os.path.join(
            args.monitor_path, "param_iter_{:05d}.h5".format(i)))
        # val loss
        val_loss = []
        for vi in tqdm.tqdm(range(val_loader.size // args.batch_size)):
            image_data_val, text_data_val = val_loader.next()
            image_val.d, text_val.d = image_data_val, text_data_val
            ground_truth_val.d = ground_truth_d_val
            total_loss_val.forward(clear_buffer=True)
            val_loss.append(total_loss_val.d)
        val_loss = np.array(val_loss).mean()
        for mon, loss in monitor_val_loss:
            mon.add(current_epoch, val_loss)


def main():
    d = datetime.datetime.now()
    monitor_path = d.strftime(f"tmp.monitor/%Y-%m-%d-%H-%M-%S")
    args = get_args(monitor_path)
    save_args(args)
    train(args)


if __name__ == '__main__':
    main()


# vision-and-language/clip/weight_convert_nn2pth.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import numpy as np
np.set_printoptions(precision=32, floatmode='fixed')

import torch
import nnabla as nn
import nnabla.parametric_functions as PF

print("nnabla version:", nn.__version__)
print("Torch version:", torch.__version__)

from collections import OrderedDict
import pickle


def param_map():
    '''map to tensor name from Nnabla default parameter names
    '''
    return {
        '/W': '/weight',
        '_W': '_weight',
        '/b': '/bias',
        '_b': '_bias',
        '/': '.'
    }


def rename_params(param_name):
    map_dict = param_map()
    for k in map_dict:
        if k in param_name:
            param_name = param_name.replace(k, map_dict[k])
    return param_name


def convert(pth_file, nn_h5, prefix=""):
    nn.clear_parameters()
    nn.load_parameters(nn_h5)
    state_dict = OrderedDict()
    named_params = nn.get_parameters()
    for k, v in named_params.items():
        print(k)
        key = rename_params(k)
        value = torch.from_numpy(v.d)
        state_dict[prefix + key] = value

    torch.save(state_dict, pth_file)


if __name__ == '__main__':
    pth_file = "./convert.pth"  # output file
    # input file to be converted
    nn_h5 = "./tmp.monitor/2022-04-20-09-03-39/param_epoch_00029.h5"
    prefix = "module."  # set prefix for each key

    convert(pth_file, nn_h5, prefix)


# vision-and-language/cliport/setup.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from setuptools import setup, find_packages

setup(
    name='nnabla_cliport',
    version='0.0.1',
    packages=find_packages(),
    include_package_data=True,
    license='Apache License 2.0',
    zip_safe=False,
    description="CLIPort implementation using nnabla.",
    author='Yu Ishihara',
    author_email='yu.ishihara@sony.com',
    url='https://github.com/sony/nnabla-examples/'
)


# responsible_ai/adversarial_debiasing_images/utils/utils.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
from sklearn.metrics import average_precision_score, f1_score, recall_score
from scipy.ndimage import gaussian_filter
import matplotlib.pyplot as plt


def compute_class_weight(loader):
    """
    Compute sample weights
    """

    cp = 0
    cn = 0
    cn_dn = 0
    cn_dp = 0
    cp_dn = 0
    cp_dp = 0

    weights = []
    print("size", loader.size)
    max_iter = int(loader.size / 32)
    for i in range(max_iter):
        _, targets = loader.next()
        class_label = targets[:, 0]
        domain_label = targets[:, 1]
        cp += class_label.sum()  # class is positive
        cn += (targets.shape[0] - class_label.sum())  # class is negative
        # class is negative, domain is negative
        cn_dn += ((class_label + domain_label) == 0).sum()
        cn_dp += ((class_label - domain_label) == -1).sum()
        cp_dn += ((class_label - domain_label) == 1).sum()
        cp_dp += ((class_label + domain_label) == 2).sum()

    for i in range(max_iter):
        _, targets = loader.next()
        class_label = targets[:, 0]
        domain_label = targets[:, 1]

        weights.append(
            (class_label * cp + (1 - class_label) * cn) /
            (2 * (
                    (1 - class_label) * (1 - domain_label) * cn_dn
                    + (1 - class_label) * domain_label * cn_dp
                    + class_label * (1 - domain_label) * cp_dn
                    + class_label * domain_label * cp_dp
            )
             )
        )
    weights = np.concatenate(weights)
    return weights


def get_bias_amplification_attribute_task(running_labels, running_preds):
    """
    compute bias amplification for attribute-task
    References:
        [1] Directional Bias Amplification(https://arxiv.org/abs/2102.12594)
        [2] Fair Attribute Classification through Latent Space
        De-biasing(https://arxiv.org/abs/2012.01469)
    Args:
        running_labels : labels on test set
        running_preds : predictions on test set
    Returns:
         bias amplification
    """
    bog_tilde = np.zeros((2, 2))
    bog_gt_g = np.zeros((2, 2))
    for i, objs in enumerate([running_labels, running_preds]):
        female = np.where(objs[:, 1] == 0)[0]  # Unprivileged class
        male = np.where(objs[:, 1] == 1)[0]  # Privileged class
        kitchen = np.where(objs[:, 0] == 0)[0]  # Unfavourable class
        sports = np.where(objs[:, 0] == 1)[0]  # Favourable class
        if i == 0:
            bog_tilde[0][0] = len(set(kitchen) & set(female))
            bog_tilde[0][1] = len(set(kitchen) & set(male))
            bog_tilde[1][0] = len(set(sports) & set(female))
            bog_tilde[1][1] = len(set(sports) & set(male))
        elif i == 1:
            bog_gt_g[0][0] = len(set(kitchen) & set(female))
            bog_gt_g[0][1] = len(set(kitchen) & set(male))
            bog_gt_g[1][0] = len(set(sports) & set(female))
            bog_gt_g[1][1] = len(set(sports) & set(male))

    total_images_train = np.sum(bog_tilde)
    data_bog = bog_tilde / np.sum(bog_tilde, axis=0)
    pred_bog = bog_gt_g / np.sum(bog_tilde, axis=0)
    p_t_a = bog_tilde / np.sum(bog_tilde, axis=0)
    p_t = np.sum(bog_tilde, axis=1) / total_images_train
    diff = np.zeros_like(data_bog)
    for i in range(len(data_bog)):
        for j in range(len(data_bog[0])):
            diff[i][j] = pred_bog[i][j] - data_bog[i][j]
            indicator = np.sign(p_t_a[i][j] - p_t[i])
            if indicator == 0:
                diff[i][j] = 0
            elif indicator == -1:
                diff[i][j] = - diff[i][j]
    value = np.nanmean(diff)
    return value


def get_bias_amplification(domain, targets, pred):
    """
    BA measures how much more often a target attribute is predicted with a
    protected attribute than the ground truth value.
    References:
        [1] Directional Bias Amplification(https://arxiv.org/abs/2102.12594)
        [2] Fair Attribute Classification through Latent Space
        De-biasing(https://arxiv.org/abs/2012.01469)
    Args:
        domain (numpy.ndarray) : actual protected attribute
        targets (numpy.ndarray) : actual target label
        pred (numpy.ndarray) : predicted label
    Returns:
        Bias amplification
    """

    test_labels = np.zeros((targets.shape[0], 2))
    test_labels[:, 0] = targets
    test_labels[:, 1] = domain
    test_pred = np.zeros((targets.shape[0], 2))
    test_pred[:, 0] = pred
    test_pred[:, 1] = domain
    auc_bias = get_bias_amplification_attribute_task(test_labels, test_pred)
    return auc_bias


def get_difference_equality_opportunity(domain, targets, pred):
    """
    Compute the absolute difference in FNR between the protected attribute group.
    Args:
        domain (numpy.ndarray) : actual protected attribute
        targets (numpy.ndarray) : actual target label
        pred (numpy.ndarray) : predicted label
    Returns:
        difference_in_equla_opportunity(float)
    """

    g0 = np.argwhere(domain == 0)
    g1 = np.argwhere(domain == 1)
    deo = np.abs(
        (1 - recall_score(targets[g0], pred[g0])) - (1 - recall_score(targets[g1], pred[g1])))

    return np.median(deo)


def get_average_precision(targets, scores):
    """
    Compute average precision score
    Args:
        targets (numpy.ndarray) : actual target label
        scores (numpy.ndarray) : predicted scores
    Returns:
        Average precission
    """

    avg_prec = average_precision_score(targets, scores)
    print(avg_prec)

    return np.median(avg_prec)


def get_f1_threshold(targets_all, scores_all):
    """
    get the f1 threshold and accuracy
    Args:
        targets_all (numpy.ndarray) : actual target label
        scores : predicted scores
    Returns:
        best_acc : best accuracy
        best_t : best threshold
    """
    best_t = -1.0
    best_acc = 0.0
    for t in range(1, 10):
        thresh = 0.1 * t
        curr_scores = np.where(scores_all > thresh, 1, 0)
        acc = f1_score(targets_all, curr_scores)
        if acc > best_acc:
            best_acc = acc
            best_t = thresh
    one_dec = best_t

    for t in range(1, 20):
        thresh = (one_dec - 0.1) + 0.01 * t
        curr_scores = np.where(scores_all > thresh, 1, 0)
        acc = f1_score(targets_all, curr_scores)
        # print(thresh, acc, best_acc, flush=True)
        if acc > best_acc:
            best_acc = acc
            best_t = thresh

    return best_acc, best_t


def calibrated_threshold(targets, scores):
    """
    Calibrated threshold
    Args:
        targets (numpy.ndarray) : actual target label
        scores (numpy.ndarray) : predicted scores
    Returns:
        calibrated threshold
    """
    cp = int(targets.sum())
    scores_copy = np.copy(scores)
    scores_copy.sort()
    thresh = scores_copy[-cp]
    return thresh


def get_kl(p, q):
    """
    Kullback-Leibler divergence D(P || Q) for discrete distributions
    Args:
         p,q (float array): Discrete probability distributions.
    Returns:
        kl
    """
    p = np.asarray(p, dtype=np.float)
    q = np.asarray(q, dtype=np.float)
    return np.sum(np.where(p != 0, p * np.log(p / q), 0))


def smoothed_hist_kl_distance(a, b, nbins=50, sigma=1):
    """
    smooth two histograms with Standard kernel (gaussian kernel with mean = 0 ,sigma=1),
    and calculate the KL distance between two smoothed histograms.
    Args:
        a (numpy.ndarray) : numpy nd array
        b (numpy.ndarray) : numpy nd array
        nbins (int) : number of histogram bins
        sigma : standard deviation for Gaussian kernel
    Returns:
        kl distance
    """
    ahist = np.histogram(a, bins=nbins)[0]
    bhist = np.histogram(b, bins=nbins)[0]

    asmooth = gaussian_filter(ahist, sigma)
    bsmooth = gaussian_filter(bhist, sigma)

    asmooth = asmooth / asmooth.sum() + 1e-6
    bsmooth = bsmooth / bsmooth.sum() + 1e-6

    return get_kl(asmooth, bsmooth), get_kl(bsmooth, asmooth)


def get_kl_divergence(domain, targets, scores):
    """
    Measures the divergence between score distributions (KL/threshold-invariant metric)
    References:
        [1] Towards Threshold Invariant Fair Classification(https://arxiv.org/abs/2102.12594)
    Args:
        domain (numpy.ndarray) : actual protected attribute
        targets (numpy.ndarray) : actual target label
        scores (numpy.ndarray) : predicted scores
    Returns:
        kl_divergence
    """

    nbin = 50  # Number of histogram bins
    a_b_pos, b_a_pos = smoothed_hist_kl_distance(scores[np.logical_and(domain == 0, targets == 1)],
                                                 scores[np.logical_and(
                                                     domain == 1, targets == 1)],
                                                 nbins=nbin)
    a_b_neg, b_a_neg = smoothed_hist_kl_distance(scores[np.logical_and(domain == 0, targets == 0)],
                                                 scores[np.logical_and(
                                                     domain == 1, targets == 0)],
                                                 nbins=nbin)

    return np.median([a_b_pos] + [a_b_neg] + [b_a_pos] + [b_a_neg])


def plot_fairness(fairness, ax, metric="DPD", title="fairness metric", bar_x_axis="Original"):
    """
    plot single fairness metric
    """
    ax.set_ylim([-0.6, 0.6])
    ax.axhline(y=0.0, color='r', linestyle='-')
    ax.bar([bar_x_axis], fairness, color="blue", width=2)
    ax.set_ylabel(metric)
    ax.set_title(title, fontsize=10)
    for index, value in enumerate(fairness):
        if value < 0:
            ax.text(index, value - 0.1,
                    str(round(value, 3)), fontweight='bold', color='red',
                    bbox=dict(facecolor='red', alpha=0.4))
        else:
            ax.text(index, value + 0.1,
                    str(round(value, 3)), fontweight='bold', color='red',
                    bbox=dict(facecolor='red', alpha=0.4))


def plot_fairness_multi(DEO, BA, AAOD, accuracy, bar_x_axis="original"):
    """
    plot fairness metrics
    """
    fig, axes = plt.subplots(1, 3, figsize=(10, 4), sharey=True)
    fig.suptitle("Model Fairness", fontsize=16)
    plot_fairness([DEO], ax=axes[0], metric="DEO",
                  title="Difference in Equal opportunity(DEO)", bar_x_axis=bar_x_axis)
    plot_fairness([BA], ax=axes[1], metric="BA",
                  title="Bias Amplification (BA)", bar_x_axis=bar_x_axis)
    plot_fairness([AAOD], ax=axes[2], metric="KL",
                  title="KL divergence", bar_x_axis=bar_x_axis)
    fig.text(0.92, 0.65, '\n'.join(
        ["Average precession score:", f"- AP : {accuracy:.3f}"]), fontsize='15')
    plt.show()


# responsible_ai/data_cleansing/datasets/create_cifar10_csv.py

# Copyright 2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''
Provide data iterator for CIFAR10 examples.
'''
from contextlib import contextmanager
import argparse
import numpy as np
import tarfile
import os
import tqdm
from imageio import imwrite
from nnabla.logger import logger
from nnabla.utils.data_iterator import data_iterator
from nnabla.utils.data_source import DataSource
from nnabla.utils.data_source_loader import download
from .utils import get_filename_to_download, save_list_to_csv, split_data_into_train_val, ensure_dir


class Cifar10DataSource(DataSource):
    '''
    Get data directly from cifar10 dataset from Internet(yann.lecun.com).
    '''

    def _get_data(self, position):
        image = self._images[self._indexes[position]]
        label = self._labels[self._indexes[position]]
        return (image, label)

    def __init__(self, train=True, shuffle=False, rng=None, output_dir=None):
        super(Cifar10DataSource, self).__init__(shuffle=shuffle)

        self._train = train
        data_uri = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"
        logger.info('Getting labeled data from {}.'.format(data_uri))
        output_file = get_filename_to_download(output_dir, data_uri)
        r = download(data_uri, output_file=output_file)  # file object returned
        with tarfile.open(fileobj=r, mode="r:gz") as fpin:
            # Training data
            if train:
                images = []
                labels = []
                for member in fpin.getmembers():
                    if "data_batch" not in member.name:
                        continue
                    fp = fpin.extractfile(member)
                    data = np.load(fp, allow_pickle=True, encoding="bytes")
                    images.append(data[b"data"])
                    labels.append(data[b"labels"])
                self._size = 50000
                self._images = np.concatenate(
                    images).reshape(self._size, 3, 32, 32)
                self._labels = np.concatenate(labels).reshape(-1, 1)
            # Validation data
            else:
                for member in fpin.getmembers():
                    if "test_batch" not in member.name:
                        continue
                    fp = fpin.extractfile(member)
                    data = np.load(fp, allow_pickle=True, encoding="bytes")
                    images = data[b"data"]
                    labels = data[b"labels"]
                self._size = 10000
                self._images = images.reshape(self._size, 3, 32, 32)
                self._labels = np.array(labels).reshape(-1, 1)
        r.close()
        logger.info('Getting labeled data from {}.'.format(data_uri))

        self._size = self._labels.size
        self._variables = ('x', 'y')
        if rng is None:
            rng = np.random.RandomState(313)
        self.rng = rng
        self.reset()

    def reset(self):
        if self._shuffle:
            self._indexes = self.rng.permutation(self._size)
        else:
            self._indexes = np.arange(self._size)
        super(Cifar10DataSource, self).reset()

    @property
    def images(self):
        """Get copy of whole data with a shape of (N, 1, H, W)."""
        return self._images.copy()

    @property
    def labels(self):
        """Get copy of whole label with a shape of (N, 1)."""
        return self._labels.copy()


@contextmanager
def data_iterator_cifar10(batch_size,
                          train=True,
                          rng=None,
                          shuffle=True,
                          with_memory_cache=False,
                          with_file_cache=False,
                          output_dir=None):
    '''
    Provide DataIterator with :py:class:`Cifar10DataSource`
    with_memory_cache, with_parallel and with_file_cache option's default value is all False,
    because :py:class:`Cifar10DataSource` is able to store all data into memory.

    For example,

    .. code-block:: python

        with data_iterator_cifar10(True, batch_size) as di:
            for data in di:
                SOME CODE TO USE data.

    '''
    with Cifar10DataSource(train=train, shuffle=shuffle, rng=rng, output_dir=output_dir) as ds, \
        data_iterator(ds,
                      batch_size,
                      rng=rng,
                      with_memory_cache=with_memory_cache,
                      with_file_cache=with_file_cache) as di:
        yield di


def data_iterator_to_csv(csv_path, data_path, data_iterator, seed=0):
    index = 0
    csv_data = []
    with data_iterator as data:
        line = ['x:image', 'y:label']
        csv_data.append(line)
        pbar = tqdm.tqdm(total=data.size, unit='images')
        initial_epoch = data.epoch
        while data.epoch == initial_epoch:
            d = data.next()
            for i in range(len(d[0])):
                label = d[1][i][0]
                file_name = data_path + \
                    '/{}'.format(label) + '/{}.png'.format(index)
                full_path = os.path.join(
                    csv_path, file_name.replace('/', os.path.sep))
                directory = os.path.dirname(full_path)
                if not os.path.exists(directory):
                    os.makedirs(directory)
                imwrite(full_path, d[0][i].reshape(
                    3, 32, 32).transpose(1, 2, 0))
                csv_data.append([file_name, label])
                index += 1
                pbar.update(1)
        pbar.close()
    return csv_data


def create_data_csv(seed):
    path = os.path.abspath(os.path.dirname(__file__))
    base_dir = os.path.join(path, 'cifar10')
    ensure_dir(base_dir)
    # Create original training set
    logger.log(99, 'Downloading CIFAR-10 dataset...')
    output_dir = os.path.join(path, 'download')
    train_di = data_iterator_cifar10(50000, output_dir=output_dir)
    logger.log(99, 'Creating "cifar10_training.csv"... ')
    train_csv = data_iterator_to_csv(base_dir, 'training', train_di)
    train_csv, val_csv = split_data_into_train_val(
        train_csv, val_size=10000, seed=seed)

    save_list_to_csv(train_csv, base_dir,
                     'cifar10_training' + '_' + str(seed) + '.csv')
    save_list_to_csv(val_csv, base_dir, 'cifar10_validation' +
                     '_' + str(seed) + '.csv')
    # Create original test set
    validation_di = data_iterator_cifar10(
        10000, False, None, False, output_dir=output_dir)
    logger.log(99, 'Creating "cifar10_test.csv"... ')
    test_csv = data_iterator_to_csv(base_dir, 'validation', validation_di)
    save_list_to_csv(test_csv, base_dir, 'cifar10_test.csv')

    logger.log(99, 'Dataset creation completed successfully.')


def main(args):
    create_data_csv(args.seed)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='csv data', formatter_class=argparse.RawTextHelpFormatter)

    parser.add_argument(
        '-s', '--seed', help='seed num', default=0, type=int, required=True)
    args = parser.parse_args()
    main(args)


# responsible_ai/data_cleansing/datasets/create_mnist_csv.py

# Copyright 2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''
Download MNIST examples as CSV.
'''
import os
import csv
import argparse
from imageio import imwrite
from contextlib import contextmanager
import numpy
import struct
import zlib
import tqdm

from nnabla.logger import logger
from nnabla.utils.data_iterator import data_iterator
from nnabla.utils.data_source import DataSource
from nnabla.utils.data_source_loader import download
from .utils import get_filename_to_download, save_list_to_csv, split_data_into_train_val, ensure_dir


class MnistDataSource(DataSource):
    '''
    Get data directly from MNIST dataset from Internet(yann.lecun.com).
    '''

    def _get_data(self, position):
        image = self._images[self._indexes[position]]
        label = self._labels[self._indexes[position]]
        return (image, label)

    def __init__(self, train=True, shuffle=False, rng=None, output_dir=None):
        super(MnistDataSource, self).__init__(shuffle=shuffle)
        self._train = train
        if self._train:
            image_uri = 'http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz'
            label_uri = 'http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz'
        else:
            image_uri = 'http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz'
            label_uri = 'http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz'

        logger.info('Getting label data from {}.'.format(label_uri))
        # With python3 we can write this logic as following, but with
        # python2, gzip.object does not support file-like object and
        # urllib.request does not support 'with statement'.
        #
        #   with request.urlopen(label_uri) as r, gzip.open(r) as f:
        #       _, size = struct.unpack('>II', f.read(8))
        #       self._labels = numpy.frombuffer(f.read(), numpy.uint8).reshape(-1, 1)
        #
        label_output_file = get_filename_to_download(output_dir, label_uri)
        # file object returned
        r = download(label_uri, output_file=label_output_file)
        data = zlib.decompress(r.read(), zlib.MAX_WBITS | 32)
        _, size = struct.unpack('>II', data[0:8])
        self._labels = numpy.frombuffer(data[8:], numpy.uint8).reshape(-1, 1)
        r.close()
        logger.info('Getting label data done.')

        logger.info('Getting image data from {}.'.format(image_uri))
        image_output_file = get_filename_to_download(output_dir, image_uri)
        r = download(image_uri, output_file=image_output_file)
        data = zlib.decompress(r.read(), zlib.MAX_WBITS | 32)
        _, size, height, width = struct.unpack('>IIII', data[0:16])
        self._images = numpy.frombuffer(data[16:], numpy.uint8).reshape(
            size, 1, height, width)
        r.close()
        logger.info('Getting image data done.')

        self._size = self._labels.size
        self._variables = ('x', 'y')
        if rng is None:
            rng = numpy.random.RandomState(313)
        self.rng = rng
        self.reset()

    def reset(self):
        if self._shuffle:
            self._indexes = self.rng.permutation(self._size)
        else:
            self._indexes = numpy.arange(self._size)
        super(MnistDataSource, self).reset()

    @property
    def images(self):
        """Get copy of whole data with a shape of (N, 1, H, W)."""
        return self._images.copy()

    @property
    def labels(self):
        """Get copy of whole label with a shape of (N, 1)."""
        return self._labels.copy()


@contextmanager
def data_iterator_mnist(batch_size,
                        train=True,
                        rng=None,
                        shuffle=True,
                        with_memory_cache=False,
                        with_file_cache=False,
                        output_dir=None):
    '''
    Provide DataIterator with :py:class:`MnistDataSource`
    with_memory_cache, with_parallel and with_file_cache option's default value is all False,
    because :py:class:`MnistDataSource` is able to store all data into memory.

    For example,

    .. code-block:: python

        with data_iterator_mnist(True, batch_size) as di:
            for data in di:
                SOME CODE TO USE data.

    '''
    with MnistDataSource(train=train, shuffle=shuffle, rng=rng, output_dir=output_dir) as ds, \
        data_iterator(ds,
                      batch_size,
                      rng=None,
                      with_memory_cache=with_memory_cache,
                      with_file_cache=with_file_cache) as di:
        yield di


def data_iterator_to_csv(csv_path, csv_file_name, data_path, data_iterator):
    index = 0
    csv_data = []
    with data_iterator as data:
        line = ['x:image', 'y:label']
        csv_data.append(line)
        pbar = tqdm.tqdm(total=data.size, unit='images')
        initial_epoch = data.epoch
        while data.epoch == initial_epoch:
            d = data.next()
            for i in range(len(d[0])):
                label = d[1][i][0]
                file_name = data_path + \
                    '/{}'.format(label) + '/{}.png'.format(index)
                full_path = os.path.join(
                    csv_path, file_name.replace('/', os.path.sep))
                directory = os.path.dirname(full_path)
                if not os.path.exists(directory):
                    os.makedirs(directory)
                imwrite(full_path, d[0][i].reshape(28, 28))
                csv_data.append([file_name, label])
                index += 1
                pbar.update(1)
        pbar.close()
    with open(os.path.join(csv_path, csv_file_name), 'w') as f:
        writer = csv.writer(f, lineterminator='\n')
        writer.writerows(csv_data)
    return csv_data


def create_onehot_dataset(source_csv, file_name, num_class):
    csv_data = []
    for i, line in enumerate(source_csv):
        if i == 0:
            new_line = [line[0]]
            for i2 in range(num_class):
                new_line.append('y__{}:{}'.format(i2, i2))
            csv_data.append(new_line)
        else:
            label = line[1]
            onehot = numpy.zeros(num_class)
            onehot[label] = 1
            new_line = [line[0]]
            new_line.extend(onehot)
            csv_data.append(new_line)
    with open(file_name, 'w') as f:
        writer = csv.writer(f, lineterminator='\n')
        writer.writerows(csv_data)


def create_data_csv(seed):
    path = os.path.abspath(os.path.dirname(__file__))
    base_dir = os.path.join(path, 'mnist')
    ensure_dir(base_dir)

    # Create original training set
    logger.log(99, 'Downloading MNIST training set images...')
    output_dir = os.path.join(path, 'download')
    train_di = data_iterator_mnist(
        60000, True, None, False, output_dir=output_dir)
    logger.log(99, 'Creating "mnist_training.csv"... ')
    train_csv = data_iterator_to_csv(
        base_dir, 'mnist_training.csv', 'training', train_di)

    # Create original test set
    logger.log(99, 'Downloading MNIST test set images...')
    validation_di = data_iterator_mnist(
        10000, False, None, False, output_dir=output_dir)
    logger.log(99, 'Creating "mnist_test.csv"... ')
    test_csv = data_iterator_to_csv(
        base_dir, 'mnist_test.csv', 'validation', validation_di)

    # Create one-hot training set
    logger.log(99, 'Creating "mnist_training_onehot.csv"... ')
    create_onehot_dataset(train_csv, os.path.join(
        base_dir, "mnist_training_onehot.csv"), 10)

    # Create one-hot test set
    logger.log(99, 'Creating "mnist_test_onehot.csv"... ')
    create_onehot_dataset(test_csv, os.path.join(
        base_dir, "mnist_test_onehot.csv"), 10)

    # Create 100 data training set for semi-supervised learning
    logger.log(99, 'Creating "mnist_training_100.csv"... ')
    labels = numpy.zeros(10)
    csv_data = []
    for i, line in enumerate(train_csv):
        if i == 0:
            csv_data.append(line)
        else:
            label = line[1]
            if labels[label] < 10:
                csv_data.append(line)
                labels[label] += 1
    with open(os.path.join(base_dir, "mnist_training_100.csv"), 'w') as f:
        writer = csv.writer(f, lineterminator='\n')
        writer.writerows(csv_data)

    # Create unlabeled training set for semi-supervised learning
    logger.log(99, 'Creating "mnist_training_unlabeled.csv"... ')
    labels = numpy.zeros(10)
    csv_data = []
    for i, line in enumerate(train_csv):
        if i == 0:
            csv_data.append(['xu:image'])
        else:
            csv_data.append([line[0]])
    with open(os.path.join(base_dir, "mnist_training_unlabeled.csv"), 'w') as f:
        writer = csv.writer(f, lineterminator='\n')
        writer.writerows(csv_data)

    # Create small training set
    logger.log(99, 'Creating "small_mnist_4or9_training.csv"... ')
    labels = numpy.zeros(2)
    csv_data = []
    for i, line in enumerate(train_csv):
        if i == 0:
            csv_data.append(['x:image', 'y:9'])
        else:
            label = line[1]
            if label == 4 or label == 9:
                if label == 4:
                    label = 0
                else:
                    label = 1
                if labels[label] < 750:
                    csv_data.append([line[0], label])
                    labels[label] += 1
    with open(os.path.join(base_dir, "small_mnist_4or9_training.csv"), 'w') as f:
        writer = csv.writer(f, lineterminator='\n')
        writer.writerows(csv_data)

    # Create small test set
    logger.log(99, 'Creating "small_mnist_4or9_test.csv"... ')
    labels = numpy.zeros(2)
    csv_data = []
    for i, line in enumerate(test_csv):
        if i == 0:
            csv_data.append(['x:image', 'y:9'])
        else:
            label = line[1]
            if label == 4 or label == 9:
                if label == 4:
                    label = 0
                else:
                    label = 1
                if labels[label] < 250:
                    csv_data.append([line[0], label])
                    labels[label] += 1
    with open(os.path.join(base_dir, "small_mnist_4or9_test.csv"), 'w') as f:
        writer = csv.writer(f, lineterminator='\n')
        writer.writerows(csv_data)

    # Create small test set with initial memory
    logger.log(99, 'Creating "small_mnist_4or9_test_w_initmemory.csv"... ')
    memory_size = 256
    for i in range(len(csv_data)):
        if i == 0:
            for i2 in range(memory_size):
                csv_data[0].append('c__{}'.format(i2))
        else:
            csv_data[i].extend(numpy.zeros(memory_size))
    with open(os.path.join(base_dir, "small_mnist_4or9_test_w_initmemory.csv"), 'w') as f:
        writer = csv.writer(f, lineterminator='\n')
        writer.writerows(csv_data)

    logger.log(99, 'Dataset creation completed successfully.')

    # split based on seed
    # Create original training set
    train_csv, val_csv = split_data_into_train_val(
        train_csv, val_size=10000, seed=seed)
    save_list_to_csv(train_csv, base_dir, 'mnist_training' +
                     '_' + str(seed) + '.csv')
    save_list_to_csv(val_csv, base_dir, 'mnist_validation' +
                     '_' + str(seed) + '.csv')


def main(args):
    create_data_csv(args.seed)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='csv data', formatter_class=argparse.RawTextHelpFormatter)

    parser.add_argument(
        '-s', '--seed', help='seed num', default=0, type=int, required=True)
    args = parser.parse_args()
    main(args)


# responsible_ai/data_cleansing/datasets/create_stl10_csv.py

# Copyright 2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import sys
import os
import numpy as np
from contextlib import contextmanager
from tqdm import tqdm
from imageio import imwrite
import csv
from nnabla.logger import logger
from nnabla.utils.data_iterator import data_iterator
from nnabla.utils.data_source import DataSource
from nnabla.utils.data_source_loader import download
import tarfile
import argparse
from .utils import split_data_into_train_val, save_list_to_csv, get_filename_to_download, ensure_dir

sys.path.append("../../utils/")
from neu.safe_extract import safe_extract


class STL10DataSource(DataSource):
    '''
    Directly get data from STL10
    implemented refering to this example
    https://github.com/sony/nnabla/blob/master/tutorial/cifar10_classification.ipynb
    '''

    def __init__(self, train=True, shuffle=False, rng=None, output_dir=None):
        super(STL10DataSource, self).__init__(shuffle=shuffle, rng=rng)
        self._train = train
        data_uri = 'http://ai.stanford.edu/~acoates/stl10/stl10_binary.tar.gz'
        logger.info('Getting labbeled data from {}'.format(data_uri))
        default_seed = 313
        output_file = get_filename_to_download(output_dir, data_uri)
        r = download(data_uri, output_file=output_file)  # file object returned
        print(r.name)
        binary_dir = os.path.join(output_dir, "stl10_binary")
        with tarfile.open(fileobj=r, mode="r:gz") as tar:
            safe_extract(tar, path=output_dir)

        for member in os.listdir(binary_dir):
            if train:
                if 'train_' not in member:
                    continue
                print(member)
                self.load_image_and_labels(os.path.join(binary_dir, member))

            # Validation data
            else:
                print(member)
                if 'test_' not in member:
                    continue
                self.load_image_and_labels(os.path.join(binary_dir, member))
        r.close
        logger.info('Getting labeled data from {}'.format(data_uri))
        self._size = self._labels.size
        self._variables = ('x', 'y')
        if rng is None:
            rng = np.random.RandomState(default_seed)
        self.rng = rng
        self.reset()

    def load_image_and_labels(self, path_to_file):
        if '_X' in path_to_file:
            print("images")
            self._images = self.load_images(path_to_file)

        elif '_y' in path_to_file:
            print("labels")
            self._labels = self.load_labels(path_to_file)

    def load_images(self, image_binary):
        with open(image_binary, "rb") as fp:
            images = np.fromfile(fp, dtype=np.uint8)
            images = images.reshape(-1, 3, 96, 96)
            return images

    def load_labels(self, label_binary):
        with open(label_binary, "rb") as fp:
            labels = np.fromfile(fp, dtype=np.uint8)
            labels = labels.reshape(-1, 1)
            # 1-10 -> 0-9
            labels = labels - 1
            return labels

    def _get_data(self, position):
        image = self._images[self._indices[position]]
        label = self._labels[self._indices[position]]
        return (image, label)

    def reset(self):
        if self._shuffle:
            self._indices = self.rng.permutation(self._size)
        else:
            self._indices = np.arange(self._size)
        super(STL10DataSource, self).reset()

    @property
    def images(self):
        '''Get copy of whole data with a shape of (N, 1, H, W).'''
        return self._images.copy()

    @property
    def labels(self):
        '''Get copy of whole label with a shape of (N, 1).'''
        return self._labels.copy()


@contextmanager
def data_iterator_stl10(batch_size,
                        train=True,
                        rng=None,
                        shuffle=True,
                        with_memory_cache=False,
                        with_file_cache=False,
                        output_dir=None):
    '''
    Provide DataIterator with :py:class:`STL10DataSource`
    with_memory_cache and with_file_cache option's default value is all False,
    because :py:class:`STL10DataSource` is able to store all data into memory.
    '''
    """
    _data_iterator = data_iterator(
        STL10DataSource(train=train, shuffle=shuffle, rng=rng),
        batch_size,
        rng,
        with_memory_cache,
        with_file_cache
    )
    return _data_iterator
    """

    with STL10DataSource(train=train, shuffle=shuffle, rng=rng, output_dir=output_dir) as ds, \
        data_iterator(ds,
                      batch_size,
                      rng=rng,
                      with_memory_cache=with_memory_cache,
                      with_file_cache=with_file_cache) as di:
        yield di


def data_iterator_to_csv(csv_path, csv_file_name, data_path, data_iterator):
    index = 0
    csv_data = []
    with data_iterator as data:
        line = ['x:image', 'y:label']
        csv_data.append(line)
        pbar = tqdm(total=data.size, unit='images')
        initial_epoch = data.epoch
        while data.epoch == initial_epoch:
            d = data.next()
            for i in range(len(d[0])):
                label = d[1][i][0]
                file_name = data_path + \
                    '/{}'.format(label) + '/{}.png'.format(index)
                full_path = os.path.join(
                    csv_path, file_name.replace('/', os.path.sep))
                directory = os.path.dirname(full_path)
                if not os.path.exists(directory):
                    os.makedirs(directory)
                imwrite(full_path, d[0][i].reshape(
                    3, 96, 96).transpose(1, 2, 0))
                csv_data.append([file_name, label])
                index += 1
                pbar.update(1)
        pbar.close()
    with open(os.path.join(csv_path, csv_file_name), 'w') as f:
        writer = csv.writer(f, lineterminator='\n')
        writer.writerows(csv_data)
    return csv_data


def create_data_csv(seed):
    path = os.path.abspath(os.path.dirname(__file__))
    base_dir = os.path.join(path, 'stl10')
    ensure_dir(base_dir)
    # Create original training set
    logger.log(99, 'Downloading STL10 dataset...')
    output_dir = os.path.join(path, 'download')
    train_di = data_iterator_stl10(
        5000, True, None, False, output_dir=output_dir)
    logger.log(99, 'Creating "stl10_training.csv"... ')
    train_csv = data_iterator_to_csv(
        base_dir, 'stl10_training.csv', 'training', train_di)
    train_csv, val_csv = split_data_into_train_val(
        train_csv, val_size=1000, seed=seed)
    save_list_to_csv(train_csv, base_dir, 'stl10_training' +
                     '_' + str(seed) + '.csv')
    save_list_to_csv(val_csv, base_dir, 'stl10_validation' +
                     '_' + str(seed) + '.csv')

    # Validation
    validation_di = data_iterator_stl10(
        8000, False, None, False, output_dir=output_dir)
    logger.log(99, 'Creating "stl10_test.csv"... ')
    _ = data_iterator_to_csv(
        base_dir, 'stl10_test.csv', 'validation', validation_di)
    logger.log(99, 'Dataset creation completed successfully.')


def main(args):
    create_data_csv(args.seed)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='csv data', formatter_class=argparse.RawTextHelpFormatter)

    parser.add_argument(
        '-s', '--seed', help='seed num', default=0, type=int, required=True)
    args = parser.parse_args()
    main(args)


# responsible_ai/data_cleansing/datasets/utils.py

# Copyright 2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import csv
import numpy as np


def get_filename_to_download(output_dir: str, data_uri: str):
    if output_dir is None:
        output_file = None
    else:
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        output_file = os.path.join(output_dir, data_uri.split('/')[-1])
    return output_file


def save_list_to_csv(csv_data, csv_path, csv_file_name):
    with open(os.path.join(csv_path, csv_file_name), 'w') as f:
        writer = csv.writer(f, lineterminator='\n')
        writer.writerows(csv_data)


def split_data_into_train_val(data_list, val_size=10000, seed=0):
    n = len(data_list) - 1  # forst row is title
    np.random.seed(seed)
    idx_val = np.random.permutation(n)[:val_size]
    idx_train = np.setdiff1d(np.arange(n), idx_val)
    data = data_list[1:]
    train_data = [data_list[0]] + [data[i] for i in idx_train]
    val_data = [data_list[0]] + [data[i] for i in idx_val]
    return train_data, val_data


def ensure_dir(path):
    if not os.path.exists(path):
        os.makedirs(path)


# responsible_ai/data_cleansing/influence_functions/infl.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import nnabla as nn
import nnabla.solvers as S
import nnabla.functions as F
import os
import functools
from tqdm import tqdm
from sgd_influence.model import setup_model
from sgd_influence.dataset import get_batch_data, init_dataset, get_data, get_image_size, get_batch_indices
from sgd_influence.utils import ensure_dir, get_indices, save_to_csv


def adjust_batch_size(model, batch_size, loss_fn=None):
    has_loss = loss_fn is not None
    if has_loss:
        loss_d, loss_g = loss_fn.d, loss_fn.g
    pred, loss_fn, input_image = model(batch_size=batch_size)
    if has_loss:
        loss_fn.d = loss_d
        loss_fn.g = loss_g
    return pred, loss_fn, input_image


def save_infl_for_analysis(infl_list, use_all_params, save_dir, infl_filename, epoch, header, data_type):
    dn = os.path.join(save_dir, 'epoch%02d' % (epoch))
    if use_all_params:
        dn = os.path.join(dn, 'infl_original')
    else:
        dn = os.path.join(dn, 'infl_arranged')
    ensure_dir(dn)
    save_to_csv(filename=os.path.join(dn, os.path.basename(infl_filename)),
                header=header, list_to_save=infl_list, data_type=data_type)


def compute_gradient(grad_model, solver, dataset, batch_size, idx_list_to_data, resize_size):
    n = len(idx_list_to_data)
    grad_idx = get_batch_indices(n, batch_size, seed=None)
    u = {}
    loss_fn = None
    for i in tqdm(grad_idx, desc='calc gradient (2/3 steps)'):
        X, y = get_batch_data(dataset, idx_list_to_data,
                              i, resize_size, test=True)
        _, loss_fn, input_image = adjust_batch_size(
            grad_model, len(X), loss_fn)
        input_image["image"].d = X
        input_image["label"].d = y
        loss_fn.forward()
        solver.zero_grad()
        loss_fn.backward(clear_buffer=True)

        for key, param in nn.get_parameters(grad_only=False).items():
            uu = u.get(key, None)
            if uu is None:
                u[key] = nn.Variable(param.shape)
                u[key].data.zero()
            u[key].d += param.g / n
    return u


def infl_icml(model_info_dict, file_dir_dict, use_all_params, need_evaluate, alpha):
    num_epochs = 2
    # params
    lr = 0.005
    seed = model_info_dict['seed']
    net_func = model_info_dict['net_func']
    batch_size = model_info_dict['batch_size']
    test_batch_size = 1000
    target_epoch = model_info_dict['num_epochs']
    # files and dirs
    save_dir = file_dir_dict['save_dir']
    infl_filename = file_dir_dict['infl_filename']
    final_model_name = file_dir_dict['model_filename']
    final_model_path = os.path.join(save_dir, 'epoch%02d' % (
        target_epoch - 1), 'weights', final_model_name)
    input_dir_name = os.path.dirname(file_dir_dict['train_csv'])

    # setup
    trainset, valset, image_shape, n_classes, ntr, nval = init_dataset(
        file_dir_dict['train_csv'], file_dir_dict['val_csv'], seed)
    n_channels, _h, _w = image_shape
    resize_size = get_image_size((_h, _w))
    idx_train = get_indices(ntr, seed)
    idx_val = get_indices(nval, seed)

    nn.load_parameters(final_model_path)
    trained_params = nn.get_parameters(grad_only=False)

    test = True

    grad_model = functools.partial(
        setup_model, net_func=net_func, n_classes=n_classes, n_channels=n_channels, resize_size=resize_size, test=test, reduction='mean')
    solver = S.Momentum(lr=lr, momentum=0.9)
    solver.set_parameters(trained_params)
    # gradient
    u = compute_gradient(grad_model, solver, valset,
                         test_batch_size, idx_val, resize_size)

    # Hinv * u with SGD
    seed_train = 0
    v = dict()
    for key, param in nn.get_parameters(grad_only=False).items():
        v[key] = nn.Variable(param.d.shape, need_grad=True)
        v[key].d = 0
        v[key].g = 0

    solver.set_parameters(v)

    loss_train = []
    loss_fn = None
    for epoch in range(num_epochs):
        # training
        seed_train = 0
        np.random.seed(epoch)
        idx = get_batch_indices(ntr, batch_size, seed=epoch)
        for j, i in enumerate(idx):
            seeds = list(range(seed_train, seed_train + i.size))
            seed_train += i.size
            X, y = get_batch_data(trainset, idx_train, i,
                                  resize_size, test=False, seeds=seeds)
            _, loss_fn, input_image = adjust_batch_size(
                grad_model, len(X), loss_fn)
            input_image["image"].d = X
            input_image["label"].d = y
            loss_fn.forward()

            grad_params = nn.grad(
                loss_fn, [param for param in nn.get_parameters(grad_only=False).values()])
            vg = 0
            for vv, g in zip(v.values(), grad_params):
                vg += F.sum(vv*g)

            for parameters in trained_params.values():
                parameters.grad.zero()

            vgrad_params = nn.grad(
                vg, [param for param in nn.get_parameters(grad_only=False).values()])
            loss_i = 0
            for vgp, vv, uu in zip(vgrad_params, v.values(), u.values()):
                loss_i += 0.5 * F.sum(vgp * vv + alpha *
                                      vv * vv) - F.sum(uu * vv)
            loss_i.forward()

            solver.zero_grad()
            loss_i.backward(clear_buffer=True)
            solver.update()
            loss_train.append(loss_i.d.copy())

    # influence
    infl_dict = dict()
    infl = np.zeros(ntr)
    for i in tqdm(range(ntr), desc='calc influence (3/3 steps)'):
        csv_idx = idx_train[i]
        file_name = trainset.get_filepath_to_data(csv_idx)
        file_name = os.path.join(input_dir_name, file_name)
        file_name = os.path.normpath(file_name)
        X, y = get_data(trainset, idx_train[i], resize_size, True, seed=i)
        _, loss_fn, input_image = adjust_batch_size(
            grad_model, len(X), loss_fn)
        input_image["image"].d = X
        input_image["label"].d = y
        loss_fn.forward()
        for parameters in trained_params.values():
            parameters.grad.zero()
        loss_fn.backward(clear_buffer=True)
        infl_i = 0
        for j, param in enumerate(nn.get_parameters(grad_only=False).values()):
            infl_i += (param.g.copy() * list(v.values())[j].d.copy()).sum()
        infl[i] = -infl_i / ntr
        infl_dict[csv_idx] = [file_name, y, infl[i]]
    infl_list = [val + [key] for key, val in infl_dict.items()]
    infl_list = sorted(infl_list, key=lambda x: (x[-2]))

    # save
    header = ['x:image', 'y:label', 'influence', 'datasource_index']
    data_type = 'object,int,float,int'
    if need_evaluate:
        save_infl_for_analysis(infl_list, use_all_params,
                               save_dir, infl_filename, epoch, header, data_type)
    save_to_csv(filename=infl_filename, header=header,
                list_to_save=infl_list, data_type=data_type)


# responsible_ai/data_cleansing/sgd_influence/analyze.py

# Copyright 2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import pandas as pd
import numpy as np
from matplotlib import pyplot as plt
from .utils import ensure_dir


def get_summary_df(file_dir, method_list, baseline_name):
    summary_df = pd.DataFrame()
    for method in method_list:
        file_method_dir = os.path.join(file_dir, method)
        score_files = os.listdir(file_method_dir)
        for score_file in score_files:
            df = pd.read_csv(os.path.join(file_method_dir, score_file))
            file_basename = score_file.split('_')
            seed = file_basename[-1].split('.')[0]
            n_drops = file_basename[-3]
            _df = df.copy()
            _df['method'] = method
            _df['seed'] = int(seed)
            _df['n_drops'] = int(n_drops)
            summary_df = pd.concat([summary_df, _df], axis=0)

    summary_df.loc[summary_df['n_drops'] == 0, 'method'] = baseline_name
    _v = summary_df.loc[summary_df['n_drops'] == 0].copy()
    for i in set(summary_df['n_drops']) - set([0]):
        _v['n_drops'] = i
        summary_df = pd.concat([summary_df, _v], axis=0)
    # calc mean
    mean_df = summary_df.groupby(['method', 'n_drops']).mean().reset_index()
    mean_df['mean'] = 'mean'
    mean_df['seed'] = np.nan
    # calc std
    std_df = summary_df.groupby(['method', 'n_drops']).std().reset_index()
    std_df['std'] = 'std'
    std_df['seed'] = np.nan
    summary_df = pd.concat([summary_df, mean_df], axis=0)
    summary_df = pd.concat([summary_df, std_df], axis=0)
    return summary_df.reset_index(drop=True)


def analyze_result(score_dir, baseline_name='No Removal'):
    # obtain method
    files = os.listdir(score_dir)
    methods = np.array(
        [f for f in files if os.path.isdir(os.path.join(score_dir, f))])
    # read data
    summary_df = get_summary_df(score_dir, methods, baseline_name)
    return summary_df


def plot_result(acc, output_dir, baseline_name, fig_save_dir_name='fig'):
    # drawing
    df = acc.loc[acc['mean'].notna()]
    hue = 'method'
    # sort labels
    labels = sorted(df[hue].drop_duplicates().values.tolist())
    labels.remove(baseline_name)
    labels.append(baseline_name)

    dataset = [(df[df[hue] == x]['n_drops'].values.tolist(
    ), df[df[hue] == x]['test_accuracy'].values.tolist()) for x in labels]
    for label, dataset in zip(labels, dataset):
        x, y = dataset[0], dataset[1]
        if label == baseline_name:
            plt.semilogx(x, y)
        else:
            plt.semilogx(x, y, marker='o')
    plt.legend(labels)

    plt.xlabel('# of instances removed')
    plt.ylabel('Accuracy')
    # save
    fig_save_dir = os.path.join(output_dir, fig_save_dir_name)
    ensure_dir(fig_save_dir)
    plt.savefig(os.path.join(fig_save_dir, 'result.png'))


def analyze(file_dir, baseline_name='No Removal'):
    # analyze and save result
    summary_df = analyze_result(os.path.join(file_dir, 'score'), baseline_name)
    summary_dir = os.path.normpath(os.path.join(file_dir, 'summary'))
    ensure_dir(summary_dir)
    summary_df.to_csv(os.path.join(summary_dir, 'summary.csv'))
    # plot result
    plot_result(summary_df, file_dir, baseline_name)


# responsible_ai/data_cleansing/sgd_influence/dataset.py

# Copyright 2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import numpy as np
import nnabla as nn
from nnabla.utils.data_source_loader import _load_functions
from nnabla.utils.data_source_implements import CsvDataSource
from nnabla.utils.image_utils import imresize
import nnabla.functions as F
import random


class MeanStd(object):
    def __new__(cls, *args, **kargs):
        if not hasattr(cls, "_instance"):
            cls._instance = super(MeanStd, cls).__new__(cls)
        return cls._instance

    def __init__(self, resize_size, n_channels):
        if n_channels == 1:
            self._init_single_channel(resize_size)
        else:
            self._init_multi_channel(resize_size)

    def _init_single_channel(self, resize_size):
        ch_mean = np.array([np.full(resize_size, 0.5)])
        ch_std = np.array([np.full(resize_size, 0.5)])
        self.ch_mean = nn.Variable.from_numpy_array(ch_mean)
        self.ch_std = nn.Variable.from_numpy_array(ch_std)

    def _init_multi_channel(self, resize_size):
        ch_mean = np.array([np.full(resize_size, 0.4914), np.full(
            resize_size, 0.4822), np.full(resize_size, 0.4465)])
        ch_std = np.array([np.full(resize_size, 0.2023), np.full(
            resize_size, 0.1994), np.full(resize_size, 0.2010)])
        self.ch_mean = nn.Variable.from_numpy_array(ch_mean)
        self.ch_std = nn.Variable.from_numpy_array(ch_std)


class CsvDataSourceFilename(CsvDataSource):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._filepath_column_idx = self._get_column_idx('filepath')
        self._label_column_idx = self._get_column_idx('label')

    def _get_column_idx(self, target):
        """
        Parameters
        ----------
        target: str
            'filepath' or 'label'

        Returns
        ----------
        i: int
            index of column that contains 'filepath' or 'label' info
        """
        if target not in ['filepath', 'label']:
            raise KeyError("target is 'filepath' or 'label'")

        def _is_filepath_idx(value):
            ext = (os.path.splitext(value)[1]).lower()
            if ext in _load_functions.keys():
                return True
            return False

        def _is_label_idx(value):
            try:
                value = float(value)
                return True
            except ValueError:
                return False

        func_dict = {
            'filepath': _is_filepath_idx,
            'label': _is_label_idx
        }

        # judge from first row
        idx = 0
        for i, column_value in enumerate(self._rows[idx]):
            # Implemented refering to below
            # https://github.com/sony/nnabla/blob/f5eff2de5329ef02c40e7a5d7344abd91b19ece8/python/src/nnabla/utils/data_source_implements.py#L402
            # https://github.com/sony/nnabla/blob/f5eff2de5329ef02c40e7a5d7344abd91b19ece8/python/src/nnabla/utils/data_source_loader.py#L343
            if func_dict[target](column_value):
                return i
        raise RuntimeError(
            '{} info is not in {}.'.format(target, self._filename))

    def get_filepath_to_data(self, idx):
        """
        Parameters
        ----------
        idx: int
            index of self._rows that includes filepath to data

        Returns
        ----------
        filepath: str
            ex: 'training/1.png'
        """
        return self._rows[idx][self._filepath_column_idx]

    def get_n_classes(self):
        """
        Parameters
        ----------
        None

        Returns
        ----------
        n_classes: int
            number of classes of label
            ex. n_classes is 10 if the dataset contains 0-9 labels 
        """

        labels = [row[self._label_column_idx] for row in self._rows]
        n_classes = len(set(labels))
        return n_classes


def setup_nnabla_dataset(filename):
    dataset = CsvDataSourceFilename(
        filename=filename,
        shuffle=False,
        normalize=True
    )
    return dataset


def get_batch_data(dataset, idx_list_to_data, idx_list_to_idx, resize_size, test, seeds=None, escape_list=[]):
    X = []
    y = []
    if seeds is None:
        seeds = [0] * len(idx_list_to_idx)
    for idx, seed in zip(idx_list_to_idx, seeds):
        if idx_list_to_data[idx] in escape_list:
            continue
        image, label = get_data(
            dataset, idx_list_to_data[idx], resize_size, test, seed)
        X.append(image)
        y.append(label)
    y = np.array(y).reshape(-1, 1)
    return X, y


def normalize(x, resize_size):
    mean_std = MeanStd(resize_size, x.shape[0])
    ch_mean = mean_std.ch_mean
    ch_std = mean_std.ch_std
    x = F.sub2(x, ch_mean)
    x = F.div2(x, ch_std)
    return x


def transform(x, resize_size, seed, test):
    x = nn.Variable.from_numpy_array(x)
    if not test:
        # random crop
        random.seed(seed)
        x = F.image_augmentation(x, shape=resize_size, pad=(4, 4), seed=seed)
        # random horizontal flip
        random.seed(seed)
        x = F.image_augmentation(x, flip_lr=True, seed=seed)
    x = normalize(x, resize_size)
    x.forward(clear_buffer=True)
    x = x.d
    return x


def init_dataset(input_train, input_val, seed):
    return init_nnabla(input_train, input_val)


def init_nnabla(input_train, input_val):
    trainset = setup_nnabla_dataset(input_train)
    valset = setup_nnabla_dataset(input_val)
    image, _ = trainset._get_data(0)
    n_classes = trainset.get_n_classes()
    ntr = trainset.size
    nval = valset.size
    return trainset, valset, image.shape, n_classes, ntr, nval


def get_data(dataset, idx, resize_size, test, seed):
    image, label = get_data_nnabla(dataset, idx, resize_size, test, seed)
    return image, label


def get_data_nnabla(dataset, idx, resize_size, test, seed):
    image, label = dataset._get_data(idx)
    image = imresize(image, resize_size, channel_first=True)
    image = transform(image, resize_size, seed, test)
    return image, label[0]


def get_batch_indices(num_data, batch_size, seed=None):
    if seed is None:
        shuffled_idx = np.arange(num_data)
    else:
        np.random.seed(seed)
        shuffled_idx = np.random.permutation(num_data)
    indices_list = []
    for i in range(0, num_data, batch_size):
        indices = shuffled_idx[i:i+batch_size]
        indices_list.append(indices)
    return indices_list


def get_image_size(image_shape, max_size=128, min_size=32):
    """
    Parameters
    ----------
    image_shape: tuple
        (int, int)
    max_size: int
    min_size: int

    Returns
    ----------
    resize_size: tuple
        (int, int)
    """
    h, w = image_shape
    short_edge = min(h, w)
    if short_edge < min_size:
        resize_size = (min_size, min_size)
    elif (short_edge >= min_size) & (short_edge <= max_size):
        resize_size = (short_edge, short_edge)
    elif short_edge > max_size:
        resize_size = (max_size, max_size)
    return resize_size


# responsible_ai/data_cleansing/sgd_influence/eval.py

# Copyright 2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import functools
from tqdm import tqdm
import numpy as np
import nnabla as nn
import nnabla.solvers as S
from .model import setup_model, adjust_batch_size
from .utils import get_context, get_indices, save_to_csv, read_csv
from .dataset import get_batch_data, init_dataset, get_image_size, setup_nnabla_dataset
from .train import eval_model


def select_model_file(retrain_all, output_path, start_epoch):
    if retrain_all:
        fn = '%s/epoch%02d/weights/model_step%04d.h5' % (output_path, 0, 0)
    else:
        fn = '%s/epoch%02d/weights/final_model.h5' % (
            output_path, start_epoch - 1)
    return fn


def retrain(model_info_dict, file_dir_dict, retrain_all, escape_list=[]):
    # params
    lr = model_info_dict['lr']
    seed = model_info_dict['seed']
    net_func = model_info_dict['net_func']
    batch_size = model_info_dict['batch_size']
    end_epoch = model_info_dict['num_epochs']
    start_epoch = model_info_dict['start_epoch']
    # files and dirs
    save_dir = file_dir_dict['save_dir']
    info_filename = file_dir_dict['info_filename']
    score_filename = file_dir_dict['score_filename']
    info_file_path = os.path.join(save_dir, info_filename)
    # setup
    trainset, valset, image_shape, n_classes, ntr, nval = init_dataset(
        file_dir_dict['train_csv'], file_dir_dict['val_csv'], seed)
    testset = setup_nnabla_dataset(file_dir_dict['test_csv'])
    ntest = testset.size
    n_channels, _h, _w = image_shape
    resize_size = get_image_size((_h, _w))

    # Create training graphs
    test = False
    train_model = functools.partial(
        setup_model, net_func=net_func, n_classes=n_classes, n_channels=n_channels, resize_size=resize_size, test=test)
    # Create validation graphs
    test = True
    val_model = functools.partial(
        setup_model, net_func=net_func, n_classes=n_classes, n_channels=n_channels, resize_size=resize_size, test=test)

    # setup optimizer (SGD)
    fn = select_model_file(retrain_all, save_dir, start_epoch)
    nn.load_parameters(fn)
    solver = S.Sgd(lr=lr)
    solver.set_parameters(nn.get_parameters(grad_only=False))
    # get shuffled index using designated seed
    idx_train = get_indices(ntr, seed)
    idx_val = get_indices(nval, seed)
    idx_test = get_indices(ntest, seed)

    # training
    info = np.load(info_file_path, allow_pickle=True)
    score = []
    loss_train = None
    for epoch in tqdm(range(start_epoch, end_epoch), desc='retrain'):
        for step_info in info[epoch]:
            idx, seeds, lr = step_info['idx'], step_info['seeds'], step_info['lr']
            X, y = get_batch_data(trainset, idx_train, idx, resize_size,
                                  test=False, seeds=seeds, escape_list=escape_list)
            if len(X) == 0:
                continue
            _, loss_train, input_image_train = adjust_batch_size(
                train_model, solver, len(X), loss_train)
            input_image_train["image"].d = X
            input_image_train["label"].d = y

            loss_train.forward()
            solver.zero_grad()
            loss_train.backward(clear_buffer=True)
            for key, param in nn.get_parameters(grad_only=False).items():
                param.g *= len(X) / idx.size
            solver.update()
        # evaluation
        loss_val, acc_val = eval_model(
            val_model, solver, valset, idx_val, batch_size, resize_size)
        loss_test, acc_test = eval_model(
            val_model, solver, testset, idx_test, batch_size, resize_size)
        score.append((loss_val, loss_test, acc_val, acc_test))
        # save
    save_to_csv(filename=score_filename, header=[
                'val_loss', 'test_loss', 'val_accuracy', 'test_accuracy', ], list_to_save=score, data_type='float,float,float,float')


def get_escape_list(infl_filename, n_to_remove):
    infl_list = read_csv(infl_filename)
    infl_arr = np.array(infl_list)
    header = infl_list[0]
    idx_ds = header.index('datasource_index')
    escape_list = infl_arr[1:1+n_to_remove, idx_ds].astype(int).tolist()
    return escape_list


def run_train_for_eval(model_info_dict, file_dir_dict, n_to_remove, retrain_all):
    os.environ['NNABLA_CUDNN_DETERMINISTIC'] = '1'
    os.environ['NNABLA_CUDNN_ALGORITHM_BY_HEURISTIC'] = '1'

    # files and dirs
    infl_filename = file_dir_dict['infl_filename']
    # gpu/cpu
    ctx = get_context(model_info_dict['device_id'])
    nn.set_default_context(ctx)

    # train
    escape_list = get_escape_list(infl_filename, n_to_remove)
    retrain(model_info_dict, file_dir_dict,
            retrain_all, escape_list=escape_list)


# responsible_ai/data_cleansing/sgd_influence/infl.py

# Copyright 2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import functools
import numpy as np
from tqdm import tqdm
import nnabla as nn
import nnabla.solvers as S
import nnabla.functions as F
from .model import setup_model, adjust_batch_size
from .dataset import get_batch_data, init_dataset, get_data, get_image_size, get_batch_indices
from .utils import ensure_dir, get_indices, save_to_csv


def select_modelfile_for_infl(use_all_params, final_model_path, save_dir, epoch, step):
    if use_all_params:
        weights_dir = 'weights'
        fn = '%s/epoch%02d/%s/model_step%04d.h5' % (
            save_dir, epoch, weights_dir, step)
    else:
        fn = final_model_path
    return fn


def save_infl_for_analysis(infl_list, use_all_params, save_dir, infl_filename, epoch, header, data_type):
    dn = os.path.join(save_dir, 'epoch%02d' % (epoch))
    if use_all_params:
        dn = os.path.join(dn, 'infl_original')
    else:
        dn = os.path.join(dn, 'infl_arranged')
    ensure_dir(dn)
    save_to_csv(filename=os.path.join(dn, os.path.basename(infl_filename)),
                header=header, list_to_save=infl_list, data_type=data_type)


def infl_sgd(model_info_dict, file_dir_dict, use_all_params, need_evaluate):
    # params
    lr = model_info_dict['lr']
    seed = model_info_dict['seed']
    net_func = model_info_dict['net_func']
    batch_size = model_info_dict['batch_size']
    end_epoch = model_info_dict['end_epoch']
    target_epoch = model_info_dict['num_epochs']
    # files and dirs
    save_dir = file_dir_dict['save_dir']
    info_filename = file_dir_dict['info_filename']
    infl_filename = file_dir_dict['infl_filename']
    final_model_name = file_dir_dict['model_filename']
    final_model_path = os.path.join(save_dir, 'epoch%02d' % (
        target_epoch - 1), 'weights', final_model_name)
    input_dir_name = os.path.dirname(file_dir_dict['train_csv'])

    # setup
    trainset, valset, image_shape, n_classes, ntr, nval = init_dataset(
        file_dir_dict['train_csv'], file_dir_dict['val_csv'], seed)
    n_channels, _h, _w = image_shape
    resize_size = get_image_size((_h, _w))
    idx_train = get_indices(ntr, seed)
    idx_val = get_indices(nval, seed)

    nn.load_parameters(final_model_path)
    trained_params = nn.get_parameters(grad_only=False)
    test = True

    grad_model = functools.partial(
        setup_model, net_func=net_func, n_classes=n_classes, n_channels=n_channels, resize_size=resize_size, test=test, reduction='sum')

    solver = S.Sgd(lr=lr)
    solver.set_parameters(trained_params)
    # gradient
    u = compute_gradient(grad_model, solver, valset,
                         batch_size, idx_val, target_epoch, resize_size)

    test = False
    infl_model = functools.partial(
        setup_model, net_func=net_func, n_classes=n_classes, n_channels=n_channels, resize_size=resize_size, test=test)
    # influence
    infl_dict = {}
    info = np.load(os.path.join(save_dir, info_filename), allow_pickle=True)
    loss_fn = None
    for epoch in tqdm(range(target_epoch - 1, end_epoch - 1, -1), desc='calc influence (3/3 steps)'):
        for step_info in info[epoch][::-1]:
            idx, seeds, lr, step = step_info['idx'], step_info['seeds'], step_info['lr'], step_info['step']
            fn = select_modelfile_for_infl(
                use_all_params, final_model_path, save_dir, epoch, step)
            _, loss_fn, input_image = adjust_batch_size(
                infl_model, solver, 1, loss_fn)
            nn.load_parameters(fn)
            params = nn.get_parameters(grad_only=False)
            solver = S.Sgd(lr=lr)
            solver.set_parameters(params)
            X = []
            y = []
            for i, seed in zip(idx, seeds):
                i = int(i)
                image, label = get_data(
                    trainset, idx_train[i], resize_size, test, seed=seed)
                X.append(image)
                y.append(label)
                input_image["image"].d = image
                input_image["label"].d = label
                loss_fn.forward()
                solver.zero_grad()
                loss_fn.backward(clear_buffer=True)

                csv_idx = idx_train[i]
                infl = infl_dict.get(csv_idx, [0.0])[-1]
                for j, (key, param) in enumerate(nn.get_parameters(grad_only=False).items()):
                    infl += lr * (u[key].d * param.g).sum() / idx.size

                # store infl
                file_name = trainset.get_filepath_to_data(csv_idx)
                file_name = os.path.join(input_dir_name, file_name)
                file_name = os.path.normpath(file_name)
                infl_dict[csv_idx] = [file_name, label, infl]

            # update u
            _, loss_fn, input_image = adjust_batch_size(
                infl_model, solver, len(idx), loss_fn)
            input_image["image"].d = X
            input_image["label"].d = np.array(y).reshape(-1, 1)
            loss_fn.forward()
            params = nn.get_parameters(grad_only=False)
            grad_params = {}
            for key, p in zip(params.keys(), nn.grad([loss_fn], params.values())):
                grad_params[key] = p
            ug = 0
            # compute H[t]u[t]
            for key, uu in u.items():
                try:
                    ug += F.sum(uu * grad_params[key])
                except TypeError:
                    # cannot calc grad with batch normalization runnning mean and var
                    pass
            ug.forward()
            solver.zero_grad()
            ug.backward(clear_buffer=True)

            for j, (key, param) in enumerate(nn.get_parameters(grad_only=False).items()):
                u[key].d -= lr * param.g / idx.size

        # sort by influence score
        infl_list = [val + [key] for key, val in infl_dict.items()]
        infl_list = sorted(infl_list, key=lambda x: (x[-2]))

        # save
        header = ['x:image', 'y:label', 'influence', 'datasource_index']
        data_type = 'object,int,float,int'
        if need_evaluate:
            save_infl_for_analysis(
                infl_list, use_all_params, save_dir, infl_filename, epoch, header, data_type)
    save_to_csv(filename=infl_filename, header=header,
                list_to_save=infl_list, data_type=data_type)


def compute_gradient(grad_model, solver, dataset, batch_size, idx_list_to_data, epoch, resize_size):
    n = len(idx_list_to_data)
    grad_idx = get_batch_indices(n, batch_size, seed=None)
    u = {}
    loss_fn = None
    for i in tqdm(grad_idx, desc='calc gradient (2/3 steps)'):
        X, y = get_batch_data(dataset, idx_list_to_data,
                              i, resize_size, test=True)
        _, loss_fn, input_image = adjust_batch_size(
            grad_model, solver, len(X), loss_fn)
        input_image["image"].d = X
        input_image["label"].d = y
        loss_fn.forward()
        solver.zero_grad()
        loss_fn.backward(clear_buffer=True)

        for j, (key, param) in enumerate(nn.get_parameters(grad_only=False).items()):
            uu = u.get(key, None)
            if uu is None:
                u[key] = nn.Variable(param.shape)
                u[key].data.zero()
            u[key].d += param.g / n
    return u


# responsible_ai/data_cleansing/sgd_influence/model.py

# Copyright 2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import functools
import numpy as np
import nnabla as nn
import nnabla.parametric_functions as PF
import nnabla.functions as F


def block(x, scope_name, n_channels, kernel, pad, test):
    with nn.parameter_scope(scope_name):
        with nn.parameter_scope('conv1'):
            h = PF.convolution(x, n_channels, kernel=kernel,
                               pad=pad, with_bias=True)
            h = PF.batch_normalization(h, batch_stat=not test)
            h = F.relu(h)

        with nn.parameter_scope('conv2'):
            h = PF.convolution(h, n_channels, kernel=kernel,
                               pad=pad, with_bias=True)
            h = F.relu(h)
            h = F.max_pooling(h, kernel=(2, 2), stride=(2, 2))
    return h


def cifarnet(x, test=False, n_classes=10):
    maps = [32, 64, 128]
    kernel = (3, 3)
    pad = (1, 1)

    h = block(x, 'block1', maps[0], kernel, pad, test)
    h = block(h, 'block2', maps[1], kernel, pad, test)
    h = block(h, 'block3', maps[2], kernel, pad, test)
    h = PF.affine(h, n_classes)
    return h


def loss_function(pred, label, reduction='mean'):
    loss_dict = {
        'mean': F.mean(F.softmax_cross_entropy(pred, label)),
        'sum': F.sum(F.softmax_cross_entropy(pred, label))
    }
    return loss_dict[reduction]


def calc_acc(pred, label, method='mean'):
    acc_sum = (np.argmax(pred, axis=1).reshape(-1, 1)
               == np.array(label).reshape(-1, 1)).sum()
    acc_dict = {
        'mean': acc_sum / len(label),
        'sum': acc_sum
    }
    return acc_dict[method]


def setup_model(net_func, n_classes, n_channels, resize_size, batch_size, test, reduction='mean'):
    prediction = functools.partial(net_func, n_classes=n_classes)
    image = nn.Variable(
        (batch_size, n_channels, resize_size[0], resize_size[1]))
    label = nn.Variable((batch_size, 1))
    pred = prediction(image, test)
    loss_fn = loss_function(pred, label, reduction)
    input_image = {"image": image, "label": label}
    return pred, loss_fn, input_image


def adjust_batch_size(model, solver, batch_size, loss_fn=None):
    params = nn.get_parameters(grad_only=False)
    has_loss = loss_fn is not None
    if has_loss:
        loss_d, loss_g = loss_fn.d, loss_fn.g
    pred, loss_fn, input_image = model(batch_size=batch_size)
    solver.set_parameters(params)
    if has_loss:
        loss_fn.d = loss_d
        loss_fn.g = loss_g
    return pred, loss_fn, input_image


def get_config(args, is_eval=False):
    class Congig(object):
        def __init__(self, cfg, is_eval):
            lr = 0.05
            batch_size = 64

            self.model_info_dict = {
                'lr': lr,
                'batch_size': batch_size,
                'device_id': cfg.device_id,
                'seed': cfg.seed,
                'net_func': cifarnet,
                'num_epochs': cfg.n_epochs,
            }

            self.file_dir_dict = {
                'train_csv': cfg.input_train,
                'val_csv': cfg.input_val,
                'score_filename': cfg.score_output,
                'save_dir': cfg.weight_output,
                'infl_filename': cfg.output,
                'model_filename': 'final_model.h5',
                'info_filename': 'info.npy'
            }
            if is_eval:
                self._init_for_eval(cfg)
            else:
                self._init_for_train_infl(cfg)

        def _init_for_eval(self, cfg):
            self.file_dir_dict['test_csv'] = cfg.input_test
            start_epoch = 0 if cfg.retrain_all else cfg.n_epochs - 1
            self.model_info_dict['start_epoch'] = start_epoch

        def _init_for_train_infl(self, cfg):
            self.calc_infl_with_all_params = cfg.calc_infl_with_all_params
            self.need_evaluate = False if cfg.score_output is None else True
            end_epoch_dict = {'last': cfg.n_epochs - 1, 'all': 0}
            self.model_info_dict['end_epoch'] = end_epoch_dict[cfg.calc_infl_method]

    return Congig(args, is_eval)


# responsible_ai/data_cleansing/sgd_influence/train.py

# Copyright 2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import numpy as np
from tqdm import tqdm
import nnabla as nn
import functools
import nnabla.solvers as S
from .model import setup_model, calc_acc, adjust_batch_size
from .dataset import get_batch_indices, get_batch_data, init_dataset, get_image_size
from .utils import ensure_dir, get_indices, save_to_csv


def save_all_params(params_dict, c, k, j, bundle_size, step_size, save_dir, epoch):
    params_dict[c] = nn.get_parameters(grad_only=False).copy()
    c += 1
    if c == bundle_size or j == step_size - 1:
        dn = os.path.join(save_dir, 'epoch%02d' % (epoch), 'weights')
        ensure_dir(dn)
        for cc, params in params_dict.items():
            fn = '%s/model_step%04d.h5' % (dn, k + cc)
            nn.save_parameters(fn, params=params, extension=".h5")
        k += c
        c = 0
        params_dict = {}
    return params_dict, c, k


def eval_model(val_model, solver, dataset, idx_list_to_data, batch_size, resize_size):
    loss = 0
    acc = 0
    n = len(idx_list_to_data)
    idx = np.array_split(np.arange(n), batch_size)
    loss_fn = None
    for _, i in enumerate(idx):
        X, y = get_batch_data(dataset, idx_list_to_data,
                              i, resize_size, test=True)
        pred, loss_fn, input_image = adjust_batch_size(
            val_model, solver, len(X), loss_fn)
        input_image["image"].d = X
        input_image["label"].d = y
        loss_fn.forward()
        loss += loss_fn.d * len(X)
        acc += calc_acc(pred.d, y, method='sum')
    loss /= n
    acc /= n
    return loss, acc


def train(model_info_dict, file_dir_dict, use_all_params, need_evaluate, bundle_size=200):
    # params
    lr = model_info_dict['lr']
    seed = model_info_dict['seed']
    net_func = model_info_dict['net_func']
    batch_size = model_info_dict['batch_size']
    num_epochs = model_info_dict['num_epochs']
    infl_end_epoch = model_info_dict['end_epoch']
    # files and dirs
    save_dir = file_dir_dict['save_dir']
    info_filename = file_dir_dict['info_filename']
    model_filename = file_dir_dict['model_filename']
    score_filename = file_dir_dict['score_filename']
    # setup
    trainset, valset, image_shape, n_classes, ntr, nval = init_dataset(
        file_dir_dict['train_csv'], file_dir_dict['val_csv'], seed)
    n_channels, _h, _w = image_shape
    resize_size = get_image_size((_h, _w))
    # Create training graphs
    test = False
    train_model = functools.partial(
        setup_model, net_func=net_func, n_classes=n_classes, n_channels=n_channels, resize_size=resize_size, test=test)
    # Create validation graphs
    test = True
    val_model = functools.partial(
        setup_model, net_func=net_func, n_classes=n_classes, n_channels=n_channels, resize_size=resize_size, test=test)
    # setup optimizer (SGD)
    solver = S.Sgd(lr=lr)
    solver.set_parameters(nn.get_parameters(grad_only=False))

    # get shuffled index using designated seed
    idx_train = get_indices(ntr, seed)
    idx_val = get_indices(nval, seed)

    # training
    seed_train = 0
    info = []
    score = []
    loss_train = None
    for epoch in tqdm(range(num_epochs), desc='training (1/3 steps)'):
        idx = get_batch_indices(ntr, batch_size, seed=epoch)
        epoch_info = []
        c = 0
        k = 0
        params_dict = {}
        for j, i in enumerate(idx):
            seeds = list(range(seed_train, seed_train + i.size))
            seed_train += i.size
            epoch_info.append({'epoch': epoch, 'step': j,
                               'idx': i, 'lr': lr, 'seeds': seeds})
            if (use_all_params) & (epoch >= infl_end_epoch):
                params_dict, c, k = save_all_params(
                    params_dict, c, k, j, bundle_size, len(idx), save_dir, epoch)
            X, y = get_batch_data(trainset, idx_train, i,
                                  resize_size, test=False, seeds=seeds)
            _, loss_train, input_image_train = adjust_batch_size(
                train_model, solver, len(X), loss_train)
            input_image_train["image"].d = X
            input_image_train["label"].d = y

            loss_train.forward()
            solver.zero_grad()
            loss_train.backward(clear_buffer=True)
            solver.update()
        info.append(epoch_info)
        # save if params are necessary for calculating influence
        if epoch >= infl_end_epoch-1:
            dn = os.path.join(save_dir, 'epoch%02d' % (epoch), 'weights')
            ensure_dir(dn)
            nn.save_parameters(os.path.join(dn, model_filename), params=nn.get_parameters(
                grad_only=False), extension=".h5")
        # evaluation
        if need_evaluate:
            loss_tr, acc_tr = eval_model(
                val_model, solver, trainset, idx_train, batch_size, resize_size)
            loss_val, acc_val = eval_model(
                val_model, solver, valset, idx_val, batch_size, resize_size)
            score.append((loss_tr, loss_val, acc_tr, acc_val))
    # save epoch and step info
    np.save(os.path.join(save_dir, info_filename), arr=info)
    # save score
    if need_evaluate:
        save_to_csv(filename=score_filename, header=[
                    'train_loss', 'val_loss', 'train_accuracy', 'val_accuracy'], list_to_save=score, data_type='float,float,float,float')


# responsible_ai/data_cleansing/sgd_influence/utils.py

# Copyright 2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import csv
import numpy as np
from nnabla.ext_utils import get_extension_context
from nnabla.utils.load import load
from shutil import rmtree


def delete_file(file_name):
    if os.path.isfile(file_name):
        os.remove(file_name)


def delete_dir(dir_name, keyword='sgd_infl_results'):
    if os.path.isdir(dir_name):
        if keyword in dir_name:
            rmtree(dir_name)


def ensure_dir(dir):
    if not os.path.exists(dir):
        os.makedirs(dir)


def get_context(device_id):
    # for cli app use
    try:
        context = 'cudnn'
        ctx = get_extension_context(context, device_id=device_id)
    except ModuleNotFoundError:
        context = 'cpu'
        ctx = get_extension_context(context, device_id=device_id)
    # for nnc use
    config_filename = 'net.nntxt'
    if os.path.isfile(config_filename):
        config_info = load([config_filename])
        ctx = config_info.global_config.default_context

    return ctx


def get_indices(n, seed):
    np.random.seed(seed)
    idx = np.random.permutation(n)
    return idx


def save_to_csv(filename, header, list_to_save, data_type):
    with open(filename, 'w') as f:
        writer = csv.writer(f, lineterminator='\n')
        writer.writerow(header)
        writer.writerows(np.array([tuple(row)
                                   for row in list_to_save], dtype=data_type))


def read_csv(filename):
    with open(filename) as f:
        reader = csv.reader(f)
        ret = [s for s in reader]
    return ret


# responsible_ai/gan_data_debiased/utils/utils.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
from sklearn.metrics import average_precision_score, f1_score, recall_score
from scipy.ndimage import gaussian_filter


def get_all_attr():
    """
    get all attributes in celebA dataset
    """
    return ['5_o_Clock_Shadow', 'Arched_Eyebrows', 'Attractive', 'Bags_Under_Eyes',
            'Bald', 'Bangs', 'Big_Lips', 'Big_Nose', 'Black_Hair', 'Blond_Hair',
            'Blurry', 'Brown_Hair', 'Bushy_Eyebrows', 'Chubby', 'Double_Chin',
            'Eyeglasses', 'Goatee', 'Gray_Hair', 'Heavy_Makeup', 'High_Cheekbones',
            'Male', 'Mouth_Slightly_Open', 'Mustache', 'Narrow_Eyes', 'No_Beard',
            'Oval_Face', 'Pale_Skin', 'Pointy_Nose', 'Receding_Hairline', 'Rosy_Cheeks',
            'Sideburns', 'Smiling', 'Straight_Hair', 'Wavy_Hair', 'Wearing_Earrings',
            'Wearing_Hat', 'Wearing_Lipstick', 'Wearing_Necklace', 'Wearing_Necktie', 'Young']


def compute_class_weight(loader):
    """
    Compute sample weights
    """

    cp = 0
    cn = 0
    cn_dn = 0
    cn_dp = 0
    cp_dn = 0
    cp_dp = 0

    weights = []
    max_iter = int(loader.size/32)
    for i in range(max_iter):
        _, targets = loader.next()
        class_label = targets[:, 0]
        domain_label = targets[:, 1]
        cp += class_label.sum()  # class is positive
        cn += (targets.shape[0] - class_label.sum())  # class is negative
        # class is negative, domain is negative
        cn_dn += ((class_label + domain_label) == 0).sum()
        cn_dp += ((class_label - domain_label) == -1).sum()
        cp_dn += ((class_label - domain_label) == 1).sum()
        cp_dp += ((class_label + domain_label) == 2).sum()
    for i in range(max_iter):
        _, targets = loader.next()
        class_label = targets[:, 0]
        domain_label = targets[:, 1]

        weights.append(
            (class_label*cp + (1-class_label)*cn) /
            (2*(
                    (1-class_label)*(1-domain_label)*cn_dn
                    + (1-class_label)*domain_label*cn_dp
                    + class_label*(1-domain_label)*cp_dn
                    + class_label*domain_label*cp_dp
                   )
             )
        )
    weights = np.concatenate(weights)
    return weights


def get_bias_amplification_attribute_task(running_labels, running_preds):
    """
    compute bias amplification for attribute-task

    References:
        [1] Directional Bias Amplification(https://arxiv.org/abs/2102.12594)
        [2] Fair Attribute Classification through Latent Space
        De-biasing(https://arxiv.org/abs/2012.01469)
    Args:
        running_labels : labels on test set
        running_preds : predictions on test set
    Returns:
         bias amplification
    """
    bog_tilde = np.zeros((2, 2))
    bog_gt_g = np.zeros((2, 2))
    for i, objs in enumerate([running_labels, running_preds]):
        female = np.where(objs[:, 1] == 0)[0]  # Unprivileged class
        male = np.where(objs[:, 1] == 1)[0]  # Privileged class
        kitchen = np.where(objs[:, 0] == 0)[0]  # Unfavourable class
        sports = np.where(objs[:, 0] == 1)[0]  # Favourable class
        if i == 0:
            bog_tilde[0][0] = len(set(kitchen) & set(female))
            bog_tilde[0][1] = len(set(kitchen) & set(male))
            bog_tilde[1][0] = len(set(sports) & set(female))
            bog_tilde[1][1] = len(set(sports) & set(male))
        elif i == 1:
            bog_gt_g[0][0] = len(set(kitchen) & set(female))
            bog_gt_g[0][1] = len(set(kitchen) & set(male))
            bog_gt_g[1][0] = len(set(sports) & set(female))
            bog_gt_g[1][1] = len(set(sports) & set(male))

    total_images_train = np.sum(bog_tilde)
    data_bog = bog_tilde / np.sum(bog_tilde, axis=0)
    pred_bog = bog_gt_g / np.sum(bog_tilde, axis=0)
    p_t_a = bog_tilde / np.sum(bog_tilde, axis=0)
    p_t = np.sum(bog_tilde, axis=1) / total_images_train
    diff = np.zeros_like(data_bog)
    for i in range(len(data_bog)):
        for j in range(len(data_bog[0])):
            diff[i][j] = pred_bog[i][j] - data_bog[i][j]
            indicator = np.sign(p_t_a[i][j] - p_t[i])
            if indicator == 0:
                diff[i][j] = 0
            elif indicator == -1:
                diff[i][j] = - diff[i][j]
    value = np.nanmean(diff)
    return value


def get_bias_amplification(domain, targets, pred):
    """
    BA measures how much more often a target attribute is predicted with a
    protected attribute than the ground truth value.

    References:
        [1] Directional Bias Amplification(https://arxiv.org/abs/2102.12594)
        [2] Fair Attribute Classification through Latent Space
        De-biasing(https://arxiv.org/abs/2012.01469)
    Args:
        domain (numpy.ndarray) : actual protected attribute
        targets (numpy.ndarray) : actual target label
        pred (numpy.ndarray) : predicted label
    Returns:
        Bias amplification
    """

    test_labels = np.zeros((targets.shape[0], 2))
    test_labels[:, 0] = targets
    test_labels[:, 1] = domain
    test_pred = np.zeros((targets.shape[0], 2))
    test_pred[:, 0] = pred
    test_pred[:, 1] = domain
    auc_bias = get_bias_amplification_attribute_task(test_labels, test_pred)
    return auc_bias


def get_difference_equality_opportunity(domain, targets, pred):
    """
    Compute the absolute difference in FNR between the protected attribute group.

    Args:
        domain (numpy.ndarray) : actual protected attribute
        targets (numpy.ndarray) : actual target label
        pred (numpy.ndarray) : predicted label
    Returns:
        difference_in_equla_opportunity(float)

    """

    g0 = np.argwhere(domain == 0)
    g1 = np.argwhere(domain == 1)
    deo = np.abs(
        (1-recall_score(targets[g0], pred[g0]))-(1-recall_score(targets[g1], pred[g1])))

    return np.median(deo)


def get_average_precision(targets, scores):
    """
    Compute average precision score
    Args:
        targets (numpy.ndarray) : actual target label
        scores (numpy.ndarray) : predicted scores
    Returns:
        Average precission
    """

    avg_prec = average_precision_score(targets, scores)

    return np.median(avg_prec)


def get_f1_threshold(targets_all, scores_all):
    """
    get the f1 threshold and accuracy
    Args:
        targets_all (numpy.ndarray) : actual target label
        scores : predicted scores
    Returns:
        best_acc : best accuracy
        best_t : best threshold
    """
    best_t = -1.0
    best_acc = 0.0
    for t in range(1, 10):
        thresh = 0.1*t
        curr_scores = np.where(scores_all > thresh, 1, 0)
        acc = f1_score(targets_all, curr_scores)
        if acc > best_acc:
            best_acc = acc
            best_t = thresh
    one_dec = best_t

    for t in range(1, 20):
        thresh = (one_dec-0.1) + 0.01*t
        curr_scores = np.where(scores_all > thresh, 1, 0)
        acc = f1_score(targets_all, curr_scores)
        # print(thresh, acc, best_acc, flush=True)
        if acc > best_acc:
            best_acc = acc
            best_t = thresh

    return best_acc, best_t


def calibrated_threshold(targets, scores):
    """
    Calibrated threshold
    Args:
        targets (numpy.ndarray) : actual target label
        scores (numpy.ndarray) : predicted scores
    Returns:
        calibrated threshold
    """
    cp = int(targets.sum())
    scores_copy = np.copy(scores)
    scores_copy.sort()
    thresh = scores_copy[-cp]
    return thresh


def get_kl(p, q):
    """
    Kullback-Leibler divergence D(P || Q) for discrete distributions
    Args:
         p,q (float array): Discrete probability distributions.
    Returns:
        kl
    """
    p = np.asarray(p, dtype=np.float)
    q = np.asarray(q, dtype=np.float)
    return np.sum(np.where(p != 0, p * np.log(p / q), 0))


def smoothed_hist_kl_distance(a, b, nbins=50, sigma=1):
    """
    smooth two histograms with Standard kernel (gaussian kernel with mean = 0 ,sigma=1),
    and calculate the KL distance between two smoothed histograms.
    Args:
        a (numpy.ndarray) : numpy nd array
        b (numpy.ndarray) : numpy nd array
        nbins (int) : number of histogram bins
        sigma : standard deviation for Gaussian kernel
    Returns:
        kl distance
    """
    ahist = np.histogram(a, bins=nbins)[0]
    bhist = np.histogram(b, bins=nbins)[0]

    asmooth = gaussian_filter(ahist, sigma)
    bsmooth = gaussian_filter(bhist, sigma)

    asmooth = asmooth/asmooth.sum() + 1e-6
    bsmooth = bsmooth/bsmooth.sum() + 1e-6

    return get_kl(asmooth, bsmooth), get_kl(bsmooth, asmooth)


def get_kl_divergence(domain, targets, scores):
    """
    Measures the divergence between score distributions (KL/threshold-invariant metric)
    References:
        [1] Towards Threshold Invariant Fair Classification(https://arxiv.org/abs/2102.12594)
        [2] Fair Attribute Classification through Latent Space De-biasing(https://arxiv.org/abs/2012.01469)
    Args:
        domain (numpy.ndarray) : actual protected attribute
        targets (numpy.ndarray) : actual target label
        scores (numpy.ndarray) : predicted scores
    Returns:
        kl_divergence
    """

    nbin = 50  # Number of histogram bins

    # a_b, b_a = smoothed_hist_kl_distance(scores[domain == 0], scores[domain == 1], nbins=nbin)
    a_b_pos, b_a_pos = smoothed_hist_kl_distance(scores[np.logical_and(domain == 0, targets == 1)],
                                                 scores[np.logical_and(domain == 1, targets == 1)], nbins=nbin)
    a_b_neg, b_a_neg = smoothed_hist_kl_distance(scores[np.logical_and(domain == 0, targets == 0)],
                                                 scores[np.logical_and(domain == 1, targets == 0)], nbins=nbin)

    return np.median([a_b_pos]+[a_b_neg]+[b_a_pos]+[b_a_neg])


# responsible_ai/prejudice_remover_regularizer_images/utils/utils.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import average_precision_score, f1_score, recall_score


def get_auc_bias(gt_label, pred_label):
    """
    compute bias amplification for attribute-task
    References:
        [1] Directional Bias Amplification(https://arxiv.org/abs/2102.12594)
        [2] Fair Attribute Classification through Latent Space
        De-biasing(https://arxiv.org/abs/2012.01469)
    Args:
        gt_label: ground truth labels
        pred_label: predicated labels
    Returns:
         bias amplification
    """
    bog_tilde = np.zeros((2, 2))
    bog_gt_g = np.zeros((2, 2))
    for i, objs in enumerate([gt_label, pred_label]):
        female = np.where(objs[:, 1] == 0)[0]  # Unprivileged class
        male = np.where(objs[:, 1] == 1)[0]  # Privileged class
        kitchen = np.where(objs[:, 0] == 0)[0]  # Unfavourable class
        sports = np.where(objs[:, 0] == 1)[0]  # Favourable class
        if i == 0:
            bog_tilde[0][0] = len(set(kitchen) & set(female))
            bog_tilde[0][1] = len(set(kitchen) & set(male))
            bog_tilde[1][0] = len(set(sports) & set(female))
            bog_tilde[1][1] = len(set(sports) & set(male))
        elif i == 1:
            bog_gt_g[0][0] = len(set(kitchen) & set(female))
            bog_gt_g[0][1] = len(set(kitchen) & set(male))
            bog_gt_g[1][0] = len(set(sports) & set(female))
            bog_gt_g[1][1] = len(set(sports) & set(male))

    total_images_train = np.sum(bog_tilde)
    data_bog = bog_tilde / np.sum(bog_tilde, axis=0)
    pred_bog = bog_gt_g / np.sum(bog_tilde, axis=0)
    p_t_a = bog_tilde / np.sum(bog_tilde, axis=0)
    p_t = np.sum(bog_tilde, axis=1) / total_images_train
    diff = np.zeros_like(data_bog)
    for i in range(len(data_bog)):
        for j in range(len(data_bog[0])):
            diff[i][j] = pred_bog[i][j] - data_bog[i][j]
            indicator = np.sign(p_t_a[i][j] - p_t[i])
            if indicator == 0:
                diff[i][j] = 0
            elif indicator == -1:
                diff[i][j] = - diff[i][j]
    value = np.nanmean(diff)
    return value


def get_bias_amplification(domain, targets, pred):
    """
    BA measures how much more often a target attribute is predicted with a
    protected attribute than the ground truth value.
    References:
        [1] Directional Bias Amplification(https://arxiv.org/abs/2102.12594)
        [2] Fair Attribute Classification through Latent Space
        De-biasing(https://arxiv.org/abs/2012.01469)
    Args:
        domain (numpy.ndarray): protected attribute
        targets (numpy.ndarray): target label
        pred (numpy.ndarray): predicted label
    Returns:
        Bias amplification
    """

    test_labels = np.zeros((targets.shape[0], 2))
    test_labels[:, 0] = targets
    test_labels[:, 1] = domain
    test_pred = np.zeros((targets.shape[0], 2))
    test_pred[:, 0] = pred
    test_pred[:, 1] = domain
    auc_bias = get_auc_bias(test_labels, test_pred)
    return auc_bias


def get_diff_in_equal_opportunity(domain, targets, pred):
    """
    Compute the absolute difference in FNR between the protected attribute group.
    Args:
        domain (numpy.ndarray) : actual protected attribute
        targets (numpy.ndarray) : actual target label
        pred (numpy.ndarray) : predicted label
    Returns:
        difference_in_equal_opportunity(float)
    """

    g0 = np.argwhere(domain == 0)
    g1 = np.argwhere(domain == 1)
    deo = np.abs(
        (1 - recall_score(targets[g0], pred[g0])) - (1 - recall_score(targets[g1], pred[g1])))

    return np.median(deo)


def get_average_precision(targets, scores):
    """
    Compute average precision score
    Args:
        targets (numpy.ndarray) : actual target label
        scores (numpy.ndarray) : predicted scores
    Returns:
        Average precission
    """

    avg_prec = average_precision_score(targets, scores)

    return np.median(avg_prec)


def get_f1_threshold(targets_all, scores_all):
    """
    get the f1 threshold and accuracy
    Args:
        targets_all (numpy.ndarray) : actual target label
        scores : predicted scores
    Returns:
        best_acc : best accuracy
        best_t : best threshold
    """
    best_t = -1.0
    best_acc = 0.0
    for t in range(1, 10):
        thresh = 0.1 * t
        curr_scores = np.where(scores_all > thresh, 1, 0)
        acc = f1_score(targets_all, curr_scores)
        if acc > best_acc:
            best_acc = acc
            best_t = thresh
    one_dec = best_t

    for t in range(1, 20):
        thresh = (one_dec - 0.1) + 0.01 * t
        curr_scores = np.where(scores_all > thresh, 1, 0)
        acc = f1_score(targets_all, curr_scores)
        # print(thresh, acc, best_acc, flush=True)
        if acc > best_acc:
            best_acc = acc
            best_t = thresh

    return best_acc, best_t


def calibrated_threshold(targets, scores):
    """
    Calibrated threshold
    Args:
        targets (numpy.ndarray) : actual target label
        scores (numpy.ndarray) : predicted scores
    Returns:
        calibrated threshold
    """
    cp = int(targets.sum())
    scores_copy = np.copy(scores)
    scores_copy.sort()
    thresh = scores_copy[-cp]
    return thresh


def get_cvs(output_f, output_m, cal_thresh):
    """
    Calders and Verwer defined a discrimination score,
    by subtracting the conditional probability of the positive class given a sensitive value
    from that given a non-sensitive value.
    Args:
        output_f (list): output of unprivileged class
        output_m (list) : output of privileged class
        cal_thresh (float) : calibrated threshold
    Returns:
        CV Score (float)
    """
    yf_pred = (output_f >= cal_thresh)
    ym_pred = (output_m >= cal_thresh)
    corr_f = np.sum(yf_pred == True)
    corr_m = np.sum(ym_pred == True)
    P_y1_s1 = corr_f / output_f.shape[0]
    P_y1_s0 = corr_m / output_m.shape[0]
    CV_score = np.abs(P_y1_s0 - P_y1_s1)
    return round(CV_score.item(), 4)


def plot_fairness(fairness, ax, metric="DPD", title="fairness metric", bar_x_axis="Original"):
    """
    plot single fairness metric
    """
    ax.set_ylim([-0.6, 0.6])
    ax.axhline(y=0.0, color='r', linestyle='-')
    ax.bar([bar_x_axis], fairness, color="blue", width=2)
    ax.set_ylabel(metric)
    ax.set_title(title, fontsize=10)
    for index, value in enumerate(fairness):
        if value < 0:
            ax.text(index, value - 0.1,
                    str(round(value, 3)), fontweight='bold', color='red',
                    bbox=dict(facecolor='red', alpha=0.4))
        else:
            ax.text(index, value + 0.1,
                    str(round(value, 3)), fontweight='bold', color='red',
                    bbox=dict(facecolor='red', alpha=0.4))


def plot_fairness_multi(DEO, CV_score, BA, accuracy, bar_x_axis="original"):
    """
    plot fairness metrics
    """
    fig, axes = plt.subplots(1, 3, figsize=(10, 4), sharey=True)
    fig.suptitle("Model Fairness", fontsize=16)
    plot_fairness([DEO], ax=axes[0], metric="DEO",
                  title="Difference in Equal opportunity(DEO)", bar_x_axis=bar_x_axis)
    plot_fairness([CV_score], ax=axes[1], metric="CV Score",
                  title="CV score", bar_x_axis=bar_x_axis)
    plot_fairness([BA], ax=axes[2], metric="BA",
                  title="Bias Amplification (BA)", bar_x_axis=bar_x_axis)

    fig.text(0.92, 0.65, '\n'.join(
        ["Average precession score:", f"- AP : {accuracy:.3f}"]), fontsize='15')
    plt.show()


# responsible_ai/rejection_option_based_classification_images/utils/utils.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import matplotlib.pyplot as plt
from sklearn import metrics
from sklearn.metrics import average_precision_score
from sklearn.metrics import confusion_matrix


def get_best_classification_thershold(val_scores, val_targets):
    """
    Compute best classfication thershold on validation data set
    Args:
        val_scores (numpy.ndarray): estimated probability predictions (targets score)
                                    as returned by a classifier on validation dataset.
        val_targets (numpy.ndarray) :  actual validation taget labels
    Returns:
        best_acc (float): best accuracy on validation dataset.
        best_classification_thershold (float): best classification thershold on validation dataset.
    """

    num_thresh = 100
    ba_arr = np.zeros(num_thresh)
    classification_thershold_array = np.linspace(0.01, 0.99, num_thresh)
    for idx, classification_thershold in enumerate(classification_thershold_array):
        y_val_pred = np.where(val_scores > classification_thershold, 1, 0)
        val_acc = metrics.accuracy_score(val_targets, y_val_pred)
        ba_arr[idx] = val_acc
    best_ind = np.where(ba_arr == np.max(ba_arr))[0][0]
    best_classification_thershold = classification_thershold_array[best_ind]
    best_acc = np.max(ba_arr)
    print("Best balanced accuracy  = %.4f" % best_acc)
    print("Optimal classification threshold = %.4f" %
          best_classification_thershold)
    return best_acc, best_classification_thershold


def get_average_precision(targets, scores):
    """
    Compute average precision score
    Args:
        targets (numpy.ndarray) : actual target label
        scores (numpy.ndarray) : predicted scores
    Returns:
        Average precission
    """

    avg_prec = average_precision_score(targets, scores)

    return np.median(avg_prec)


def get_fairness(y, privileged_group, preds):
    """
    Compute multiple fairness metrics for the classifier:
    1. Demographic parity
    2. Equal opportunity
    3. Equalized odds
    Args:
        y (numpy.ndarray): True data (or target, ground truth)
        privileged_group (numpy.ndarray): list of privileged group values
        preds (numpy.ndarray): data predicted (calculated, output) by your model
    Returns:
        demographic_parity_difference (float): Demographic parity
        equal_opportunity_difference (float): Equal opportunity
        average_abs_odds_difference(float): Equalized odds
    """

    y_unprivileged, preds_unprivileged = y[privileged_group ==
                                           False], preds[privileged_group == False]
    y_privileged, preds_privileged, = y[privileged_group], preds[privileged_group]
    cm_unprivileged = confusion_matrix(y_unprivileged, preds_unprivileged)
    cm_privileged = confusion_matrix(y_privileged, preds_privileged)
    unprivileged_PR = (
        cm_unprivileged[1, 1] + cm_unprivileged[0, 1]) / cm_unprivileged.sum()
    privileged_PR = (cm_privileged[1, 1] +
                     cm_privileged[0, 1]) / cm_privileged.sum()

    # compute demographic parity
    demographic_parity_difference = unprivileged_PR - privileged_PR
    unprivileged_TPR = cm_unprivileged[1, 1] / cm_unprivileged[1].sum()
    privileged_TPR = cm_privileged[1, 1] / cm_privileged[1].sum()

    # compute equal opportunity
    equal_opportunity_difference = unprivileged_TPR - privileged_TPR
    unprivileged_FPR = cm_unprivileged[0, 1] / cm_unprivileged[0].sum()
    privileged_FPR = cm_privileged[0, 1] / cm_privileged[0].sum()

    # compute Equalized odds
    average_abs_odds_difference = 0.5 * \
        (abs(unprivileged_FPR - privileged_FPR) +
         abs(unprivileged_TPR - privileged_TPR))

    return demographic_parity_difference, equal_opportunity_difference, average_abs_odds_difference


def plot_fairness(fairness, ax, metric="DPD", title="fairness metric", bar_x_axis="Original"):
    """
    plot single fairness metric
    """

    ax.set_ylim([-0.6, 0.6])
    ax.axhline(y=0.0, color='r', linestyle='-')
    ax.bar([bar_x_axis], fairness, color="blue", width=2)
    ax.set_ylabel(metric)
    ax.set_title(title, fontsize=10)

    for index, value in enumerate(fairness):
        if value < 0:
            ax.text(index, value - 0.1, str(round(value, 3)), fontweight='bold',
                    color='red', bbox=dict(facecolor='red', alpha=0.4))
        else:
            ax.text(index, value + 0.1, str(round(value, 3)), fontweight='bold',
                    color='red', bbox=dict(facecolor='red', alpha=0.4))


def plot_fairness_multi(dpd, eod, aaod, accuracy, bar_x_axis="original"):
    """
    plot three fairness metrics
    Args:
        dpd (float) : demographic parity difference
        eod (float) : equal opportunity difference
        aaod (float) : absolute average odd difference
        accuracy (float) : accuracy of the model
        bar_x_axis (str): bar name of the X-axis
    """
    fig, axes = plt.subplots(1, 3, figsize=(10, 4), sharey=True)
    fig.suptitle("Model Fairness", fontsize=16)
    plot_fairness([dpd], ax=axes[0], metric="DPD",
                  title="Demographic Parity", bar_x_axis=bar_x_axis)
    plot_fairness([eod], ax=axes[1], metric="EOD",
                  title="Equal Opportunity", bar_x_axis=bar_x_axis)
    plot_fairness([aaod], ax=axes[2], metric="AAOD",
                  title="Equalized Odds", bar_x_axis=bar_x_axis)
    fig.text(0.92, 0.65, '\n'.join(
        ["Accuracy of the model:", f"- accuracy : {accuracy:.3f}"]), fontsize='15')
    plt.show()


# data_augmentation/MixedDataLearning.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import nnabla as nn
import nnabla.functions as F


def softmax_cross_entropy_loss_vlabel(pred, vlabel):
    # The shape of vlabel is supposed to be (batch_size, n_class)
    logp = F.log_softmax(pred)
    loss = -1.0 * F.mean(F.sum(vlabel * logp, axis=1))
    return loss

# Abstract class for learning with mixed data


class MixedDataLearning(object):
    def __init__(self):
        # Set params for mixing data
        return NotImplemented

    def set_mix_ratio():
        return NotImplemented

    def mix_data(self, x, y):
        # Mix data given x and y, and return mix_x and mix_y
        # x, y, mix_x, and mix_y are supposed to be nn.Variable
        return NotImplemented

    def loss(self, pred, mix_y):
        # Calculate a classification loss given mix_y and prediction results of mix_x.
        # Both pred and mix_y are supposed to be a nn.Variable
        return NotImplemented


# Mixup
class MixupLearning(MixedDataLearning):
    def __init__(self, batch_size, alpha=0.5):
        # Set params for mixing data
        # For mixup, set alpha for the beta distribution that generates interpolation ratios.
        self._batch_size = batch_size
        self._alpha = alpha
        self._lam = nn.Variable((batch_size, 1))

    def set_mix_ratio(self):
        if self._alpha > 0.0:
            self._lam.d = np.random.beta(
                self._alpha, self._alpha, self._batch_size).reshape((self._batch_size, 1))
        else:
            self._lam.d = np.ones((self._batch_size, 1))

    def mix_data(self, x, y):
        # Mix data given x and y, and return mix_x and mix_y
        # Both y and mix_y are supposed to be nn.Variable((batch_size, n_class))
        batch_size = x.shape[0]
        ind = np.random.permutation(batch_size)
        x0 = x
        y0 = y
        x1 = x0[ind]
        y1 = y0[ind]
        mix_x = self._lam.reshape((-1, 1, 1, 1)) * x0 + \
            (1.0-self._lam.reshape((-1, 1, 1, 1))) * x1
        mix_y = self._lam * y0 + (1.0-self._lam) * y1
        return mix_x, mix_y

    def loss(self, pred, mix_y):
        # Calculate a classification loss given mix_y and prediction results of mix_x.
        # Both pred and mix_y are supposed to be a nn.Variable
        return softmax_cross_entropy_loss_vlabel(pred, mix_y)


# Cutmix
class CutmixLearning(MixedDataLearning):
    def __init__(self, shape_of_batch, alpha=0.5, cutmix_prob=0.5):
        # Set params for mixing data
        # For cutmix, set alpha for the beta distribution that generates cutting area ratio.
        # cutmix_prob controls a probablity to conduct cutmix for each batch
        # shape_of_batch should be (batch_size, n_channels, height, width)
        self._batch_size = shape_of_batch[0]
        self._alpha = alpha
        self._cutmix_prob = cutmix_prob
        self._mask = nn.Variable(shape_of_batch)
        self._lam = nn.Variable((shape_of_batch[0], 1))

    def set_mix_ratio(self):
        # How to get a random bounding box
        def rand_bbox(shape_of_x, lam):
            width = shape_of_x[3]
            height = shape_of_x[2]
            cut_ratio = np.sqrt(1.0 - lam)
            cut_w = np.int(width * cut_ratio)
            cut_h = np.int(height * cut_ratio)

            cx = np.random.randint(width)
            cy = np.random.randint(height)

            bbx0 = np.clip(cx - cut_w//2, 0, width)
            bby0 = np.clip(cy - cut_h//2, 0, height)
            bbx1 = np.clip(cx + cut_w//2, 0, width)
            bby1 = np.clip(cy + cut_h//2, 0, height)

            return bbx0, bby0, bbx1, bby1

        def get_mask(shape_of_x, bbx0, bby0, bbx1, bby1):
            mask = np.zeros(shape_of_x)
            mask[:, :, bby0:bby1, bbx0:bbx1] = 1.0
            return mask

        if self._alpha > 0.0 and np.random.rand() <= self._cutmix_prob:
            lam_tmp = np.random.beta(self._alpha, self._alpha)
            bbx0, bby0, bbx1, bby1 = rand_bbox(self._mask.shape, lam_tmp)
            self._mask.d = get_mask(self._mask.shape, bbx0, bby0, bbx1, bby1)
            self._lam.d = (1.0 - ((bbx1-bbx0)*(bby1-bby0)/(
                self._mask.shape[2]*self._mask.shape[3]))) * np.ones((self._batch_size, 1))
        else:
            self._mask.d = np.zeros(self._mask.shape)
            self._lam.d = np.ones((self._batch_size, 1))

    def mix_data(self, x, y):
        # Mix data given x and y, and return mix_x and mix_y
        # Both y and mix_y are supposed to be nn.Variable((batch_size, n_class))
        batch_size = x.shape[0]
        ind = np.random.permutation(batch_size)
        x0 = x
        y0 = y
        x1 = x0[ind]
        y1 = y0[ind]
        mix_x = (1.0 - self._mask) * x0 + self._mask * x1
        mix_y = self._lam * y0 + (1.0 - self._lam) * y1
        return mix_x, mix_y

    def loss(self, pred, mix_y):
        # Calculate a classification loss given mix_y and prediction results of mix_x.
        # Both pred and mix_y are supposed to be a nn.Variable
        return softmax_cross_entropy_loss_vlabel(pred, mix_y)


# VH-Mixup
class VHMixupLearning(MixedDataLearning):
    def __init__(self, shape_of_batch, alpha=0.5):
        # Set params for mixing data
        # For vh-mixup, set alpha for the beta distribution that generates interpolation ratios.
        # shape_of_batch should be (batch_size, n_channels, height, width)
        self._batch_size = shape_of_batch[0]
        self._maskv = nn.Variable(shape_of_batch)
        self._maskh = nn.Variable(shape_of_batch)
        self._lamx = nn.Variable([shape_of_batch[0], 1])
        self._lamy = nn.Variable([shape_of_batch[0], 1])
        self._alpha = alpha

    def set_mix_ratio(self):
        # How to concatenate images
        def get_maskv(alpha):
            if alpha <= 0.0:
                return np.ones(self._maskv.shape), 1.0
            mask = np.zeros(self._maskv.shape)
            lam = np.random.beta(self._alpha, self._alpha)
            lh = np.int(lam * self._maskv.shape[2])
            mask[:, :, 0:lh, :] = 1.0
            return mask, lam

        def get_maskh(alpha):
            if alpha <= 0.0:
                return np.ones(self._maskh.shape), 1.0
            mask = np.zeros(self._maskh.shape)
            lam = np.random.beta(self._alpha, self._alpha)
            lw = np.int(lam * self._maskh.shape[3])
            mask[:, :, :, 0:lw] = 1.0
            return mask, lam

        self._maskv.d, lam1 = get_maskv(self._alpha)
        self._maskh.d, lam2 = get_maskh(self._alpha)
        if self._alpha > 0.0:
            self._lamx.d = np.random.beta(
                self._alpha, self._alpha, self._batch_size).reshape((self._batch_size, 1))
        else:
            self._lamx.d = np.ones((self._batch_size, 1))
        self._lamy.d = lam1 * self._lamx.d + lam2 * (1.0 - self._lamx.d)

    def mix_data(self, x, y):
        # Mix data given x and y, and return mix_x and mix_y
        # Both y and mix_y are supposed to be nn.Variable((batch_size, n_class))
        batch_size = x.shape[0]
        ind = np.random.permutation(batch_size)
        x0 = x
        y0 = y
        x1 = x0[ind]
        y1 = y0[ind]
        x_hcat = self._maskh * x0 + (1.0 - self._maskh) * x1
        x_vcat = self._maskv * x0 + (1.0 - self._maskv) * x1
        mix_x = self._lamx.reshape(
            (-1, 1, 1, 1)) * x_hcat + (1.0-self._lamx.reshape((-1, 1, 1, 1))) * x_vcat
        mix_y = self._lamy * y0 + (1.0-self._lamy) * y1
        return mix_x, mix_y

    def loss(self, pred, mix_y):
        # Calculate a classification loss given mix_y and prediction results of mix_x.
        # Both pred and mix_y are supposed to be a nn.Variable
        return softmax_cross_entropy_loss_vlabel(pred, mix_y)


# data_augmentation/args.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


def get_args(save_path='tmp.monitor', max_iter=80000, learning_rate=1e-1, batch_size=128):
    """
    Get command line arguments.

    Arguments set the default values of command line arguments.
    """
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--batch-size", "-b", type=int, default=batch_size)
    parser.add_argument("--learning-rate", "-l",
                        type=float, default=learning_rate)
    parser.add_argument("--max-iter", "-i", type=int, default=max_iter)
    parser.add_argument("--val-interval", "-v", type=int, default=1000)
    parser.add_argument("--device-id", "-g", type=str, default='0',
                        help='Device ID the training run on. This is only valid if you specify `-c cudnn`.')
    parser.add_argument("--type-config", "-t", type=str, default='float',
                        help='Type of computation. e.g. "float", "half".')
    parser.add_argument("--model-save-interval", "-s", type=int, default=5000)
    parser.add_argument("--save-path", "-p",
                        type=str, default=save_path)
    parser.add_argument('--context', '-c', type=str,
                        default='cudnn', help="Extension path. ex) cpu, cudnn.")
    parser.add_argument("--net", "-n", type=str,
                        default='resnet18',
                        help="Neural network architecture type.")
    parser.add_argument("--solver", "-o", type=str,
                        default='Adam',
                        help="Solver for training DNN")

    # Mix-up
    parser.add_argument("--alpha", "-a", type=float, default=0.5)
    parser.add_argument("--mixtype", "-mt", type=str, default="mixup")

    return parser.parse_args()


# data_augmentation/classification_MDL.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import os
import time
from args import get_args
from data_cifar10 import *
import nnabla as nn
from nnabla.ext_utils import get_extension_context
import nnabla.functions as F
import nnabla.parametric_functions as PF
import nnabla.solvers as S
import numpy as np
import functools
from models import *
from MixedDataLearning import *


def single_image_augment(image):
    image = F.image_augmentation(image, contrast=1.0, angle=0.0, flip_lr=True)
    image = F.random_shift(image, shifts=(4, 4), border_mode="reflect")
    return image


def train():
    """
    Main script.

    Steps:

    * Parse command line arguments.
    * Specify contexts for computation.
    * Initialize DataIterator.
    * Construct a computation graph for training and one for validation.
    * Initialize solver and set parameter variables to that.
    * Create monitor instances for saving and displaying training stats.
    * Training loop
      * Computate error rate for validation data (periodically)
      * Get a next minibatch.
      * Execute forwardprop
      * Set parameter gradients zero
      * Execute backprop.
      * Solver updates parameters by using gradients computed by backprop.
      * Compute training error
    """
    # Parse args
    args = get_args()
    n_valid_samples = 10000
    bs_valid = args.batch_size
    extension_module = args.context
    ctx = get_extension_context(
        extension_module, device_id=args.device_id, type_config=args.type_config)
    nn.set_default_context(ctx)

    # Dataset
    data_iterator = data_iterator_cifar10
    n_class = 10

    # Model architecture
    if args.net == "resnet18":
        prediction = functools.partial(
            resnet18_prediction, ncls=n_class, nmaps=64, act=F.relu)
    if args.net == "resnet34":
        prediction = functools.partial(
            resnet34_prediction, ncls=n_class, nmaps=64, act=F.relu)

    # Create training graphs
    test = False
    if args.mixtype == "mixup":
        mdl = MixupLearning(args.batch_size, alpha=args.alpha)
    elif args.mixtype == "cutmix":
        mdl = CutmixLearning((args.batch_size, 3, 32, 32),
                             alpha=args.alpha, cutmix_prob=1.0)
    elif args.mixtype == "vhmixup":
        mdl = VHMixupLearning((args.batch_size, 3, 32, 32), alpha=args.alpha)
    else:
        print("[ERROR] Unknown mixtype: " + args.mixtype)
        return
    image_train = nn.Variable((args.batch_size, 3, 32, 32))
    label_train = nn.Variable((args.batch_size, 1))
    mix_image, mix_label = mdl.mix_data(single_image_augment(
        image_train), F.one_hot(label_train, (n_class, )))
    pred_train = prediction(mix_image, test)
    loss_train = mdl.loss(pred_train, mix_label)
    input_train = {"image": image_train, "label": label_train}

    # Create validation graph
    test = True
    image_valid = nn.Variable((bs_valid, 3, 32, 32))
    pred_valid = prediction(image_valid, test)
    input_valid = {"image": image_valid}

    # Solvers
    if args.solver == "Adam":
        solver = S.Adam()
    elif args.solver == "Momentum":
        solver = S.Momentum(lr=args.learning_rate)
    solver.set_parameters(nn.get_parameters())

    # Create monitor
    from nnabla.monitor import Monitor, MonitorSeries, MonitorTimeElapsed
    monitor = Monitor(args.save_path)
    monitor_loss = MonitorSeries("Training loss", monitor, interval=10)
    monitor_time = MonitorTimeElapsed("Training time", monitor, interval=10)
    monitor_verr = MonitorSeries("Test error", monitor, interval=1)

    # Data Iterator
    tdata = data_iterator(args.batch_size, True)
    vdata = data_iterator(args.batch_size, False)

    print("Size of the training data: %d " % tdata.size)
    # Training-loop
    for i in range(args.max_iter):
        # Forward/Zerograd/Backward
        image, label = tdata.next()
        input_train["image"].d = image
        input_train["label"].d = label
        mdl.set_mix_ratio()
        loss_train.forward()
        solver.zero_grad()
        loss_train.backward()

        # Model update by solver
        if args.solver == "Momentum":
            if i == args.max_iter / 2:
                solver.set_learning_rate(args.learning_rate / 10.0)
            if i == args.max_iter / 4 * 3:
                solver.set_learning_rate(args.learning_rate / 10.0**2)
        solver.update()

        # Validation
        if (i+1) % args.val_interval == 0 or i == 0:
            ve = 0.
            vdata._reset()
            vdata_pred = np.zeros((n_valid_samples, n_class))
            vdata_label = np.zeros((n_valid_samples, 1), dtype=np.int32)
            for j in range(0, n_valid_samples, args.batch_size):
                image, label = vdata.next()
                input_valid["image"].d = image
                pred_valid.forward()
                vdata_pred[j:min(j+args.batch_size, n_valid_samples)
                           ] = pred_valid.d[:min(args.batch_size, n_valid_samples-j)]
                vdata_label[j:min(j+args.batch_size, n_valid_samples)
                            ] = label[:min(args.batch_size, n_valid_samples-j)]
            ve = categorical_error(vdata_pred, vdata_label)
            monitor_verr.add(i+1, ve)

        if int((i+1) % args.model_save_interval) == 0:
            nn.save_parameters(os.path.join(
                args.save_path, 'params_%06d.h5' % (i+1)))

        # Monitering
        monitor_loss.add(i+1, loss_train.d.copy())
        monitor_time.add(i+1)

    nn.save_parameters(os.path.join(args.save_path,
                                    'params_%06d.h5' % (args.max_iter)))


if __name__ == '__main__':
    """
    $ python classification.py --context "cudnn" -b 64

    """
    train()


# data_augmentation/data_cifar10.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''
Provide data iterator for CIFAR10 examples.
'''
from contextlib import contextmanager
import numpy as np
import struct
import tarfile
import zlib
import time
import os
import errno

from nnabla.logger import logger
from nnabla.utils.data_iterator import data_iterator
from nnabla.utils.data_source import DataSource
from nnabla.utils.data_source_loader import download, get_data_home


class Cifar10DataSource(DataSource):
    '''
    Get data directly from cifar10 dataset from Internet(yann.lecun.com).
    '''

    def _get_data(self, position):
        image = self._images[self._indexes[position]]
        label = self._labels[self._indexes[position]]
        return (image, label)

    def __init__(self, train=True, shuffle=False, rng=None):
        super(Cifar10DataSource, self).__init__(shuffle=shuffle, rng=rng)

        self._train = train
        data_uri = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"
        logger.info('Getting labeled data from {}.'.format(data_uri))
        r = download(data_uri)  # file object returned
        with tarfile.open(fileobj=r, mode="r:gz") as fpin:
            # Training data
            if train:
                images = []
                labels = []
                for member in fpin.getmembers():
                    if "data_batch" not in member.name:
                        continue
                    fp = fpin.extractfile(member)
                    data = np.load(fp, encoding="bytes", allow_pickle=True)
                    images.append(data[b"data"])
                    labels.append(data[b"labels"])
                self._size = 50000
                self._images = np.concatenate(
                    images).reshape(self._size, 3, 32, 32)
                self._labels = np.concatenate(labels).reshape(-1, 1)
            # Validation data
            else:
                for member in fpin.getmembers():
                    if "test_batch" not in member.name:
                        continue
                    fp = fpin.extractfile(member)
                    data = np.load(fp, encoding="bytes", allow_pickle=True)
                    images = data[b"data"]
                    labels = data[b"labels"]
                self._size = 10000
                self._images = images.reshape(self._size, 3, 32, 32)
                self._labels = np.array(labels).reshape(-1, 1)
        r.close()
        logger.info('Getting labeled data from {}.'.format(data_uri))

        self._size = self._labels.size
        self._variables = ('x', 'y')
        if rng is None:
            rng = np.random.RandomState(313)
        self.rng = rng
        self.reset()

    def reset(self):
        if self._shuffle:
            self._indexes = self.rng.permutation(self._size)
        else:
            self._indexes = np.arange(self._size)
        super(Cifar10DataSource, self).reset()

    @property
    def images(self):
        """Get copy of whole data with a shape of (N, 1, H, W)."""
        return self._images.copy()

    @property
    def labels(self):
        """Get copy of whole label with a shape of (N, 1)."""
        return self._labels.copy()


def data_iterator_cifar10(batch_size,
                          train=True,
                          rng=None,
                          shuffle=True,
                          with_memory_cache=False,
                          with_file_cache=False):
    '''
    Provide DataIterator with :py:class:`Cifar10DataSource`
    with_memory_cache and with_file_cache option's default value is all False,
    because :py:class:`Cifar10DataSource` is able to store all data into memory.
    '''
    return data_iterator(Cifar10DataSource(train=train, shuffle=shuffle, rng=rng),
                         batch_size,
                         rng,
                         with_memory_cache,
                         with_file_cache)


# data_augmentation/models.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import os
import time
from args import get_args
import nnabla as nn
import nnabla.communicators as C
from nnabla.ext_utils import get_extension_context
import nnabla.functions as F
import nnabla.parametric_functions as PF
import nnabla.solvers as S
import numpy as np


def categorical_error(pred, label):
    """
    Compute categorical error given score vectors and labels as
    numpy.ndarray.
    """
    pred_label = pred.argmax(1)
    return (pred_label != label.flat).mean()


def resnet18_prediction(image, test=False, ncls=10, nmaps=64, act=F.relu):
    """
    Construct ResNet 18
    """
    # Residual Unit
    def res_unit(x, nmap_out, scope_name, stride=1):
        nmap_in = x.shape[1]
        with nn.parameter_scope(scope_name):
            # Conv -> BN -> Nonlinear
            with nn.parameter_scope("conv1"):
                h = PF.convolution(x, nmap_out, kernel=(3, 3), pad=(1, 1),
                                   with_bias=False, stride=(stride, stride))
                h = PF.batch_normalization(h, batch_stat=not test)
                h = act(h)
            # Conv -> BN -> Nonlinear
            with nn.parameter_scope("conv2"):
                h = PF.convolution(h, nmap_out, kernel=(3, 3), pad=(1, 1),
                                   with_bias=False)
                h = PF.batch_normalization(h, batch_stat=not test)
            # Conv -> BN
            if nmap_in != nmap_out:
                with nn.parameter_scope("conv3"):
                    x2 = PF.convolution(x, nmap_out, kernel=(1, 1), pad=(0, 0),
                                        with_bias=False, stride=(stride, stride))
                    x2 = PF.batch_normalization(x2, batch_stat=not test)
            else:
                x2 = x
            # Residual -> Nonlinear
            h = act(F.add2(h, x2))
            return h
    # Conv -> BN -> Nonlinear
    with nn.parameter_scope("conv1"):
        h = PF.convolution(image, nmaps, kernel=(3, 3),
                           pad=(1, 1), with_bias=False)
        h = PF.batch_normalization(h, batch_stat=not test)
        h = act(h)

    h = res_unit(h, nmaps, "conv2-1", 1)    # -> 32x32
    h = res_unit(h, nmaps, "conv2-2", 1)    # -> 32x32
    h = res_unit(h, nmaps*2, "conv3-1", 2)  # -> 16x16
    h = res_unit(h, nmaps*2, "conv3-2", 1)  # -> 16x16
    h = res_unit(h, nmaps*4, "conv4-1", 2)  # -> 8x8
    h = res_unit(h, nmaps*4, "conv4-2", 1)  # -> 8x8
    h = res_unit(h, nmaps*8, "conv5-1", 2)  # -> 4x4
    h = res_unit(h, nmaps*8, "conv5-2", 1)  # -> 4x4
    h = F.average_pooling(h, kernel=(4, 4))  # -> 1x1
    h = PF.affine(h, 1000, name="bottleneck")  # -> 1x1000
    h = act(h)
    pred = PF.affine(h, ncls)

    return pred


def resnet34_prediction(image, test=False, ncls=10, nmaps=64, act=F.relu):
    """
    Construct ResNet 34
    """
    # Residual Unit
    def res_unit(x, nmap_out, scope_name, stride=1):
        nmap_in = x.shape[1]
        with nn.parameter_scope(scope_name):
            # Conv -> BN -> Nonlinear
            with nn.parameter_scope("conv1"):
                h = PF.convolution(x, nmap_out, kernel=(3, 3), pad=(1, 1),
                                   with_bias=False, stride=(stride, stride))
                h = PF.batch_normalization(h, batch_stat=not test)
                h = act(h)
            # Conv -> BN -> Nonlinear
            with nn.parameter_scope("conv2"):
                h = PF.convolution(h, nmap_out, kernel=(3, 3), pad=(1, 1),
                                   with_bias=False)
                h = PF.batch_normalization(h, batch_stat=not test)
            # Conv -> BN
            if nmap_in != nmap_out:
                with nn.parameter_scope("conv3"):
                    x2 = PF.convolution(x, nmap_out, kernel=(1, 1), pad=(0, 0),
                                        with_bias=False, stride=(stride, stride))
                    x2 = PF.batch_normalization(x2, batch_stat=not test)
            else:
                x2 = x
            # Residual -> Nonlinear
            h = act(F.add2(h, x2))
            return h
    # Conv -> BN -> Nonlinear
    with nn.parameter_scope("conv1"):
        h = PF.convolution(image, nmaps, kernel=(3, 3),
                           pad=(1, 1), with_bias=False)
        h = PF.batch_normalization(h, batch_stat=not test)
        h = act(h)

    h = res_unit(h, nmaps, "conv2-1", 1)    # -> 32x32
    h = res_unit(h, nmaps, "conv2-2", 1)    # -> 32x32
    h = res_unit(h, nmaps, "conv2-3", 1)    # -> 32x32
    h = res_unit(h, nmaps*2, "conv3-1", 2)  # -> 16x16
    h = res_unit(h, nmaps*2, "conv3-2", 1)  # -> 16x16
    h = res_unit(h, nmaps*2, "conv3-3", 1)  # -> 16x16
    h = res_unit(h, nmaps*2, "conv3-4", 1)  # -> 16x16
    h = res_unit(h, nmaps*4, "conv4-1", 2)  # -> 8x8
    h = res_unit(h, nmaps*4, "conv4-2", 1)  # -> 8x8
    h = res_unit(h, nmaps*4, "conv4-3", 1)  # -> 8x8
    h = res_unit(h, nmaps*4, "conv4-4", 1)  # -> 8x8
    h = res_unit(h, nmaps*4, "conv4-5", 1)  # -> 8x8
    h = res_unit(h, nmaps*4, "conv4-6", 1)  # -> 8x8
    h = res_unit(h, nmaps*8, "conv5-1", 2)  # -> 4x4
    h = res_unit(h, nmaps*8, "conv5-2", 1)  # -> 4x4
    h = res_unit(h, nmaps*8, "conv5-3", 1)  # -> 4x4
    h = F.average_pooling(h, kernel=(4, 4))  # -> 1x1
    h = PF.affine(h, 1000, name="bottleneck")  # -> 1x1000
    h = act(h)
    pred = PF.affine(h, ncls)

    return pred


# diffusion-models/generate.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import os
from queue import Queue
from threading import Thread

import hydra
import nnabla as nn
import numpy as np
from neu.misc import init_nnabla
from neu.reporter import get_tiled_image, save_tiled_image
from nnabla.logger import logger
from nnabla.utils.image_utils import imsave
from omegaconf import OmegaConf

import config
from dataset.common import SimpleDataIterator
from diffusion_model.model import Model


def save_img(q: Queue, channel_last: bool):
    while True:
        mode, img, file_path = q.get()

        if mode == "STOP":
            break

        assert isinstance(img, np.ndarray)

        if mode == "TILED":
            assert len(img.shape) == 4
            save_tiled_image(img.astype(np.uint8),
                             file_path,
                             channel_last=channel_last)

        else:
            if len(img.shape) == 4:
                assert img.shape[0] == 1
                img = img[0]

            imsave(file_path, img.astype(np.uint8),
                   channel_first=not channel_last)


@hydra.main(version_base=None,
            config_path="config/yaml",
            config_name="config_generate")
def main(conf: config.GenScriptConfig):
    # load diffusion and model config
    loaded_conf: config.LoadedConfig = config.load_saved_conf(
        conf.generate.config)

    comm = init_nnabla(ext_name="cudnn",
                       device_id=conf.runtime.device_id,
                       type_config="float",
                       random_pseed=True)

    # update respacing parameter based on given config
    loaded_conf.diffusion.respacing_step = conf.generate.respacing_step

    model = Model(diffusion_conf=loaded_conf.diffusion,
                  model_conf=loaded_conf.model)

    # load parameters
    assert os.path.exists(
        conf.generate.h5), f"{conf.generate.h5} is not found. Please make sure the h5 file exists."
    nn.parameter.load_parameters(conf.generate.h5)

    # Generate

    # setup data iterator for lowres samples
    B = conf.generate.batch_size

    model_kwargs = {}
    is_upsample = loaded_conf.model.low_res_size is not None
    if is_upsample:
        lowres_conf = config.DatasetConfig(batch_size=B,
                                           dataset_root_dir=conf.generate.base_samples_dir,
                                           image_size=loaded_conf.model.low_res_size,
                                           shuffle_dataset=False,
                                           channel_last=loaded_conf.model.channel_last)

        def lowres_label_cb(path):
            # assume file name is "gen_{class_id}_{image id}_{device_id}_{...}.png
            return int(os.path.basename(path).split(".")[0].split("_")[1])

        data_lowres = SimpleDataIterator(lowres_conf,
                                         comm=comm,
                                         label_creator_callback=lowres_label_cb)

        # calculate number of iterations based on number of lowres samples.
        num_iter = (data_lowres.size + B - 1) // B

    else:
        num_samples_per_iter = B * comm.n_procs
        num_iter = (conf.generate.samples +
                    num_samples_per_iter - 1) // num_samples_per_iter

    if loaded_conf.model.class_cond:
        if is_upsample:
            model_kwargs["class_label"] = nn.Variable(shape=(B, ))
        elif conf.generate.gen_class_id is None:
            # random class
            import nnabla.functions as F
            model_kwargs["class_label"] = F.randint(
                low=0, high=loaded_conf.model.num_classes, shape=(B, ))
        else:
            assert conf.generate.gen_class_id in list(range(0, loaded_conf.model.num_classes)), \
                f"invalid class_id: '{conf.generate_gen_class_id}'. Must be an integer in [0, {loaded_conf.model.num_classes})."
            # deterministic class
            model_kwargs["class_label"] = nn.Variable.from_numpy_array(
                [conf.generate.gen_class_id for _ in range(B)])

        model_kwargs["class_label"].persistent = True

    # sampling
    local_saved_cnt = 0

    # setup output dir and show config
    if comm.rank == 0:
        os.makedirs(conf.generate.output_dir, exist_ok=True)

        logger.info("===== script config =====")
        print(OmegaConf.to_yaml(conf))

        logger.info("===== loaded config =====")
        print(OmegaConf.to_yaml(loaded_conf))

    comm.barrier()

    # launch thread to save img
    q = Queue()
    thread = Thread(target=save_img, args=(q, loaded_conf.model.channel_last))
    thread.start()

    for i in range(num_iter):
        logger.info(f"Generate samples {i + 1} / {num_iter}.")

        if is_upsample:
            lowres, label = data_lowres.next()
            lowres_var = nn.Variable.from_numpy_array(lowres)

            if loaded_conf.model.noisy_low_res:
                if conf.generate.lowres_aug_timestep is None:
                    aug_timestep = nn.Variable.from_numpy_array(
                        np.zeros(shape=(B, )))
                    model_kwargs["input_cond_aug_timestep"] = aug_timestep
                else:
                    aug_timestep = nn.Variable.from_numpy_array(
                        [conf.generate.lowres_noise_timestep for _ in range(B)])
                    lowres_var, aug_level = model.gaussian_conditioning_augmentation(
                        lowres_var, aug_timestep)
                    model_kwargs["input_cond_aug_timestep"] = aug_level

            model_kwargs["input_cond"] = lowres_var

            if loaded_conf.model.class_cond:
                model_kwargs["class_label"].d = label

        sample_out, xt_samples, x_starts = model.sample(shape=[B, ] + loaded_conf.model.image_shape,
                                                        x_init=None,
                                                        dump_interval=-1,
                                                        use_ema=conf.generate.ema,
                                                        progress=comm.rank == 0,
                                                        sampler=conf.generate.sampler,
                                                        classifier_free_guidance_weight=conf.generate.classifier_free_guidance_weight,
                                                        model_kwargs=model_kwargs)

        # scale back to [0, 255]
        sample_out = (sample_out.d + 1) * 127.5
        sample_out = np.clip(sample_out, 0, 255)

        if conf.generate.tiled:
            if conf.generate.gen_class_id is not None:
                save_path = os.path.join(
                    conf.generate.output_dir, f"class_{conf.generate.gen_class_id}_CFG_{conf.generate.classifier_free_guidance_weight}.png")
            else:
                save_path = os.path.join(
                    conf.generate.output_dir, f"gen_{local_saved_cnt}_{comm.rank}.png")
            q.put(["TILED", sample_out, save_path])
            local_saved_cnt += 1
        else:
            for b in range(B):
                if loaded_conf.model.class_cond:
                    class_id = int(model_kwargs["class_label"].d[b])
                    save_path = os.path.join(
                        conf.generate.output_dir, f"gen_{class_id}_{local_saved_cnt}_{comm.rank}.png")
                else:
                    save_path = os.path.join(
                        conf.generate.output_dir, f"gen_{local_saved_cnt}_{comm.rank}.png")

                q.put(["SINGLE", sample_out[b], save_path])

                local_saved_cnt += 1

        # create video for x_starts
        if conf.generate.save_xstart:
            import moviepy.editor as mp

            clips = []
            for i in range(len(x_starts)):
                xstart = x_starts[i][1]
                assert isinstance(xstart, np.ndarray)
                im = get_tiled_image(
                    np.clip((xstart + 1) * 127.5, 0, 255), channel_last=conf.channel_last).astype(np.uint8)
                clips.append(im)

            clip = mp.ImageSequenceClip(clips, fps=5)
            clip.write_videofile(os.path.join(
                conf.generate.output_dir, f"pred_x0_along_time_{local_saved_cnt}_{comm.rank}.mp4"))

    q.put(["STOP", None, None])
    thread.join()


if __name__ == "__main__":
    main()


# diffusion-models/launch_inference_server.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import hydra

from config import InferenceServerScriptConfig
import ui

from neu.misc import init_nnabla


@hydra.main(version_base=None,
            config_path="config/yaml",
            config_name="config_inference_server")
def main(conf: InferenceServerScriptConfig):
    # initialize nnabla context
    init_nnabla(ext_name="cudnn",
                device_id=conf.runtime.device_id,
                type_config="float")

    # create demo
    demo = ui.create_demo(conf.server)

    # launch server
    demo.queue(concurrency_count=1)
    demo.launch(server_name="0.0.0.0",
                server_port=conf.server.port)


if __name__ == "__main__":
    main()


# diffusion-models/sdedit.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import os

import hydra
import nnabla as nn
from nnabla.logger import logger
import numpy as np
from omegaconf import OmegaConf
from neu.misc import init_nnabla
from neu.reporter import get_tiled_image, save_tiled_image
from nnabla.utils.image_utils import imsave, imread

import config
from diffusion_model.model import Model


@hydra.main(version_base=None,
            config_path="config/yaml",
            config_name="config_generate")
def main(conf: config.GenScriptConfig):
    assert conf.generate.x_start is not None, "generate.x_start must be supecified for sdedit.py."
    assert os.path.exists(
        conf.generate.x_start), f"{conf.generate.x_start} is not found."

    # load diffusion and model config
    loaded_conf: config.LoadedConfig = config.load_saved_conf(
        conf.generate.config)

    comm = init_nnabla(ext_name="cudnn",
                       device_id=conf.runtime.device_id,
                       type_config="float",
                       random_pseed=True)

    if comm.n_procs > 1:
        raise ValueError("Currentely sdedit.py doesn't support multiGPU.")

    # update t_start if provided
    if conf.generate.t_start is not None:
        assert conf.generate.t_start < loaded_conf.diffusion.max_timesteps, \
            "t_start must be smaller than max_timesteps." \
            f"(t_start: {conf.generate.t_start}, T: {loaded_conf.diffusion.max_timesteps})"
        loaded_conf.diffusion.t_start = conf.generate.t_start

    model = Model(diffusion_conf=loaded_conf.diffusion,
                  model_conf=loaded_conf.model)

    # load parameters
    assert os.path.exists(
        conf.generate.h5), f"{conf.generate.h5} is not found. Please make sure the h5 file exists."
    nn.parameter.load_parameters(conf.generate.h5)

    # load image
    # todo: support multiple images
    if loaded_conf.model.channel_last:
        h, w = loaded_conf.model.image_shape[:-1]
    else:
        h, w = loaded_conf.model.image_shape[1:]
    x0 = imread(conf.generate.x_start, size=(w, h),
                channel_first=not loaded_conf.model.channel_last)
    x0 = x0[np.newaxis]

    # setup output dir
    os.makedirs(conf.generate.output_dir, exist_ok=True)

    # dump configs
    logger.info("===== script config =====")
    print(OmegaConf.to_yaml(conf))

    logger.info("===== loaded config =====")
    print(OmegaConf.to_yaml(loaded_conf))

    # Generate
    # todo: apply SDEdit several times
    x0_var = nn.Variable.from_numpy_array(x0 / 127.5 - 1)
    T_var = nn.Variable(shape=(1, ))
    T_var.data.fill(model.diffusion.num_timesteps - 1)
    sample_out, xts, x_starts = model.sample(shape=[1, ] + loaded_conf.model.image_shape,
                                             x_init=model.diffusion.q_sample(
                                                 x0_var, T_var),
                                             dump_interval=1,
                                             use_ema=conf.generate.ema,
                                             progress=comm.rank == 0,
                                             sampler=conf.generate.sampler)

    # scale back to [0, 255]
    sample_out = (sample_out + 1) * 127.5
    save_path = os.path.join(
        conf.generate.output_dir, f"gen.png")
    imsave(save_path, sample_out[0].astype(
        np.uint8), channel_first=not loaded_conf.model.channel_last)

    # create video for x_starts
    if conf.generate.save_xstart:
        import moviepy.editor as mp
        clips = []
        for i in range(len(x_starts)):
            xstart = x_starts[i][1]
            assert isinstance(xstart, np.ndarray)
            im = get_tiled_image(
                np.clip((xstart + 1) * 127.5, 0, 255), channel_last=False).astype(np.uint8)
            clips.append(im)

        clip = mp.ImageSequenceClip(clips, fps=5)
        clip.write_videofile(os.path.join(
            conf.generate.output_dir, f"pred_x0_along_time.mp4"))


if __name__ == "__main__":
    main()


# diffusion-models/train.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import fnmatch
import os
import queue
import sys

import hydra
import nnabla as nn
import nnabla.solvers as S
import numpy as np
from neu.checkpoint_util import load_checkpoint, save_checkpoint
from neu.misc import get_current_time, init_nnabla
from neu.mixed_precision import MixedPrecisionManager
from neu.reporter import KVReporter, save_tiled_image
from neu.solvers import PackedParameterSolver
from nnabla.logger import logger
from omegaconf import OmegaConf

import config
from dataset import get_dataset
from diffusion_model.diffusion import is_learn_sigma
from diffusion_model.layers import adaptive_pooling_2d
from diffusion_model.model import Model
from diffusion_model.utils import (create_ema_op,
                                   get_lr_scheduler,
                                   init_checkpoint_queue,
                                   sum_grad_norm)


def refine_monitor_files(dir_path, start_iter):
    for filename in os.listdir(dir_path):
        if filename.endswith(("series.txt", "timer.txt")):
            contents = []
            with open(os.path.join(dir_path, filename), "r") as f:
                for line in f.readlines():
                    iter = line.strip().split(" ")[0]
                    if int(iter) >= start_iter:
                        break

                    contents.append(line)

            with open(os.path.join(dir_path, filename), "w") as f:
                f.write("".join(contents))


def setup_resume(output_dir, dataset, solvers, is_master=False):
    # return
    start_iter = 0
    new_output_dir = output_dir

    parent = os.path.dirname(os.path.abspath(output_dir))

    if not os.path.exists(parent):
        # no previous checkpoint, train from scratch.
        return start_iter, output_dir

    all_logs = sorted(fnmatch.filter(
        os.listdir(parent), "*{}*".format(dataset)))
    if len(all_logs):
        latest_dir = os.path.join(parent, all_logs[-1])
        checkpoints = sorted(fnmatch.filter(
            os.listdir(latest_dir), "checkpoint_*.json"))

        # extract iteration count and use it as key
        iter_cp = [(int(x.split(".")[0].split("_")[-1]), x)
                   for x in checkpoints]

        for iter, checkpoint in reversed(sorted(iter_cp)):
            try:
                cp_path = os.path.join(latest_dir, checkpoint)
                start_iter = load_checkpoint(cp_path, solvers)

                for sname, slv in solvers.items():
                    slv.zero_grad()

                if is_master:
                    refine_monitor_files(latest_dir, start_iter)
                    init_checkpoint_queue(latest_dir)

                new_output_dir = latest_dir
                logger.info(f"Load checkpoint from {cp_path}")

                break
            except:
                logger.warning(
                    f"{checkpoint} is broken. Try to load previous checkpoints.")
        else:
            logger.warning("No valid checkpoint. Train from scratch")

    return start_iter, new_output_dir


def get_output_dir_name(org, dataset):
    return os.path.join(org, f"{get_current_time()}_{dataset}")


def augmentation(x, channel_last, random_flip):
    import nnabla.functions as F
    aug = x

    if random_flip:
        aug = F.random_flip(aug, axes=[2, ] if channel_last else None)

    return aug


def create_gen_config(conf: config.TrainScriptConfig,
                      respacing_step: int) -> config.GenScriptConfig:
    conf_gen: config.GenScriptConfig \
         = OmegaConf.masked_copy(conf, ["diffusion", "model"])

    # setup respacing
    conf_gen.diffusion.respacing_step = respacing_step

    # disable dropout
    conf_gen.model.dropout = 0

    OmegaConf.set_readonly(conf_gen, True)

    return conf_gen


@hydra.main(version_base=None, config_path="config/yaml", config_name="config_train")
def main(conf: config.TrainScriptConfig):
    # setup output dir
    conf.train.output_dir = get_output_dir_name(conf.train.output_dir,
                                                conf.dataset.name)

    # initialize nnabla runtime and get communicator
    comm = init_nnabla(ext_name="cudnn",
                       device_id=conf.runtime.device_id,
                       type_config="float",  # mixed precision is handled in unet.py
                       random_pseed=True)

    # create data iterator
    data_iterator = get_dataset(conf.dataset, comm)

    # build graph
    model = Model(diffusion_conf=conf.diffusion,
                  model_conf=conf.model)

    # setup input image
    # assume data_iterator returns [-1, 1]
    x = nn.Variable((conf.train.batch_size, ) + conf.model.image_shape)
    x_aug = augmentation(x, conf.model.channel_last,
                         random_flip=conf.dataset.random_flip)

    # create low-resolution image
    model_kwargs = {}
    if conf.model.low_res_size is not None:
        assert len(conf.model.low_res_size) == 2
        x_low_res = adaptive_pooling_2d(x_aug,
                                        conf.model.low_res_size,
                                        mode="average",
                                        channel_last=conf.model.channel_last)

        # gaussian conditioning augmentation
        if conf.train.noisy_low_res:
            x_low_res, aug_level = model.gaussian_conditioning_augmentation(
                x_low_res)
            model_kwargs["input_cond_aug_timestep"] = aug_level

        # create model_kwargs
        model_kwargs["input_cond"] = x_low_res

    # set cond drop rate for classifier-free guidance
    model_kwargs["cond_drop_rate"] = conf.train.cond_drop_rate

    # setup class condition
    if conf.model.class_cond:
        model_kwargs["class_label"] = nn.Variable((conf.train.batch_size, ))

    # setup text condition
    if conf.model.text_cond:
        assert conf.model.text_emb_shape is not None
        model_kwargs["text_emb"] = nn.Variable(
            (conf.train.batch_size, ) + conf.model.text_emb_shape)

    loss_dict, t = model.build_train_graph(x_aug,
                                           loss_scaling=None if conf.train.loss_scaling == 1.0 else conf.train.loss_scaling,
                                           model_kwargs=model_kwargs)
    assert loss_dict.batched_loss.shape == (conf.train.batch_size, )
    assert t.shape == (conf.train.batch_size, )

    # optimizer
    solver = S.Adam()
    solver.set_parameters(nn.get_parameters())

    # Packing parameters to make solver faster
    solver: PackedParameterSolver \
        = PackedParameterSolver(solver, use_ema=True)

    # learning rate scheduler
    lr_scheduler = get_lr_scheduler(conf.train)

    # for ema update
    # Note: this should be defined after solver.set_parameters() to avoid update by solver.
    ema_op, ema_params = create_ema_op(nn.get_parameters(), 0.9999)
    dummy_solver_ema = S.Sgd()
    dummy_solver_ema.set_learning_rate(0)  # just in case
    dummy_solver_ema.set_parameters(ema_params)
    assert len(nn.get_parameters(grad_only=True)) == len(ema_params)
    assert len(nn.get_parameters(grad_only=False)) == 2 * len(ema_params)

    # for checkpoint
    solvers = {
        "main": solver,
        "ema": dummy_solver_ema,
    }

    start_iter = 0  # exclusive
    if conf.train.resume:
        # when resume, use the previous output_dir having the last checkpoint.
        start_iter, output_dir = setup_resume(conf.train.output_dir,
                                              conf.dataset.name,
                                              solvers, is_master=comm.rank == 0)
        conf.train.output_dir = output_dir

    # setup output directory for generated images during training
    image_dir = os.path.join(conf.train.output_dir, "image")
    if comm.rank == 0:
        os.makedirs(image_dir, exist_ok=True)

    comm.barrier()

    # Reporter
    reporter = KVReporter(comm, save_path=conf.train.output_dir,
                          skip_kv_to_monitor=False)
    # set all keys before to prevent synchronization error
    for i in range(4):
        reporter.set_key(f"loss_q{i}")
        if is_learn_sigma(conf.model.model_var_type):
            reporter.set_key(f"vlb_q{i}")

    if conf.train.progress:
        from tqdm import trange
        piter = trange(start_iter + 1, conf.train.n_iters + 1,
                       disable=comm.rank > 0, ncols=0)
    else:
        piter = range(start_iter + 1, conf.train.n_iters + 1)

    # freeze config to disable setting or updating values in conf.
    OmegaConf.resolve(conf)
    OmegaConf.set_readonly(conf, True)

    # setup model for generation
    conf_gen = create_gen_config(conf, respacing_step=4)
    gen_model = Model(diffusion_conf=conf_gen.diffusion,
                      model_conf=conf_gen.model)

    # dump config
    if comm.rank == 0:
        # show configs in stdout
        logger.info("===== configs =====")
        print(OmegaConf.to_yaml(conf))

        # save configs
        OmegaConf.save(conf, os.path.join(
            conf.train.output_dir, "config_train.yaml"))
        OmegaConf.save(conf_gen, os.path.join(
            conf.train.output_dir, "config_gen.yaml"))

    comm.barrier()

    # setup mixed precision training if needed
    mpm = MixedPrecisionManager(
        use_fp16=conf.model.use_mixed_precision,
        initial_log_loss_scale=10)

    # Queue to keep data instances for input_cond during sampling
    num_gen = 16
    data_queue = queue.Queue(maxsize=num_gen)

    for i in piter:
        # update learning rate
        if lr_scheduler is not None:
            cur_lr = lr_scheduler._get_lr(i, None)
        else:
            cur_lr = conf.train.lr

        # rescale lr to cancel backward accumulation
        solver.set_learning_rate(cur_lr / conf.train.accum)

        # evaluate graph
        dummy_solver_ema.zero_grad()  # just in case
        solver.zero_grad()

        retry_cnt = 0
        accum_cnt = 0
        while accum_cnt < conf.train.accum:
            batch = data_iterator.next()
            data = batch[0]
            x.d = data

            label = batch[1]  # if laion, this should be caption.

            text_emb = [None, ] * conf.train.batch_size
            if conf.model.text_cond:
                text_emb = batch[2]
                model_kwargs["text_emb"].d = text_emb

            if conf.model.class_cond:
                model_kwargs["class_label"].d = label

            # keep data for input_condition in generation step
            for data_idx in range(conf.train.batch_size):
                if data_queue.full():
                    break

                data_queue.put(
                    (data[data_idx], label[data_idx], text_emb[data_idx]))

            loss_dict.loss.forward(clear_no_need_grad=True)
            mpm.backward(loss_dict.loss,
                         clear_buffer=True)

            # Retry from the first accumulation step if overflow happens.
            if mpm.is_grad_overflow(solver):
                retry_cnt += 1

                # Raise if retry happens too many times.
                if retry_cnt == 100:
                    raise ValueError("Overflow happens too many times.")

                # mpm.backward resets grad of all params to zero in this case.
                accum_cnt = 0
                continue

            # fwd/bwd successes. Count up accum_cnt.
            accum_cnt += 1

            # reporting
            reporter.kv_mean("loss", loss_dict.loss)

            if is_learn_sigma(conf.model.model_var_type):
                reporter.kv_mean("vlb", loss_dict.vlb)

            # loss for each quantile
            for j in range(conf.train.batch_size):
                ti = t.d[j]
                q_level = int(ti) * 4 // conf.diffusion.max_timesteps
                assert q_level in (
                    0, 1, 2, 3), f"q_level should be one of [0, 1, 2, 3], but {q_level} is given."
                reporter.kv_mean(f"loss_q{q_level}", float(
                    loss_dict.batched_loss.d[j]))

                if is_learn_sigma(conf.model.model_var_type):
                    reporter.kv_mean(f"vlb_q{q_level}", loss_dict.vlb.d[j])

        # gradient accumulation
        assert accum_cnt == conf.train.accum

        # update
        mpm.scale_grad(solver)
        if isinstance(solver, PackedParameterSolver):
            comm.all_reduce([x.grad for x in solver.packed_params],
                            division=True, inplace=True)
        else:
            comm.all_reduce(
                [x.grad for x in solver.get_parameters().values()], division=True, inplace=True)
        mpm.update(solver, clip_grad=conf.train.clip_grad)

        # update ema params
        if isinstance(solver, PackedParameterSolver) and solver.use_ema:
            solver.updata_ema_params(decay=0.9999)
        else:
            ema_op.forward(clear_no_need_grad=True)

        # grad norm
        if conf.train.dump_grad_norm:
            gnorm = sum_grad_norm(solver.get_parameters().values())
            reporter.kv_mean("grad", gnorm)

        # samples
        reporter.kv("samples", i * conf.train.batch_size *
                    comm.n_procs * conf.train.accum)

        # iteration (only for no-progress)
        if not conf.train.progress:
            reporter.kv("iteration", i)

        if i % conf.train.show_interval == 0:
            if conf.train.progress:
                desc = reporter.desc(
                    reset=True, sync=True)
                piter.set_description(desc=desc)
            else:
                reporter.dump(file=sys.stdout if comm.rank ==
                              0 else None, reset=True, sync=False)  # True

            reporter.flush_monitor(i)

        if i % conf.train.save_interval == 0:
            if comm.rank == 0:
                save_checkpoint(conf.train.output_dir, i, solvers,
                                n_keeps=3,
                                split_h5_per_solver=True)

            comm.barrier()

        if conf.train.gen_interval > 0 and i % conf.train.gen_interval == 0:
            # sampling
            gen_model_kwargs = {}

            # setup condition vectors
            assert data_queue.qsize() == num_gen, \
                f"queue size is smaller than expected. ({data_queue.qsize()} < {num_gen})"
            image_list = []
            label_list = []
            text_emb_list = []
            for _ in range(num_gen):
                image, label, text_emb = data_queue.get()
                image_list.append(image)
                label_list.append(label)
                text_emb_list.append(text_emb)

            if conf.model.low_res_size is not None:
                # create lowres input from training dataset
                input_cond = nn.Variable.from_numpy_array(np.stack(image_list))

                input_cond_lowres = adaptive_pooling_2d(input_cond, conf.model.low_res_size,
                                                        mode="average",
                                                        channel_last=conf.model.channel_last)
                input_cond_lowres.forward(clear_buffer=True)

                gen_model_kwargs["input_cond"] = input_cond_lowres.get_unlinked_variable(
                    need_grad=False)
                gen_model_kwargs["input_cond_aug_timestep"] = nn.Variable.from_numpy_array(
                    np.zeros((num_gen, )))

            # disable dropping condition for generation
            gen_model_kwargs["cond_drop_rate"] = 0

            # class cond
            if conf.model.class_cond:
                if conf.model.low_res_size is not None:
                    # use class id for lowres image for the upsampler
                    gen_class_label = nn.Variable.from_numpy_array(
                        np.stack(label_list))
                else:
                    # use random class id for the base model
                    gen_class_label = nn.Variable.from_numpy_array(np.random.randint(low=0,
                                                                                     high=conf.model.num_classes,
                                                                                     size=(num_gen, )))

                gen_model_kwargs["class_label"] = gen_class_label

            if conf.model.text_cond:
                gen_model_kwargs["text_emb"] = nn.Variable.from_numpy_array(
                    np.asarray(text_emb_list))

                assert gen_model_kwargs["text_emb"].shape == (
                    num_gen, ) + model_kwargs["text_emb"].shape[1:]

            sample_out, _, _ = gen_model.sample(shape=(num_gen, ) + x.shape[1:],
                                                model_kwargs=gen_model_kwargs,
                                                use_ema=True, progress=False)
            assert sample_out.shape == (num_gen, ) + x.shape[1:]

            # scale back to [0, 255]
            sample_out = (sample_out + 1) * 127.5

            save_path = os.path.join(image_dir, f"gen_{i}_{comm.rank}.png")

            save_tiled_image(sample_out.astype(np.uint8),
                             save_path, channel_last=conf.model.channel_last)

            if conf.model.text_cond:
                with open(os.path.join(image_dir, f"gen_{i}_{comm.rank}_script.txt"), "w") as f:
                    f.write("\n".join(label_list))


if __name__ == "__main__":
    main()


# interactive-demos/colab_utils.py

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from IPython.display import Image, display, Javascript
from google.colab.output import eval_js
from base64 import b64decode


def take_photo(filename="photo.png", cam_width=None, cam_height=None):
    """
        take a photo using a camera. If the machine has multiple cameras,
        you need to switch the camera
        --- Usage ---
        from IPython.display import Image
        try:
            filename = take_photo(cam_width=256, cam_height=256) # better to pass the default resolutions
            print('Saved to {}'.format(filename))
            # Show the image which was just taken.
            display(Image(filename))
        except Exception as err:
            # Errors will be thrown if the user does not have a webcam or if they do not
            # grant the page permission to access it.
            print(str(err))
    """
    js = Javascript('''
    async function takePhoto(cam_width, cam_height) {
      // default config
      var CONSTRAINTS = {
        audio: false,
        video: {
          width: {ideal: cam_width},
          height: {ideal: cam_height},
          deviceId: null
        }
      };

      function syncCamera(video, dev_id){
        CONSTRAINTS.video.deviceId = dev_id;
        if(curSTREAM !== null){
          curSTREAM.getVideoTracks().forEach((camera) => {
            camera.stop();
          });
      }

      navigator.mediaDevices.getUserMedia(CONSTRAINTS)
        .then( (stream) => {
          curSTREAM = stream;
          video.srcObject = stream;
          video.onloadedmetadata = (e) => {
            video.play();
          };
        })
      .catch((err) => {
        //console.log(err.name + ": " + err.message);
        return;
        });
      }

      // current Stream
      var curSTREAM = null;

      const div = document.createElement('div');
      const capture = document.createElement('button');
      capture.textContent = 'Capture';
      div.appendChild(capture);
        if (!navigator.mediaDevices || !navigator.mediaDevices.enumerateDevices) {
        //console.log("enumerateDevices() not supported.");
        return;
      }

      navigator.mediaDevices.enumerateDevices()
      .then(function(devices) {
        devices.forEach(function(device) {
          if (device.kind.includes("video")) {
            const dev_button = document.createElement('button');
            dev_button.textContent = device.label;
            const dev_id = device.deviceId;
            div.appendChild(dev_button);
            dev_button.onclick = function() {
            syncCamera(video, dev_id);
            }
          };
        });
      })
      .catch(function(err) {
        //console.log(err.name + ": " + err.message);
        return;
      });

      const video = document.createElement('video');
      video.style.display = 'block';
      curSTREAM = await navigator.mediaDevices.getUserMedia(CONSTRAINTS);
      document.body.appendChild(div);
      div.appendChild(video);
      video.srcObject = curSTREAM;
      await video.play();

      // Resize the output to fit the video element.
      google.colab.output.setIframeHeight(document.documentElement.scrollHeight, true);

      // Wait for Capture to be clicked.
      await new Promise((resolve) => capture.onclick = resolve);

      const canvas = document.createElement('canvas');
      canvas.width = video.videoWidth;
      canvas.height = video.videoHeight;
      canvas.getContext('2d').drawImage(video, 0, 0);
      curSTREAM.getVideoTracks()[0].stop();
      div.remove();
      return canvas.toDataURL();
      }
    ''')

    if not cam_width:
        cam_width = "null"
    if not cam_height:
        cam_height = "null"

    display(js)
    data = eval_js(f'takePhoto({cam_width}, {cam_height})')
    binary = b64decode(data.split(',')[1])
    with open(filename, 'wb') as f:
        f.write(binary)
    return filename


# responsible_ai/adversarial_debiasing/utils.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pandas as pd
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
import matplotlib.pyplot as plt
from IPython.display import Markdown, display
import sys
import numpy as np


def load_adult_data():
    """
    Census Income Dataset; target is to predict whether income exceeds $50K/yr
    (For more info : https://archive.ics.uci.edu/ml/datasets/adult)

    Returns:
        features (pandas.DataFrame): features of the Census income dataset
        target (pandas.Series): target values of the Census income dataset
        sensitive_attributes(pandas.DataFrame): sensitive attributes values of Census income dataset
    """
    file_path = r'https://archive.ics.uci.edu/ml/machine-learning-databases/adult/adult.data'
    column_names = ['age', 'workclass', 'fnlwgt', 'education', 'education_num',
                    'marital_status', 'occupation', 'relationship', 'race', 'sex',
                    'capital_gain', 'capital_loss', 'hours_per_week', 'country', 'target']
    try:
        # Restrict races to White and Black
        input_data = (pd.read_csv(file_path, names=column_names,
                                  na_values="?", sep=r'\s*,\s*', engine='python').loc[lambda df: df['race'].isin(['White', 'Black'])]).dropna()
    except Exception as e:
        print("Error : ", e)
        sys.exit(1)

    # sensitive attributes; we identify 'race' and 'sex' as sensitive attributes
    # privileged class for race and sex is White and Male respectively
    sensitive_attribs = ['race', 'sex']
    sensitive_attributes = (input_data.loc[:, sensitive_attribs]
                            .assign(race=lambda df: (df['race'] == 'White').astype(int),
                                    sex=lambda df: (df['sex'] == 'Male').astype(int)))

    # targets; 1 when someone makes over 50k , otherwise 0
    target = (input_data['target'] == '>50K').astype(int)

    # features; note that the 'target' and sensitive attribute columns are dropped
    features = (input_data
                .drop(columns=['target', 'race', 'sex', 'fnlwgt'])
                .pipe(pd.get_dummies, drop_first=True))

    display(Markdown(
        f"features : {features.shape[0]} samples, {features.shape[1]} attributes"))
    display(Markdown(f"targets : {target.shape[0]} samples"))
    display(Markdown(
        f"sensitives attributes : {sensitive_attributes.shape[0]} samples, {sensitive_attributes.shape[1]} attributes"))
    return features, target, sensitive_attributes


def get_p_rule(y_pred, z_values, threshold=0.5):
    """
    Compute the p%-rules for the classifier
    Note : threshold our classifier at 0.5 to make its prediction it binary.
    (For more info : https://arxiv.org/pdf/1507.05259.pdf)

    Args:
        y_pred: data predicted (calculated, output) by your model
        z_values: protected attribute values
        threshold: threshold of classifier

    Returns:
        p% rule (float): p%
    """

    y_z_1 = y_pred[z_values ==
                   1] > threshold if threshold else y_pred[z_values == 1]
    y_z_0 = y_pred[z_values ==
                   0] > threshold if threshold else y_pred[z_values == 0]
    odds = y_z_1.mean() / y_z_0.mean()
    return np.min([odds, 1/odds]) * 100


def get_fairness(y, privileged_group, preds):
    """
    Compute multiple fairness metrics for the classifier:
    1. Demographic parity
    2. Equal opportunity
    3. Equalized odds

    Args:
        y : True data (or target, ground truth)
        privileged_group : list of privileged group values
        preds : data predicted (calculated, output) by your model

    Returns:
        demographic_parity_difference (float): Demographic parity
        equal_opportunity_difference (float): Equal opportunity
        average_abs_odds_difference(float): Equalized odds

    """

    y_unprivileged, preds_unprivileged = y[privileged_group ==
                                           False], preds[privileged_group == False]
    y_privileged, preds_privileged, = y[privileged_group], preds[privileged_group]
    cm_unprivileged = confusion_matrix(y_unprivileged, preds_unprivileged)
    cm_privileged = confusion_matrix(y_privileged, preds_privileged)
    unprivileged_PR = (
        cm_unprivileged[1, 1] + cm_unprivileged[0, 1]) / cm_unprivileged.sum()
    privileged_PR = (cm_privileged[1, 1] +
                     cm_privileged[0, 1]) / cm_privileged.sum()

    # compute demographic parity
    demographic_parity_difference = unprivileged_PR - privileged_PR
    unprivileged_TPR = cm_unprivileged[1, 1] / cm_unprivileged[1].sum()
    privileged_TPR = cm_privileged[1, 1] / cm_privileged[1].sum()

    # compute equal opportunity
    equal_opportunity_difference = unprivileged_TPR - privileged_TPR
    unprivileged_FPR = cm_unprivileged[0, 1] / cm_unprivileged[0].sum()
    privileged_FPR = cm_privileged[0, 1] / cm_privileged[0].sum()

    # compute Equalized odds
    average_abs_odds_difference = 0.5 * (
                abs(unprivileged_FPR - privileged_FPR) + abs(unprivileged_TPR - privileged_TPR))

    return demographic_parity_difference, equal_opportunity_difference, average_abs_odds_difference


def plot_fairness(fairness, ax, metric="DPD", title="fairness metric", bar_x_axis="Original"):
    """
    plot single fairness metric
    """
    ax.set_ylim([-0.6, 0.6])
    ax.axhline(y=0.0, color='r', linestyle='-')
    ax.bar([bar_x_axis], fairness, color="blue", width=2)
    ax.set_ylabel(metric)
    ax.set_title(title, fontsize=10)
    for index, value in enumerate(fairness):
        if value < 0:
            ax.text(index, value - 0.1,
                    str(round(value, 3)), fontweight='bold', color='red', bbox=dict(facecolor='red', alpha=0.4))
        else:
            ax.text(index, value + 0.1,
                    str(round(value, 3)), fontweight='bold', color='red', bbox=dict(facecolor='red', alpha=0.4))


def plot_fairness_multi(DPD, EOD, AAOD, p_rule, accuracy, bar_x_axis="original"):
    """
    plot three fairness metrics
    """
    fig, axes = plt.subplots(1, 3, figsize=(10, 4), sharey=True)
    fig.suptitle("Model Fairness", fontsize=16)
    plot_fairness([DPD], ax=axes[0], metric="DPD",
                  title="Demographic Parity", bar_x_axis=bar_x_axis)
    plot_fairness([EOD], ax=axes[1], metric="EOD",
                  title="Equal Opportunity", bar_x_axis=bar_x_axis)
    plot_fairness([AAOD], ax=axes[2], metric="AAOD",
                  title="Equalized Odds", bar_x_axis=bar_x_axis)
    fig.text(0.92, 0.8, '\n'.join(
        ["Satisfied p%-rules:", f"- p% : {p_rule:.3f}"]), fontsize='15')
    fig.text(0.92, 0.65, '\n'.join(
        ["Accuracy of the model:", f"- accuracy : {accuracy:.3f}"]), fontsize='15')
    plt.show()


# responsible_ai/adversarial_debiasing_images/adversarial.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import pickle
import numpy as np
import nnabla as nn
import nnabla.functions as F
import nnabla.parametric_functions as PF
import nnabla.solvers as S
from nnabla.models.imagenet import ResNet50
from sklearn.metrics import average_precision_score
from utils import utils


def clf_resnet50(layer, n_classes=1, train=True):
    """
    This function uses ResNet-50 pretrained on ImageNet as the base architecture
    and replaces the linear layer in ResNet with two linear layers with the
    hidden layer of size 2,048. Dropout and ReLU are applied
    between these. mainly this function use to classify the target network.
    """
    layer_1 = F.relu(PF.affine(layer, 2048, name='classifier_1'))
    if train:
        layer_1 = F.dropout(layer_1, 0.5)
    out = PF.affine(layer_1, n_classes, name='classifier_2')
    return layer_1, out


def domain_network(features, n_classes):
    """
    This function uses to classify the domain network.
    """
    with nn.parameter_scope('domainnetwork'):
        out2 = PF.affine(features, n_classes, name='domain')
    return out2


class adversarial:
    """
    This class is use to train the Adversarial model.

    This class uses ResNet-50 pretrained on ImageNet as the base architecture
    and replaces the linear layer in ResNet with two linear layers with the
    hidden layer of size 2,048. Dropout and ReLU are applied
    between these.
    The adversarial debiasing procedure takes inspiration from GANs (Goodfellow et al. 2014)
    for training a fair classifier. In GANs they introduce a system of two neural networks in which
    the two neural networks compete with each other to become more accurate in their predictions.
    Likewise, in adversarial debiasing, we build two classifier models
    1. Target Classifier
    2. Domain Classifier
    """

    def __init__(self, batch_size=32, learning_rate=1e-4, max_iter=5086,
                 total_epochs=50, training_ratio=3, lamda=5,
                 monitor_path=None, validation_weight=None, model_load_path=None):
        """
        Construct all the necessary attributes for the adversarial model.
        Args:
            batch_size (int): number of samples contained in each generated batch
            learning_rate (float) : learning rate
            max_iter (int) : maximum iterations for an epoch
            total_epochs (int) : total epochs to train the model
            training_ratio (int) : training ratio b/w the classifier and adversary n/w
            lamda (int) : lambda parameter weighs the adversarial loss of each class
            validation_weight : sample weights
            monitor_path (str) : model parameter to be saved
            model_load_path (str) : load the model
        """
        self.batch_size = batch_size
        self.total_epochs = total_epochs
        self.max_iter = max_iter
        self.training_ratio = training_ratio
        self.lamda = lamda
        self.acc = 0
        self.monitor_path = monitor_path
        self.validation_weight = validation_weight

        model = ResNet50()
        self.input_image = nn.Variable((self.batch_size,) + (3, 224, 224))
        self.label = nn.Variable([self.batch_size, 1])
        # fine tuning
        pool = model(self.input_image, training=True, use_up_to='pool')
        self.feature, self.clf = clf_resnet50(pool, n_classes=1)
        self.clf.persistent = True
        # loss
        self.loss = F.mean(F.sigmoid_cross_entropy(self.clf, self.label))
        self.loss.persistent = True
        # hyper parameters
        self.solver = S.Adam(learning_rate)
        self.solver.set_parameters(nn.get_parameters())

        self.domain_label = nn.Variable([self.batch_size, 1])
        self.domain_output = domain_network(self.feature, 2)
        self.domain_output.persistent = True
        self.domain_loss = F.mean(F.softmax_cross_entropy(
            self.domain_output, self.domain_label))
        self.domain_loss.persistent = True
        self.log_softmax = F.log_softmax(self.domain_output, axis=1)
        self.confusion_loss = -F.mean(self.log_softmax)
        self.confusion_loss.persistent = True

        self.adv_loss = self.loss + self.lamda * self.confusion_loss
        self.adv_loss.persistent = True
        self.domain_solver = S.Adam(learning_rate)

        with nn.parameter_scope("domainnetwork"):
            self.domain_solver.set_parameters(nn.get_parameters())

        self.x_v = nn.Variable((self.batch_size,) + (3, 224, 224))
        pool_v = model(self.x_v, training=False, use_up_to='pool')
        self.feature_v, self.v_clf = clf_resnet50(pool_v, train=False)
        self.v_clf_out = F.sigmoid(self.v_clf)

        if model_load_path is not None:
            _ = nn.load_parameters(model_load_path)

    def train(self, train_loader, val_loader):
        """
        train the adversarial model
        Args:
             train_loader (NNabla data iterator) : specify training dataset
             val_loader (NNabla data iterator) : specify validation dataset
        """

        model_save = True
        print_freq = 100
        for epoch in range(self.total_epochs):
            if epoch % self.training_ratio == 0:
                self.feature.need_grad = True
                self.clf.need_grad = True
                self.domain_output.need_grad = False
            else:
                self.feature.need_grad = False
                self.clf.need_grad = False
                self.domain_output.need_grad = True

            for i in range(self.max_iter):
                image, targets = train_loader.next()
                class_target = targets[:, 0].reshape(-1, 1)
                domain_target = targets[:, -1].reshape(-1, 1)
                self.input_image.d = image
                self.label.d = class_target
                self.domain_label.d = domain_target
                if epoch % self.training_ratio == 0:
                    self.loss.forward(clear_no_need_grad=True)
                    self.solver.zero_grad()
                    self.adv_loss.forward(clear_no_need_grad=True)
                    self.adv_loss.backward(clear_buffer=True)
                    self.solver.update()
                else:
                    self.domain_solver.zero_grad()
                    self.domain_loss.forward(clear_no_need_grad=True)
                    self.domain_loss.backward(clear_buffer=True)
                    self.domain_solver.update()
                if print_freq and (i % print_freq == 0):
                    print('Training epoch {}: [{}|{}], loss:{}, domain loss {}'.format(
                        epoch, i + 1, self.max_iter, self.loss.d, self.domain_loss.d), flush=True)

            if model_save and (epoch % self.training_ratio == 0) and (epoch > 1):
                acc, val_targets, val_scores = self.check_avg_precision(
                    val_loader, weights=self.validation_weight)
                if not os.path.exists(self.monitor_path):
                    os.makedirs(self.monitor_path)
                # computing model fairness
                cal_thresh = utils.calibrated_threshold(val_targets[:, 0],
                                                        val_scores)
                _, f1_thresh = utils.get_f1_threshold(val_targets[:, 0],
                                                      val_scores)
                val_pred = np.where(val_scores > cal_thresh, 1, 0)
                average_precision = utils.get_average_precision(val_targets[:, 0],
                                                                val_scores)
                deo = utils.get_difference_equality_opportunity(val_targets[:, 1],
                                                                val_targets[:, 0], val_pred)
                bias_amplification = utils.get_bias_amplification(val_targets[:, 1],
                                                                  val_targets[:, 0], val_pred)
                kl_divergence = utils.get_kl_divergence(val_targets[:, 1],
                                                        val_targets[:, 0], val_scores)
                val_results = {
                    'AP': average_precision,
                    'DEO': deo,
                    'BA': bias_amplification,
                    'KL': kl_divergence,
                    'f1_thresh': f1_thresh,
                    'cal_thresh': cal_thresh,
                    'epoch': epoch,
                    'accuracy': acc
                }
                if acc > self.acc:
                    self.acc = acc
                    path_best_acc_model = "{}/best_acc.h5".format(
                        self.monitor_path)
                    nn.save_parameters(path_best_acc_model)
                    print("val_results : ", val_results)
                    with open(self.monitor_path + '/val_results.pkl', 'wb+') as handle:
                        pickle.dump(val_results, handle)
                    with open(self.monitor_path + '/val_scores.pkl', 'wb+') as handle:
                        pickle.dump(val_scores, handle)
                    with open(self.monitor_path + '/val_targets.pkl', 'wb+') as handle:
                        pickle.dump(val_targets, handle)
                else:
                    path_adv_model = "{}/adversarial_model.h5".format(
                        self.monitor_path)
                    nn.save_parameters(path_adv_model)
                    with open(self.monitor_path + '/val_results_adv.pkl', 'wb+') as handle:
                        pickle.dump(val_results, handle)
                    with open(self.monitor_path + '/val_scores_adv.pkl', 'wb+') as handle:
                        pickle.dump(val_scores, handle)
                    with open(self.monitor_path + '/val_targets_adv.pkl', 'wb+') as handle:
                        pickle.dump(val_targets, handle)

    def check_avg_precision(self, loader, weights=None, print_out=True):
        """
        Compute average precision (AP) from prediction scores.
        Args:
            loader (NNabla data iterator) : specify the data iterator which to compute the AP
            weights : sample weights
            print_out (bool) : print the logs
        Return:
            acc (float) : accuracy of the model
            y_all (numpy.ndarray) : actual all target labels
            pred_all (numpy.ndarray) : predicted score
        """
        y_all, pred_all = self.get_scores(loader)
        acc = average_precision_score(
            y_all[:, 0], pred_all, sample_weight=weights)
        if print_out:
            print('Avg precision all = {}'.format(acc))
        return acc, y_all, pred_all

    def get_scores(self, loader):
        """
        get predicted scores of the  model
        Args:
            loader : specify the data loader which to predicted the scores
        Returns:
            y_all (numpy.ndarray) : actual all target labels
            pred_all (numpy.ndarray) : all predicted score

        """
        # only for validation and test set
        max_iter = int(loader.size / self.batch_size)
        rem_iter = 0
        if loader.size % self.batch_size != 0:
            rem_iter += 1
            rem_values = loader.size % self.batch_size
        y_all = []
        scores_all = []
        for _ in range(max_iter):
            image, targets = loader.next()
            self.x_v.d = image
            self.v_clf_out.forward(clear_buffer=True)
            y_all.append(targets)
            scores_all.append(self.v_clf_out.d.squeeze())
        for _ in range(rem_iter):
            loader._reset()
            image, targets = loader.next()
            self.x_v.d = image
            self.v_clf_out.forward(clear_buffer=True)
            y_all.append(targets[:rem_values])
            scores_all.append(self.v_clf_out.d.squeeze()[:rem_values])

        y_all = np.concatenate(y_all)
        pred_all = np.concatenate(scores_all)
        return y_all, pred_all


# responsible_ai/adversarial_debiasing_images/args.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse


def get_args():
    """
     parse args
     """
    description = "Adversarial debiasing"
    parser = argparse.ArgumentParser(description)
    parser.add_argument("-d", "--device-id", type=int, default=0,
                        help="Device ID of the GPU for training")
    parser.add_argument("-c", "--context", type=str, default="cudnn",
                        help="Extension path: cpu or cudnn.")
    parser.add_argument("--type-config", "-t", type=str, default='float',
                        help='Type of computation. e.g. "float", "half".')

    parser.add_argument('--model_train', choices=['baseline', 'adversarial'],
                        type=str, default='adversarial',
                        help="baseline model or Adversarial de-biasing ")
    parser.add_argument("--celeba_image_train_dir", type=str,
                        default="./data/train", help="Path of training directory")
    parser.add_argument("--celeba_image_valid_dir", type=str,
                        default="./data/valid", help="Path of validation directory")
    parser.add_argument("--attr_path", type=str, default=r'./data/celeba/list_attr_celeba.txt',
                        help="celebA attribute file path (ex: list_attr_celeba.txt)")

    parser.add_argument('--attribute', type=str, default='Attractive',
                        help="choose the target attribute to train "
                             "(e.g., Attractive, Arched EyeBrows, Bushy Eyebrows, smilling ,etc..")
    parser.add_argument('--protected_attribute', type=str, default='Male',
                        help="choose the protected attributes (e.g., Male , Pale_Skin)")

    # as per author's citation, the baseline model is trained on the CelebA training set
    # with 162,770 images.
    parser.add_argument("--train-beg", type=int, default=0,
                        help="start the training samples")
    parser.add_argument("--valid-beg", type=int, default=162752,
                        help="start the validation sample")
    parser.add_argument("--test-beg", type=int, default=182592,
                        help="start the test sample")
    parser.add_argument("--batch-size", "-b", type=int, default=32,
                        help="Batch size.")
    parser.add_argument("--learning-rate", "-l", type=float, default=1e-4,
                        help="learning rate is a configurable hyper parameter "
                             "used in the training of neural network")
    parser.add_argument('--max_iter', type=int, default=5086,
                        help='max iteration to train the model for an epoch')
    parser.add_argument('--total_epochs', type=int, default=50,
                        help='total no of epochs to train the model')
    # Adversary traning parametrs
    parser.add_argument('--training_ratio', type=int, default=3,
                        help='training ratio b/w the classifier and adversary n/w')
    parser.add_argument('--lamda', type=int, default=5,
                        help='lambda parameter weighs the adversarial loss of each class')

    # save model
    parser.add_argument("--model-save-path", "-o",
                        type=str, default='./results',
                        help='path where the model parameters saved.')
    parser.add_argument("--model-load-path",
                        type=str, default=None,
                        help='path where to load the .nnp model')

    opt = vars(parser.parse_args())
    return opt


# responsible_ai/adversarial_debiasing_images/basemodel.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import pickle
from sklearn.metrics import average_precision_score
import numpy as np
import nnabla as nn
import nnabla.solvers as S
import nnabla.functions as F
from nnabla.models.imagenet import ResNet50
import nnabla.parametric_functions as PF
from utils import utils


def clf_resnet50(layer, n_classes=1, train=True):
    """
    This function uses ResNet-50 pretrained on ImageNet as the base architecture
    and replaces the linear layer in ResNet with two linear layers with the
    hidden layer of size 2,048. Dropout and ReLU are applied
    between these
    """
    layer_1 = F.relu(PF.affine(layer, 2048, name='classifier_1'))
    if train:
        layer_1 = F.dropout(layer_1, 0.5)
    out = PF.affine(layer_1, n_classes, name='classifier_2')
    return out


class attribute_classifier:
    """
    This class uses ResNet-50 pretrained on ImageNet as the base architecture
    and replaces the linear layer in ResNet with two linear layers with the
    hidden layer of size 2,048. Dropout and ReLU are applied
    between these
    """

    def __init__(self, batch_size=32, learning_rate=1e-4, max_iter=5086,
                 total_epochs=50, monitor_path=None,
                 validation_weight=None, model_load_path=None):
        """
        Construct all the necessary attributes for the attribute classifier.
        Args:
            batch_size (int): number of samples contained in each generated batch
            learning_rate (float) : learning rate
            max_iter (int) : maximum iterations for an epoch
            total_epochs (int) : total epochs to train the model
            validation_weight : sample weights
            monitor_path (str) : model parameter to be saved
            model_load_path (str) : load the model
        """
        self.batch_size = batch_size
        # Resnet 50
        # training graph
        model = ResNet50()
        self.input_image = nn.Variable((self.batch_size,) + model.input_shape)
        self.label = nn.Variable([self.batch_size, 1])
        # fine tuning
        pool = model(self.input_image, training=True, use_up_to='pool')
        self.clf = clf_resnet50(pool, n_classes=1)
        self.clf.persistent = True
        # loss
        self.loss = F.mean(F.sigmoid_cross_entropy(self.clf, self.label))
        # hyper parameters
        self.solver = S.Adam(learning_rate)
        self.solver.set_parameters(nn.get_parameters())

        # validation graph
        self.x_v = nn.Variable((self.batch_size,) + model.input_shape)
        pool_v = model(self.x_v, training=False, use_up_to='pool')
        self.v_clf = clf_resnet50(pool_v, train=False)
        self.v_clf_out = F.sigmoid(self.v_clf)
        self.validation_weight = validation_weight
        # val params
        self.acc = 0.0
        self.total_epochs = total_epochs
        self.max_iter = max_iter
        self.monitor_path = monitor_path
        if model_load_path is not None:
            _ = nn.load_parameters(model_load_path)

    def train(self, train_loader, val_loader):
        """
        train the attribute classifier and save model with best accuracy on validation set
        Args:
             train_loader (NNabla data iterator) : specify training dataset
             val_loader (NNabla data iterator) : specify validation dataset
        """

        model_save = True
        print_freq = 100
        for epoch in range(self.total_epochs):
            for i in range(self.max_iter):
                image, targets = train_loader.next()
                targets = targets[:, 0].reshape(-1, 1)
                self.input_image.d = image
                self.label.d = targets
                self.solver.zero_grad()
                self.loss.forward(clear_no_need_grad=True)
                self.loss.backward(clear_buffer=True)
                self.solver.update()

                if print_freq and (i % print_freq == 0):
                    print('Training epoch {}: [{}|{}], loss:{}'.format(
                        epoch, i + 1, self.max_iter, self.loss.d), flush=True)
            if model_save:
                acc, val_targets, val_scores = \
                    self.check_avg_precision(
                        val_loader, weights=self.validation_weight)
                if acc > self.acc:
                    self.acc = acc
                    if not os.path.exists(self.monitor_path):
                        os.makedirs(self.monitor_path)
                    path_best_acc_model = "{}/best_acc.h5".format(
                        self.monitor_path)
                    nn.save_parameters(path_best_acc_model)

                    # computing model fairness
                    cal_thresh = utils.calibrated_threshold(val_targets[:, 0],
                                                            val_scores)
                    _, f1_thresh = utils.get_f1_threshold(val_targets[:, 0],
                                                          val_scores)
                    val_pred = np.where(val_scores > cal_thresh, 1, 0)
                    with open(self.monitor_path + '/val_scores.pkl', 'wb+') as handle:
                        pickle.dump(val_scores, handle)
                    with open(self.monitor_path + '/val_targets.pkl', 'wb+') as handle:
                        pickle.dump(val_targets, handle)
                    average_precision = utils.get_average_precision(val_targets[:, 0],
                                                                    val_scores)
                    deo = utils.get_difference_equality_opportunity(val_targets[:, 1],
                                                                    val_targets[:, 0], val_pred)
                    bias_amplification = utils.get_bias_amplification(val_targets[:, 1],
                                                                      val_targets[:, 0], val_pred)
                    kl_divergence = utils.get_kl_divergence(val_targets[:, 1],
                                                            val_targets[:, 0], val_scores)
                    val_results = {
                        'AP': average_precision,
                        'DEO': deo,
                        'BA': bias_amplification,
                        'KL': kl_divergence,
                        'f1_thresh': f1_thresh,
                        'cal_thresh': cal_thresh,
                        'epoch': epoch,
                        'accuracy': acc
                    }
                    with open(self.monitor_path + '/val_results.pkl', 'wb+') as handle:
                        pickle.dump(val_results, handle)

    def check_avg_precision(self, loader, weights=None, print_out=True):
        """
        Compute average precision (AP) from prediction scores.
        Args:
            loader (NNabla data iterator) : specify the data iterator which to compute the AP
            weights : sample weights
            print_out (bool) : print the logs
        Return:
            acc (float) : accuracy of the model
            y_all (numpy.ndarray) : actual all target labels
            pred_all (numpy.ndarray) : predicted score
        """
        y_all, pred_all = self.get_scores(loader)
        acc = average_precision_score(
            y_all[:, 0], pred_all, sample_weight=weights)
        if print_out:
            print('Avg precision all = {}'.format(acc))
        return acc, y_all, pred_all

    def get_scores(self, loader):
        """
        get predicted scores of the  model
        Args:
            loader : specify the data loader which to predicted the scores
        Returns:
            y_all (numpy.ndarray) : actual all target labels
            pred_all (numpy.ndarray) : all predicted score
        """
        # only for validation and test set
        max_iter = int(loader.size/self.batch_size)
        rem_iter = 0
        if loader.size % self.batch_size != 0:
            rem_iter += 1
            rem_values = loader.size % self.batch_size
        y_all = []
        scores_all = []
        for _ in range(max_iter):
            image, targets = loader.next()
            self.x_v.d = image
            self.v_clf_out.forward(clear_buffer=True)
            y_all.append(targets)
            scores_all.append(self.v_clf_out.d.squeeze())
        for _ in range(rem_iter):
            loader._reset()
            image, targets = loader.next()
            self.x_v.d = image
            self.v_clf_out.forward(clear_buffer=True)
            y_all.append(targets[:rem_values])
            scores_all.append(self.v_clf_out.d.squeeze()[:rem_values])

        y_all = np.concatenate(y_all)
        pred_all = np.concatenate(scores_all)
        return y_all, pred_all


# responsible_ai/adversarial_debiasing_images/data_loader.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import numpy as np
from PIL import Image
import albumentations as A
from nnabla.utils.data_iterator import data_iterator_simple


def data_iterator_celeba(img_path, attr_path, batch_size,
                         target_attribute='Attractive', protected_attribute='Male',
                         num_samples=-1, augment=False, shuffle=False, rng=None):
    """
    create celebA data iterator
    Args:
        img_path (str) : image path directory
        attr_path (str) : celebA attribute file path (ex: list_attr_celeba.txt)
        batch_size (int) :  number of samples contained in each generated batch
        target_attribute (str) : target attribute (ex: Arched EyeBrows,
                                Bushy Eyebrows, smiling,etc..)
        protected_attribute (str): protected attribute (ex: Male, Pale_Skin)
        num_samples (int) : number of samples taken in data loader
                            (if num_samples=-1, it will take all the images in the dataset)
        augment (bool) : data augmentation (True for training)
        shuffle (bool) : shuffle the data (True /False)
        rng : None
    Returns:
        simple data iterator
    """

    imgs = []
    for file in sorted(os.listdir(img_path), key=lambda x: int(x.split(".")[0])):
        imgs.append(os.path.join(img_path, file))
    with open(attr_path, 'r') as f:
        lines = f.readlines()

    attr_list = lines[1].strip().split()
    attr_idx_dict = {attr: i for i, attr in enumerate(attr_list)}
    labels_dict = {}
    for line in lines[2:]:
        line = line.strip().split()
        key = line[0]
        attr = line[1:]
        labels_dict[key] = np.array([int((int(attr[attr_idx_dict[target_attribute]]) + 1) / 2),
                                     int((int(attr[attr_idx_dict[protected_attribute]]) + 1) / 2)])

    # as per the author's citation, we have transformed the input image
    # (resize to , 256 * 256, 224 * 224)
    pre_process = [(256, 256), (224, 224)]
    mean_normalize = (0.485, 0.456, 0.406)
    std_normalize = (0.229, 0.224, 0.225)

    if augment:
        transform = A.Compose([
            A.Resize(pre_process[0][0], pre_process[0][1]),
            A.RandomCrop(width=pre_process[1][0], height=pre_process[1][1]),
            A.HorizontalFlip(p=0.5),
            A.Normalize(mean=mean_normalize, std=std_normalize)
        ])
    else:
        transform = A.Compose([
            A.Resize(pre_process[0][0], pre_process[0][1]),
            A.CenterCrop(width=pre_process[1][0], height=pre_process[1][1]),
            A.Normalize(mean=mean_normalize, std=std_normalize)
        ])
    if num_samples == -1:
        num_samples = len(imgs)
    else:
        print("Num. of data ({}) is used for debugging".format(num_samples))

    def load_func(i):
        img = Image.open(imgs[i])
        img = np.array(img.convert('RGB'))
        # transform
        transformed_image = transform(image=img)['image'].transpose(2, 0, 1)
        return transformed_image, labels_dict[os.path.basename(imgs[i])]

    return data_iterator_simple(load_func, num_samples, batch_size,
                                shuffle=shuffle, rng=rng, with_file_cache=False)


# responsible_ai/adversarial_debiasing_images/main.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import args
import numpy as np
import nnabla as nn
from nnabla.ext_utils import get_extension_context
import adversarial as ad
import basemodel as bm
import data_loader as di
from utils import utils


def main():
    """
    main method
    """

    opt = args.get_args()
    ctx = get_extension_context(
        opt['context'], device_id=opt['device_id'], type_config=opt['type_config'])

    nn.set_default_context(ctx)
    # model configurations
    batch_size = opt['batch_size']
    learning_rate = opt['learning_rate']
    max_iter = opt['max_iter']
    total_epochs = opt['total_epochs']
    monitor_path = opt['model_save_path']
    rng = np.random.RandomState(1)
    train_loader = di.data_iterator_celeba(
        opt['celeba_image_train_dir'], opt['attr_path'], batch_size,
        target_attribute=opt['attribute'], protected_attribute=opt['protected_attribute'],
        augment=True, shuffle=True, rng=rng)
    valid_loader = di.data_iterator_celeba(
        opt['celeba_image_valid_dir'], opt['attr_path'], batch_size,
        target_attribute=opt['attribute'], protected_attribute=opt['protected_attribute'],
        rng=rng)
    validation_weight = utils.compute_class_weight(valid_loader)
    if opt["model_train"] == 'adversarial':
        training_ratio = opt['training_ratio']
        lamda = opt['lamda']
        adv = ad.adversarial(batch_size=batch_size, total_epochs=total_epochs,
                             learning_rate=learning_rate, max_iter=max_iter,
                             training_ratio=training_ratio, lamda=lamda,
                             monitor_path=monitor_path, validation_weight=validation_weight)
        adv.train(train_loader, valid_loader)
    elif opt["model_train"] == 'baseline':
        clf = bm.attribute_classifier(batch_size=batch_size, total_epochs=total_epochs,
                                      learning_rate=learning_rate,
                                      max_iter=max_iter,
                                      validation_weight=validation_weight,
                                      monitor_path=monitor_path)
        clf.train(train_loader, valid_loader)
    else:
        print("Please provide proper argument")
        return


if __name__ == '__main__':
    main()


# responsible_ai/adversarial_debiasing_images/split_dataset.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import argparse
import shutil


def parse_args():
    """
    parse args
    """
    description = "Split the CelebA dataset "
    parser = argparse.ArgumentParser(description)
    parser.add_argument("--img_path", type=str, default=r'./data/celeba/images',
                        help="Source image path directory")
    parser.add_argument("--attr_path", type=str, default=r'./data/celeba/list_attr_celeba.txt',
                        help="celebA attribute file path (ex: list_attr_celeba.txt)")
    parser.add_argument("--out_dir", type=str, default=r'./train',
                        help='Path where the split data to be saved')
    parser.add_argument("--split", type=str, default='train',
                        help='split the dataset depends on the '
                             'split attribute(train, valid, and test)')
    opt = vars(parser.parse_args())

    return opt


def split_celeba_dataset(img_path, attr_path, out_dir, split="test"):
    """
    split the celebA dataset
    Args:
        img_path (str) : image path directory 
        attr_path (str) : celebA attribute file path (ex: list_attr_celeba.txt)
        out_dir (str) : Path where the split data to be saved
        split (string) : split the dataset depends on the split attribute(train, valid, and test)
    """
    # as per the author's citation, we are splitting the dataset
    train_beg = 0  # train starts from
    valid_beg = 162770  # valid starts from
    test_beg = 182610  # test starts from

    with open(attr_path, 'r') as label_file:
        label_file = label_file.readlines()

    # skipping the first two rows for header
    total_samples = len(label_file) - 2
    if split == 'train':
        number_samples = valid_beg - train_beg
        beg = train_beg

    elif split == 'valid':
        number_samples = test_beg - valid_beg
        beg = valid_beg

    elif split == 'test':
        number_samples = total_samples - test_beg
        beg = test_beg
    else:
        print('Error')
        return
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)

    for i in range(beg + 2, beg + number_samples + 2):
        temp = label_file[i].strip().split()
        src_dir = os.path.join(img_path, temp[0])
        dst_dir = os.path.join(out_dir, temp[0])
        shutil.copy(src_dir, dst_dir)
    print("splitting completed")


if __name__ == '__main__':
    opt = parse_args()
    split_celeba_dataset(opt['img_path'], opt['attr_path'],
                         opt['out_dir'], opt['split'])


# responsible_ai/attention_branch_network/abn.py

import nnabla.functions as F
import nnabla.parametric_functions as PF


def abn(x, test=False):
    # Input:x -> 3,32,32

    if not test:
        h1 = F.image_augmentation(x, shape=(3, 32, 32), min_scale=1, max_scale=1, angle=0.26,
                                  aspect_ratio=1, distortion=0, flip_lr=True, brightness=0, contrast=1, noise=0)
        h1 = F.random_shift(h1, (4, 4))
        h1 = F.random_flip(h1, (3,))
    else:
        h1 = x

    # MulScalar_2
    h1 = F.mul_scalar(h1, 0.01735)
    # AddScalar_2
    h1 = F.add_scalar(h1, -1.99)
    # Convolution -> 16,32,32
    h1 = PF.convolution(h1, 16, (3, 3), (1, 1),
                        with_bias=False, name='Convolution')
    # RepeatStart_1
    for i in range(18):

        # BatchNormalization_2
        h2 = PF.batch_normalization(
            h1, (1,), 0.5, 0.0001, not test, name='BatchNormalization_2[' + str(i) + ']')
        # ReLU_2
        h2 = F.relu(h2, True)
        # Convolution_2
        h2 = PF.convolution(h2, 16, (3, 3), (1, 1), with_bias=False,
                            name='Convolution_2[' + str(i) + ']')
        # BatchNormalization_11
        h2 = PF.batch_normalization(
            h2, (1,), 0.5, 0.0001, not test, name='BatchNormalization_11[' + str(i) + ']')
        # ReLU_6
        h2 = F.relu(h2, True)
        # Convolution_3
        h2 = PF.convolution(h2, 16, (3, 3), (1, 1), with_bias=False,
                            name='Convolution_3[' + str(i) + ']')

        # Add2 -> 16,32,32
        h3 = F.add2(h1, h2, False)
        # RepeatEnd_1
        h1 = h3
    # BatchNormalization_3
    h3 = PF.batch_normalization(
        h3, (1,), 0.5, 0.0001, not test, name='BatchNormalization_3')
    # ReLU_3
    h3 = F.relu(h3, True)

    # Convolution_15 -> 32,16,16
    h4 = PF.convolution(h3, 32, (3, 3), (1, 1), (2, 2),
                        with_bias=False, name='Convolution_15')

    # Convolution_4 -> 32,16,16
    h5 = PF.convolution(h3, 32, (3, 3), (1, 1), (2, 2),
                        with_bias=False, name='Convolution_4')
    # BatchNormalization_4
    h5 = PF.batch_normalization(
        h5, (1,), 0.5, 0.0001, not test, name='BatchNormalization_4')
    # ReLU_4
    h5 = F.relu(h5, True)
    # Convolution_10
    h5 = PF.convolution(h5, 32, (3, 3), (1, 1),
                        with_bias=False, name='Convolution_10')

    # Add2_5 -> 32,16,16
    h5 = F.add2(h5, h4, True)
    # RepeatStart_2
    for i in range(17):

        # BatchNormalization_5
        h6 = PF.batch_normalization(
            h5, (1,), 0.5, 0.0001, not test, name='BatchNormalization_5[' + str(i) + ']')
        # ReLU_5
        h6 = F.relu(h6, True)
        # Convolution_5
        h6 = PF.convolution(h6, 32, (3, 3), (1, 1), with_bias=False,
                            name='Convolution_5[' + str(i) + ']')
        # BatchNormalization_6
        h6 = PF.batch_normalization(
            h6, (1,), 0.5, 0.0001, not test, name='BatchNormalization_6[' + str(i) + ']')
        # ReLU_7
        h6 = F.relu(h6, True)
        # Convolution_6
        h6 = PF.convolution(h6, 32, (3, 3), (1, 1), with_bias=False,
                            name='Convolution_6[' + str(i) + ']')

        # Add2_2 -> 32,16,16
        h7 = F.add2(h5, h6, False)
        # RepeatEnd_2
        h5 = h7
    # BatchNormalization_7
    h7 = PF.batch_normalization(
        h7, (1,), 0.5, 0.0001, not test, name='BatchNormalization_7')
    # ReLU_9
    h7 = F.relu(h7, True)

    # Convolution_17 -> 64,8,8
    h8 = PF.convolution(h7, 64, (3, 3), (1, 1), (2, 2),
                        with_bias=False, name='Convolution_17')

    # Convolution_7 -> 64,8,8
    h9 = PF.convolution(h7, 64, (3, 3), (1, 1), (2, 2),
                        with_bias=False, name='Convolution_7')
    # BatchNormalization_12
    h9 = PF.batch_normalization(
        h9, (1,), 0.5, 0.0001, not test, name='BatchNormalization_12')
    # ReLU_11
    h9 = F.relu(h9, True)
    # Convolution_11
    h9 = PF.convolution(h9, 64, (3, 3), (1, 1),
                        with_bias=False, name='Convolution_11')

    # Add2_6 -> 64,8,8
    h10 = F.add2(h9, h8, True)

    # RepeatStart_3
    for i in range(17):

        # BatchNormalization_14
        h11 = PF.batch_normalization(
            h10, (1,), 0.5, 0.0001, not test, name='BatchNormalization_14[' + str(i) + ']')
        # ReLU_13
        h11 = F.relu(h11, True)
        # Convolution_16
        h11 = PF.convolution(
            h11, 64, (3, 3), (1, 1), with_bias=False, name='Convolution_16[' + str(i) + ']')
        # BatchNormalization_15
        h11 = PF.batch_normalization(
            h11, (1,), 0.5, 0.0001, not test, name='BatchNormalization_15[' + str(i) + ']')
        # ReLU_14
        h11 = F.relu(h11, True)
        # Convolution_18
        h11 = PF.convolution(
            h11, 64, (3, 3), (1, 1), with_bias=False, name='Convolution_18[' + str(i) + ']')

        # Add2_4 -> 64,8,8
        h12 = F.add2(h10, h11, False)
        # RepeatEnd_3
        h10 = h12
    # BatchNormalization_16
    h12 = PF.batch_normalization(
        h12, (1,), 0.9, 0.0001, not test, name='BatchNormalization_16')
    # ReLU_15
    h12 = F.relu(h12, True)
    # Convolution_14
    h12 = PF.convolution(h12, 64, (1, 1), (0, 0), name='Convolution_14')
    # BatchNormalization_10
    h12 = PF.batch_normalization(
        h12, (1,), 0.9, 0.0001, not test, name='BatchNormalization_10')
    # ReLU
    h12 = F.relu(h12, True)

    # Convolution_12 -> 10,8,8
    h13 = PF.convolution(h12, 10, (1, 1), (0, 0), name='Convolution_12')

    # Convolution_13 -> 1,8,8
    h14 = PF.convolution(h12, 1, (3, 3), (1, 1), name='Convolution_13')
    # BatchNormalization
    h14 = PF.batch_normalization(
        h14, (1,), 0.9, 0.0001, not test, name='BatchNormalization')
    # GlobalAveragePooling -> 10,1,1
    h13 = F.global_average_pooling(h13)
    # Reshape -> 10
    h13 = F.reshape(h13, (h13.shape[0], 10,))
    # Sigmoid_Attention_Map
    h14 = F.sigmoid(h14)

    # Mul2 -> 64,8,8
    h15 = F.mul2(h14, h10)

    # Add2_7 -> 64,8,8
    h16 = F.add2(h10, h15, False)
    # RepeatStart_4
    for i in range(17):

        # BatchNormalization_8
        h17 = PF.batch_normalization(
            h16, (1,), 0.5, 0.0001, not test, name='BatchNormalization_8[' + str(i) + ']')
        # ReLU_8
        h17 = F.relu(h17, True)
        # Convolution_8
        h17 = PF.convolution(
            h17, 64, (3, 3), (1, 1), with_bias=False, name='Convolution_8[' + str(i) + ']')
        # BatchNormalization_9
        h17 = PF.batch_normalization(
            h17, (1,), 0.5, 0.0001, not test, name='BatchNormalization_9[' + str(i) + ']')
        # ReLU_10
        h17 = F.relu(h17, True)
        # Convolution_9
        h17 = PF.convolution(
            h17, 64, (3, 3), (1, 1), with_bias=False, name='Convolution_9[' + str(i) + ']')

        # Add2_3 -> 64,8,8
        h18 = F.add2(h16, h17, False)
        # RepeatEnd_4
        h16 = h18
    # BatchNormalization_13
    h18 = PF.batch_normalization(
        h18, (1,), 0.5, 0.0001, not test, name='BatchNormalization_13')
    # AveragePooling -> 64,1,1
    h18 = F.average_pooling(h18, (8, 8), (8, 8))
    # ReLU_12
    h18 = F.relu(h18, True)
    # Affine -> 10
    h18 = PF.affine(h18, (10,), name='Affine')
    return h13, h18


# responsible_ai/attention_branch_network/data.py

import argparse
import numpy as np
import tarfile
import os
import tqdm
import csv
from imageio import imwrite
from nnabla.logger import logger
from nnabla.utils.data_iterator import data_iterator
from nnabla.utils.data_source import DataSource
from nnabla.utils.data_source_loader import download


class Cifar10DataSource(DataSource):
    '''
    Get data directly from cifar10 dataset from Internet(yann.lecun.com).
    '''

    def _get_data(self, position):
        image = self._images[self._indexes[position]]
        label = self._labels[self._indexes[position]]
        return (image, label)

    def __init__(self, train=True, shuffle=False, rng=None):
        super(Cifar10DataSource, self).__init__(shuffle=shuffle)

        self._train = train
        data_uri = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"
        logger.info('Getting labeled data from {}.'.format(data_uri))
        r = download(data_uri)  # file object returned
        with tarfile.open(fileobj=r, mode="r:gz") as fpin:
            # Training data
            if train:
                images = []
                labels = []
                for member in fpin.getmembers():
                    if "data_batch" not in member.name:
                        continue
                    fp = fpin.extractfile(member)
                    data = np.load(fp, allow_pickle=True, encoding="bytes")
                    images.append(data[b"data"])
                    labels.append(data[b"labels"])
                self._size = 50000
                self._images = np.concatenate(
                    images).reshape(self._size, 3, 32, 32)
                self._labels = np.concatenate(labels).reshape(-1, 1)
            # Validation data
            else:
                for member in fpin.getmembers():
                    if "test_batch" not in member.name:
                        continue
                    fp = fpin.extractfile(member)
                    data = np.load(fp, allow_pickle=True, encoding="bytes")
                    images = data[b"data"]
                    labels = data[b"labels"]
                self._size = 10000
                self._images = images.reshape(self._size, 3, 32, 32)
                self._labels = np.array(labels).reshape(-1, 1)
        r.close()
        logger.info('Getting labeled data from {}.'.format(data_uri))

        self._size = self._labels.size
        self._variables = ('x', 'y')
        if rng is None:
            rng = np.random.RandomState(313)
        self.rng = rng
        self.reset()

    def reset(self):
        if self._shuffle:
            self._indexes = self.rng.permutation(self._size)
        else:
            self._indexes = np.arange(self._size)
        super(Cifar10DataSource, self).reset()

    @property
    def images(self):
        """Get copy of whole data with a shape of (N, 1, H, W)."""
        return self._images.copy()

    @property
    def labels(self):
        """Get copy of whole label with a shape of (N, 1)."""
        return self._labels.copy()


def data_iterator_cifar10(batch_size,
                          train=True,
                          rng=None,
                          shuffle=True,
                          with_memory_cache=False,
                          with_file_cache=False):
    '''
    Provide DataIterator with :py:class:`Cifar10DataSource`
    with_memory_cache, with_parallel and with_file_cache option's default value is all False,
    because :py:class:`Cifar10DataSource` is able to store all data into memory.

    For example,
    '''

    ds = Cifar10DataSource(train=train, shuffle=shuffle, rng=rng)
    di = data_iterator(ds, batch_size, rng=rng,
                       with_memory_cache=with_memory_cache, with_file_cache=with_file_cache)
    return di


def data_iterator_to_csv(csv_path, csv_file_name, data_path, data_iterator):
    index = 0
    csv_data = []
    with data_iterator as data:
        line = ['x:image', 'y:label']
        csv_data.append(line)
        pbar = tqdm.tqdm(total=data.size, unit='images')
        initial_epoch = data.epoch
        while data.epoch == initial_epoch:
            d = data.next()
            for i in range(len(d[0])):
                label = d[1][i][0]
                file_name = data_path + \
                    '/{}'.format(label) + '/{}.png'.format(index)
                full_path = os.path.join(
                    csv_path, file_name.replace('/', os.path.sep))
                directory = os.path.dirname(full_path)
                if not os.path.exists(directory):
                    os.makedirs(directory)
                imwrite(full_path, d[0][i].reshape(
                    3, 32, 32).transpose(1, 2, 0))
                csv_data.append([file_name, label])
                index += 1
                pbar.update(1)
        pbar.close()
    with open(os.path.join(csv_path, csv_file_name), 'w') as f:
        writer = csv.writer(f, lineterminator='\n')
        writer.writerows(csv_data)
    return csv_data


def func(args):
    path = args.output_dir

    # Create original training set
    logger.log(99, 'Downloading CIFAR-10 dataset...')
    train_di = data_iterator_cifar10(50000, True, None, False)
    logger.log(99, 'Creating "cifar10_training.csv"... ')
    data_iterator_to_csv(
        path, 'cifar10_training.csv', './training', train_di)

    # Create original test set
    validation_di = data_iterator_cifar10(10000, False, None, False)
    logger.log(99, 'Creating "cifar10_test.csv"... ')
    data_iterator_to_csv(
        path, 'cifar10_test.csv', './validation', validation_di)

    logger.log(99, 'Dataset creation completed successfully.')


def main():
    parser = argparse.ArgumentParser(
        description='CIFAR10\n\n' +
        'Download CIFAR-10 dataset from dl.sony.com (original file is from https://www.cs.toronto.edu/~kriz/cifar.html).\n\n',
        formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument(
        '-o',
        '--output-dir',
        help='path to write NNC dataset CSV format (dir) default=CIFAR10',
        required=True)
    parser.set_defaults(func=func)

    args = parser.parse_args()
    args.func(args)


if __name__ == '__main__':
    main()


# responsible_ai/attention_branch_network/infer.py

import functools
from tqdm import tqdm
import nnabla as nn
from nnabla.ext_utils import get_extension_context
from responsible_ai.attention_branch_network.data import data_iterator_cifar10
from responsible_ai.attention_branch_network.resnet110 import resnet110
from responsible_ai.attention_branch_network.abn import abn


def infer(model="abn", parameter="./params_30.h5", device_id=0, type_config='float'):
    extension_module = "cudnn"
    ctx = get_extension_context(
        extension_module, device_id=device_id, type_config=type_config
    )
    nn.set_default_context(ctx)

    mch = 0
    path = parameter
    nn.clear_parameters()

    input_data = data_iterator_cifar10(
        1, train=False, shuffle=False)

    x = nn.Variable((1, 3, 32, 32))
    x.data.zero()

    t = nn.Variable((1, 1))

    if model == "resnet110":
        prediction = functools.partial(resnet110)
        nn.load_parameters(path)
        y = prediction(x=x, test=True)
    elif model == "abn":
        prediction = functools.partial(abn)
        nn.load_parameters(path)
        _, y = prediction(x=x, test=True)

    nn.get_parameters()

    for _ in tqdm(range(input_data.size)):
        x.d, t.d = input_data.next()
        y.forward()

        if t.d == y.d.argmax(axis=1):
            mch += 1

    print("Accuracy:{}".format(mch / input_data.size))


# responsible_ai/attention_branch_network/resnet110.py

import nnabla.functions as F
import nnabla.parametric_functions as PF


def resnet110(x, test=False):
    # Input:x -> 3,32,32

    # ImageAugmentation
    if not test:
        h = F.image_augmentation(x, shape=(3, 32, 32), min_scale=1, max_scale=1, angle=0.26,
                                 aspect_ratio=1, distortion=0, flip_lr=True, brightness=0, contrast=1, noise=0)
        h = F.random_shift(h, shifts=(4, 4))
        h = F.random_flip(h, axes=(3,))
    else:
        h = x

    # MulScalar_2
    h = F.mul_scalar(h, val=0.01735)
    # AddScalar_2
    h = F.add_scalar(h, val=-1.99)
    # Convolution -> 16,32,32
    h = PF.convolution(h, 16, (3, 3), (1, 1),
                       with_bias=False, name='Convolution')
    # RepeatStart_1
    for i in range(18):

        # BatchNormalization_2
        h1 = PF.batch_normalization(
            h, (1,), 0.5, 0.0001, not test, name='BatchNormalization_2[' + str(i) + ']')
        # ReLU_2
        h1 = F.relu(h1, True)
        # Convolution_2
        h1 = PF.convolution(h1, 16, (3, 3), (1, 1), with_bias=False,
                            name='Convolution_2[' + str(i) + ']')
        # BatchNormalization_11
        h1 = PF.batch_normalization(
            h1, (1,), 0.5, 0.0001, not test, name='BatchNormalization_11[' + str(i) + ']')
        # ReLU_6
        h1 = F.relu(h1, True)
        # Convolution_3
        h1 = PF.convolution(h1, 16, (3, 3), (1, 1), with_bias=False,
                            name='Convolution_3[' + str(i) + ']')

        # Add2 -> 16,32,32
        h2 = F.add2(h, h1, inplace=False)
        # RepeatEnd_3
        h = h2
    # BatchNormalization_3
    h2 = PF.batch_normalization(
        h2, (1,), 0.5, 0.0001, not test, name='BatchNormalization_3')
    # ReLU_3
    h2 = F.relu(h2, inplace=True)

    # Convolution_15 -> 32,16,16
    h3 = PF.convolution(h2, 32, (3, 3), (1, 1), (2, 2),
                        with_bias=False, name='Convolution_15')

    # Convolution_4 -> 32,16,16
    h4 = PF.convolution(h2, 32, (3, 3), (1, 1), (2, 2),
                        with_bias=False, name='Convolution_4')
    # BatchNormalization_4
    h4 = PF.batch_normalization(
        h4, (1,), 0.5, 0.0001, not test, name='BatchNormalization_4')
    # ReLU_4
    h4 = F.relu(h4, True)
    # Convolution_10
    h4 = PF.convolution(h4, 32, (3, 3), (1, 1),
                        with_bias=False, name='Convolution_10')

    # Add2_5 -> 32,16,16
    h4 = F.add2(h4, h3, inplace=True)
    # RepeatStart
    for i in range(17):

        # BatchNormalization_5
        h5 = PF.batch_normalization(
            h4, (1,), 0.5, 0.0001, not test, name='BatchNormalization_5[' + str(i) + ']')
        # ReLU_5
        h5 = F.relu(h5, True)
        # Convolution_5
        h5 = PF.convolution(h5, 32, (3, 3), (1, 1), with_bias=False,
                            name='Convolution_5[' + str(i) + ']')
        # BatchNormalization_6
        h5 = PF.batch_normalization(
            h5, (1,), 0.5, 0.0001, not test, name='BatchNormalization_6[' + str(i) + ']')
        # ReLU_7
        h5 = F.relu(h5, True)
        # Convolution_6
        h5 = PF.convolution(h5, 32, (3, 3), (1, 1), with_bias=False,
                            name='Convolution_6[' + str(i) + ']')

        # Add2_2 -> 32,16,16
        h6 = F.add2(h4, h5, inplace=False)
        # RepeatEnd
        h4 = h6
    # BatchNormalization_7
    h6 = PF.batch_normalization(
        h6, (1,), 0.5, 0.0001, not test, name='BatchNormalization_7')
    # ReLU_9
    h6 = F.relu(h6, inplace=True)

    # Convolution_17 -> 64,8,8
    h7 = PF.convolution(h6, 64, (3, 3), (1, 1), (2, 2),
                        with_bias=False, name='Convolution_17')

    # Convolution_7 -> 64,8,8
    h8 = PF.convolution(h6, 64, (3, 3), (1, 1), (2, 2),
                        with_bias=False, name='Convolution_7')
    # BatchNormalization_12
    h8 = PF.batch_normalization(
        h8, (1,), 0.5, 0.0001, not test, name='BatchNormalization_12')
    # ReLU_11
    h8 = F.relu(h8, True)
    # Convolution_11
    h8 = PF.convolution(h8, 64, (3, 3), (1, 1),
                        with_bias=False, name='Convolution_11')

    # Add2_6 -> 64,8,8
    h8 = F.add2(h8, h7, inplace=True)
    # RepeatStart_2
    for i in range(17):

        # BatchNormalization_8
        h9 = PF.batch_normalization(
            h8, (1,), 0.5, 0.0001, not test, name='BatchNormalization_8[' + str(i) + ']')
        # ReLU_8
        h9 = F.relu(h9, inplace=True)
        # Convolution_8
        h9 = PF.convolution(h9, 64, (3, 3), (1, 1), with_bias=False,
                            name='Convolution_8[' + str(i) + ']')
        # BatchNormalization_9
        h9 = PF.batch_normalization(
            h9, (1,), 0.5, 0.0001, not test, name='BatchNormalization_9[' + str(i) + ']')
        # ReLU_10
        h9 = F.relu(h9, inplace=True)
        # Convolution_9
        h9 = PF.convolution(h9, 64, (3, 3), (1, 1), with_bias=False,
                            name='Convolution_9[' + str(i) + ']')

        # Add2_3 -> 64,8,8
        h10 = F.add2(h8, h9, inplace=False)
        # RepeatEnd_2
        h8 = h10
    # BatchNormalization_13
    h10 = PF.batch_normalization(
        h10, (1,), 0.5, 0.0001, not test, name='BatchNormalization_13')
    # AveragePooling -> 64,1,1
    h10 = F.average_pooling(h10, kernel=(8, 8), stride=(8, 8))
    # ReLU_12
    h10 = F.relu(h10, inplace=True)
    # Affine -> 10
    h10 = PF.affine(h10, n_outmaps=(10,), name='Affine')
    return h10


# responsible_ai/attention_branch_network/train.py

import os
import functools
import math
import numpy as np
from tqdm import tqdm
import nnabla as nn
import nnabla.functions as F
import nnabla.solvers as S
import nnabla.utils.save as save
from nnabla.ext_utils import get_extension_context
from nnabla.utils.data_iterator import data_iterator_csv_dataset
from nnabla.utils.learning_rate_scheduler import StepScheduler
from nnabla.monitor import Monitor, MonitorSeries, MonitorTimeElapsed
from responsible_ai.attention_branch_network.data import data_iterator_cifar10, data_iterator_to_csv
from responsible_ai.attention_branch_network.resnet110 import resnet110
from responsible_ai.attention_branch_network.abn import abn
from utils.neu.save_nnp import save_nnp


def categorical_error(pred, label):
    pred_label = pred.argmax(1)
    return (pred_label != label.flat).mean()


def loss_function(pred, label):
    loss = F.mean(F.softmax_cross_entropy(pred, label))
    return loss


def train(batch_size=128, device_id=0, type_config='float',
          model_save_path="./", train_epochs=30, model="abn"):
    bs_train = batch_size
    bs_valid = batch_size
    extension_module = "cudnn"
    ctx = get_extension_context(
        extension_module, device_id=device_id, type_config=type_config
    )
    nn.set_default_context(ctx)

    n_train_samples = 50000
    n_val_samples = 10000

    # Data Iterator
    train_di = data_iterator_cifar10(50000, True, None, False)
    train_csv = data_iterator_to_csv(
        './', 'cifar10_training.csv', './training', train_di)
    # Create original test set
    validation_di = data_iterator_cifar10(10000, False, None, False)
    test_csv = data_iterator_to_csv(
        './', 'cifar10_test.csv', './validation', validation_di)

    train_loader = data_iterator_csv_dataset(
        './cifar10_training.csv', batch_size=bs_train, shuffle=True, rng=np.random.RandomState(0), normalize=False, with_file_cache=False)
    val_loader = data_iterator_csv_dataset(
        './cifar10_test.csv', batch_size=bs_valid, shuffle=False, normalize=False, with_file_cache=False)

    if model == "resnet110":
        prediction = functools.partial(resnet110)

        # Create training graphs
        test = False
        image_train = nn.Variable((bs_train, 3, 32, 32))
        label_train = nn.Variable((bs_train, 1))
        pred_train = prediction(x=image_train, test=test)
        loss_train = loss_function(pred_train, label_train)

        # Create validation graph
        test = True
        image_valid = nn.Variable((bs_valid, 3, 32, 32))
        label_valid = nn.Variable((bs_valid, 1))
        pred_valid = prediction(x=image_valid, test=test)
        loss_val = loss_function(pred_valid, label_valid)

    elif model == "abn":
        prediction = functools.partial(abn)

        # Create training graphs
        test = False
        image_train = nn.Variable((bs_train, 3, 32, 32))
        label_train = nn.Variable((bs_train, 1))
        pred_att, pred_train = prediction(x=image_train, test=test)
        loss_att = loss_function(pred_att, label_train)
        loss_per = loss_function(pred_train, label_train)
        loss_train = loss_att + loss_per

        # Create validation graph
        test = True
        image_valid = nn.Variable((bs_valid, 3, 32, 32))
        label_valid = nn.Variable((bs_valid, 1))
        _, pred_valid = prediction(x=image_valid, test=test)
        loss_val = loss_function(pred_valid, label_valid)

    for param in nn.get_parameters().values():
        param.grad.zero()

    solver = S.Momentum(lr=0.1, momentum=0.9)
    solver.set_parameters(nn.get_parameters())

    # Learning rate scheduler
    first_step = 150
    second_step = 225
    iters_one_epoch = 390
    learning_rate_scheduler = StepScheduler(
        init_lr=0.1, gamma=0.1, iter_steps=[iters_one_epoch*first_step, iters_one_epoch*second_step])

    start_point = 0

    # Create monitor
    monitor = Monitor('tmp.monitor')
    monitor_loss = MonitorSeries("Training loss", monitor, interval=1)
    monitor_err = MonitorSeries("Training error", monitor, interval=1)
    monitor_time = MonitorTimeElapsed("Training time", monitor, interval=1)
    monitor_verr = MonitorSeries("Test error", monitor, interval=1)
    monitor_vloss = MonitorSeries("Test loss", monitor, interval=1)

    # save_nnp
    contents = save_nnp({"x": image_valid}, {"y": pred_valid}, bs_valid)
    save.save(
        os.path.join(model_save_path,
                     (model+"_epoch0_result.nnp")), contents
    )

    train_iter = math.ceil(n_train_samples / bs_train)
    val_iter = math.ceil(n_val_samples / bs_valid)

    min_loss = np.inf

    # Training-loop
    for i in tqdm(range(start_point, train_epochs)):
        # Forward/Zerograd/Backward
        print("## Training")
        e = 0.0
        loss = 0.0
        for k in range(train_iter):
            total_iter = train_iter * i + k
            lr = learning_rate_scheduler.get_learning_rate(total_iter)
            solver.set_learning_rate(lr)
            image, label = train_loader.next()
            image_train.d = image
            label_train.d = label
            loss_train.forward(clear_no_need_grad=True)
            solver.zero_grad()
            loss_train.backward(clear_buffer=True)
            solver.weight_decay(decay_rate=0.0001)
            solver.update()
            e += categorical_error(pred_train.d, label_train.d)
            loss += loss_train.data.data.copy() * bs_train
        e /= train_iter
        loss /= n_train_samples

        monitor_loss.add(i, loss)
        monitor_err.add(i, e)
        monitor_time.add(i)

        # Validation
        ve = 0.0
        vloss = 0.0

        print("## Validation")
        for j in range(val_iter):
            image, label = val_loader.next()
            image_valid.d = image
            label_valid.d = label
            loss_val.forward()
            vloss += loss_val.data.data.copy() * bs_valid
            ve += categorical_error(pred_valid.d, label)
        ve /= val_iter
        vloss /= n_val_samples
        monitor_verr.add(i, ve)
        monitor_vloss.add(i, vloss)

        if min_loss > vloss:
            nn.save_parameters(os.path.join(
                model_save_path, "{}_best_params.h5".format(model)))
            min_loss = vloss

        if i == 5 or i % 10 == 0:
            contents = save_nnp({"x": image_valid}, {
                                "y": pred_valid}, bs_valid)
            save.save(os.path.join(model_save_path,
                                   (model+"_result_{}.nnp".format(i))), contents)

    # save_nnp_lastepoch
    contents = save_nnp({"x": image_valid}, {"y": pred_valid}, bs_valid)
    save.save(os.path.join(model_save_path,
              (model+"_result_{}.nnp".format(train_epochs))), contents)


# responsible_ai/data_cleansing/eval_influence_functions.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import os
import argparse
import shutil
import pandas as pd
import nnabla as nn
from sgd_influence.utils import ensure_dir, get_context
from sgd_influence.eval import run_train_for_eval
from sgd_influence.model import get_config
from sgd_influence.train import train
from influence_functions.infl import infl_icml
from sgd_influence.analyze import analyze
from datasets.create_cifar10_csv import create_data_csv as create_cifar10_csv
from datasets.create_stl10_csv import create_data_csv as create_stl10_csv
from datasets.create_mnist_csv import create_data_csv as create_mnist_csv


method_dict = {
    # infl_Modified_retrain_last
    'Influence_functions':
        {
            'calc_infl_method': 'last',
            'calc_infl_with_all_params': True,
            'shuffle_infl': False,
        },
    # random drop
    'Random':
        {
            'calc_infl_method': 'last',
            'calc_infl_with_all_params': False,
            'shuffle_infl': True,
        },
}


def has_enough_data(data_csv_dir, num_of_data):
    _df = pd.read_csv(data_csv_dir)
    if num_of_data >= len(_df):
        return False
    return True


def make_random_file(infl_file):
    infl = pd.read_csv(infl_file)
    infl = infl.sample(frac=1)
    infl.to_csv(infl_file, index=False)


def train_and_infl(args, shuffle_infl):
    args.score_output = None
    config = get_config(args)
    model_info_dict = config.model_info_dict
    file_dir_dict = config.file_dir_dict
    need_evaluate = config.need_evaluate
    calc_infl_with_all_params = args.calc_infl_with_all_params
    print(file_dir_dict['infl_filename'])
    if os.path.exists(file_dir_dict['infl_filename']):
        return
    # train
    # only for the round where all params are calculated to avoid duplicated calculation
    if calc_infl_with_all_params:
        train(model_info_dict, file_dir_dict,
              calc_infl_with_all_params, need_evaluate)
    # infl
    if args.dataset == "cifar10":
        alpha = 10.0
    else:
        alpha = 0.1
    if shuffle_infl:
        from_path = os.path.join(
            os.path.dirname(os.path.dirname(args.output)),
            next(iter(method_dict.keys()))
        )
        files = os.listdir(from_path)
        fn = [f for f in files if f'_{str(args.seed)}' in f][0]
        shutil.copy(os.path.join(from_path, fn), args.output)
        make_random_file(args.output)
    else:
        infl_icml(model_info_dict, file_dir_dict,
                  calc_infl_with_all_params, need_evaluate, alpha)


def eval(args, method_score_dir, score_basename):
    for idx, n_to_remove in enumerate(args.remove_n_list):
        args.score_output = os.path.join(
            method_score_dir, score_basename + str(n_to_remove) + '_seed_' + str(args.seed) + '.csv')
        print(args.score_output)
        if os.path.exists(args.score_output):
            print('already exists')
            continue
        config = get_config(args, is_eval=True)
        model_info_dict = config.model_info_dict
        file_dir_dict = config.file_dir_dict
        if not has_enough_data(file_dir_dict['train_csv'], n_to_remove):
            continue
        run_train_for_eval(model_info_dict, file_dir_dict,
                           n_to_remove, args.retrain_all)


def create_csv_files(dataset_name, seed: int):
    func_dict = {
        'cifar10': create_cifar10_csv,
        'stl10': create_stl10_csv,
        'mnist': create_mnist_csv,
    }
    func_dict[dataset_name](seed)


def generate_dataset(dataset_name, seed):
    fd = {}
    base_dir = os.path.join(os.path.abspath(
        os.path.dirname(__file__)), 'datasets', dataset_name)
    fd['training'] = os.path.join(
        base_dir, f'{dataset_name}_training_{str(seed)}.csv')
    fd['validation'] = os.path.join(
        base_dir, f'{dataset_name}_validation_{str(seed)}.csv')
    fd['test'] = os.path.join(base_dir, f'{dataset_name}_test.csv')
    if os.path.exists(fd['training']) & os.path.exists(fd['validation']):
        pass
    else:
        create_csv_files(dataset_name, seed)
    return fd


def run(args):
    os.environ['NNABLA_CUDNN_DETERMINISTIC'] = '1'
    os.environ['NNABLA_CUDNN_ALGORITHM_BY_HEURISTIC'] = '1'
    # gpu/cpu
    ctx = get_context(args.device_id)
    nn.set_default_context(ctx)

    seeds = [i for i in range(args.n_trials)]
    weight_dir = 'infl_results_sgd'
    score_basename = 'score_delete_worst_'
    output_dir = args.output_dir

    for i, seed in enumerate(seeds):
        args.seed = seed
        fd = generate_dataset(args.dataset, seed)
        args.input_train, args.input_val, args.input_test = fd[
            'training'], fd['validation'], fd['test']
        for method_name, calc_setting in method_dict.items():
            # setup params
            method_score_dir = os.path.join(output_dir, 'score', method_name)
            method_infl_dir = os.path.join(
                output_dir, 'influence', method_name)
            ensure_dir(method_score_dir)
            ensure_dir(method_infl_dir)
            args.calc_infl_method = calc_setting['calc_infl_method']
            args.calc_infl_with_all_params = calc_setting['calc_infl_with_all_params']
            shuffle_infl = calc_setting['shuffle_infl']
            args.weight_output = os.path.join(weight_dir, 'seed_%02d' % (seed))
            ensure_dir(args.weight_output)
            args.output = os.path.join(
                method_infl_dir, 'influence_' + method_name.lower() + '_' + str(seed) + '.csv')
            # train and infl
            train_and_infl(args, shuffle_infl)
            # eval
            eval(args, method_score_dir, score_basename)

    # show result
    analyze(output_dir)


def main():
    parser = argparse.ArgumentParser(
        description='check performance of SGD-influence', formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument(
        '-o', '--output_dir', help='path to output dir', default='output_influence_functions')
    parser.add_argument(
        '-e', '--n_epochs', help="epoch for SGD influence calculation(int) default=20", default=20, type=int)
    parser.add_argument(
        '-nt', '--n_trials', help='number or trials ', default=6, type=int)
    parser.add_argument(
        '-r', '--remove_n_list', help="list of n of samples to remove. ex: '-r 10 20' makes [10, 20]",
        type=int, nargs='+', default=[0, 1, 10, 40, 100, 200, 600, 1000, 3000, 5000, 10000, 30000])
    parser.add_argument(
        '-ds', '--dataset', choices=['cifar10', 'stl10', 'mnist'], default='cifar10')
    parser.add_argument(
        '-di', '--device-id', help='device id of gpu', default=0)
    parser.add_argument(
        '-ra', '--retrain_all', help='if True, retrain all.', action='store_true')
    parser.set_defaults(func=run)
    args = parser.parse_args()
    args.func(args)


if __name__ == '__main__':
    main()


# responsible_ai/data_cleansing/eval_sgd_influence.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import argparse
import shutil
import pandas as pd
import nnabla as nn
from sgd_influence.utils import ensure_dir, get_context
from sgd_influence.eval import run_train_for_eval
from sgd_influence.model import get_config
from sgd_influence.train import train
from sgd_influence.infl import infl_sgd
from sgd_influence.analyze import analyze
from datasets.create_cifar10_csv import create_data_csv as create_cifar10_csv
from datasets.create_stl10_csv import create_data_csv as create_stl10_csv
from datasets.create_mnist_csv import create_data_csv as create_mnist_csv


method_dict = {
    # sgd_Hara_retrain_last
    'Hara':
        {
            'calc_infl_method': 'last',
            'calc_infl_with_all_params': True,
            'shuffle_infl': False,
        },
    # sgd_Modified_retrain_last
    'Ours':
        {
            'calc_infl_method': 'last',
            'calc_infl_with_all_params': False,
            'shuffle_infl': False,
        },
    # random drop
    'Random':
        {
            'calc_infl_method': 'last',
            'calc_infl_with_all_params': False,
            'shuffle_infl': True,
        },
}


def has_enough_data(data_csv_dir, num_of_data):
    _df = pd.read_csv(data_csv_dir)
    if num_of_data >= len(_df):
        print(
            f'n_to_remove ({num_of_data}) >= total dataset size ({len(_df)})')
        return False
    return True


def make_random_file(infl_file):
    infl = pd.read_csv(infl_file)
    infl = infl.sample(frac=1)
    infl.to_csv(infl_file, index=False)


def train_and_infl(args, shuffle_infl):
    args.score_output = None
    config = get_config(args)
    model_info_dict = config.model_info_dict
    file_dir_dict = config.file_dir_dict
    need_evaluate = config.need_evaluate
    calc_infl_with_all_params = args.calc_infl_with_all_params
    print(file_dir_dict['infl_filename'])
    if os.path.exists(file_dir_dict['infl_filename']):
        return
    # train
    # only for the round where all params are calculated to avoid duplicated calculation
    if calc_infl_with_all_params:
        train(model_info_dict, file_dir_dict,
              calc_infl_with_all_params, need_evaluate)
    # infl
    if shuffle_infl:
        from_path = os.path.join(
            os.path.dirname(os.path.dirname(args.output)),
            next(iter(method_dict.keys()))
        )
        files = os.listdir(from_path)
        fn = [f for f in files if f'_{str(args.seed)}' in f][0]
        shutil.copy(os.path.join(from_path, fn), args.output)
        make_random_file(args.output)
    else:
        infl_sgd(model_info_dict, file_dir_dict,
                 calc_infl_with_all_params, need_evaluate)


def eval(args, method_score_dir, score_basename):
    for idx, n_to_remove in enumerate(args.remove_n_list):
        args.score_output = os.path.join(
            method_score_dir, score_basename + str(n_to_remove) + '_seed_' + str(args.seed) + '.csv')
        print(args.score_output)
        if os.path.exists(args.score_output):
            print('already exists')
            continue
        config = get_config(args, is_eval=True)
        model_info_dict = config.model_info_dict
        file_dir_dict = config.file_dir_dict
        if not has_enough_data(file_dir_dict['train_csv'], n_to_remove):
            continue
        run_train_for_eval(model_info_dict, file_dir_dict,
                           n_to_remove, args.retrain_all)


def create_csv_files(dataset_name, seed: int):
    func_dict = {
        'cifar10': create_cifar10_csv,
        'stl10': create_stl10_csv,
        'mnist': create_mnist_csv,
    }
    func_dict[dataset_name](seed)


def generate_dataset(dataset_name, seed):
    fd = {}
    base_dir = os.path.join(os.path.abspath(
        os.path.dirname(__file__)), 'datasets', dataset_name)
    fd['training'] = os.path.join(
        base_dir, f'{dataset_name}_training_{str(seed)}.csv')
    fd['validation'] = os.path.join(
        base_dir, f'{dataset_name}_validation_{str(seed)}.csv')
    fd['test'] = os.path.join(base_dir, f'{dataset_name}_test.csv')
    if os.path.exists(fd['training']) & os.path.exists(fd['validation']):
        pass
    else:
        create_csv_files(dataset_name, seed)
    return fd


def run(args):
    os.environ['NNABLA_CUDNN_DETERMINISTIC'] = '1'
    os.environ['NNABLA_CUDNN_ALGORITHM_BY_HEURISTIC'] = '1'
    # gpu/cpu
    ctx = get_context(args.device_id)
    nn.set_default_context(ctx)

    seeds = [i for i in range(args.n_trials)]
    weight_dir = 'infl_results_sgd'
    score_basename = 'score_delete_worst_'
    output_dir = args.output_dir

    for i, seed in enumerate(seeds):
        args.seed = seed
        fd = generate_dataset(args.dataset, seed)
        args.input_train, args.input_val, args.input_test = fd[
            'training'], fd['validation'], fd['test']
        for method_name, calc_setting in method_dict.items():
            # setup params
            method_score_dir = os.path.join(output_dir, 'score', method_name)
            method_infl_dir = os.path.join(
                output_dir, 'influence', method_name)
            ensure_dir(method_score_dir)
            ensure_dir(method_infl_dir)
            args.calc_infl_method = calc_setting['calc_infl_method']
            args.calc_infl_with_all_params = calc_setting['calc_infl_with_all_params']
            shuffle_infl = calc_setting['shuffle_infl']
            args.weight_output = os.path.join(weight_dir, 'seed_%02d' % (seed))
            ensure_dir(args.weight_output)
            args.output = os.path.join(
                method_infl_dir, 'influence_' + method_name.lower() + '_' + str(seed) + '.csv')
            # train and infl
            train_and_infl(args, shuffle_infl)
            # eval
            eval(args, method_score_dir, score_basename)

    # show result
    analyze(output_dir)


def main():
    parser = argparse.ArgumentParser(
        description='check performance of SGD-influence', formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument(
        '-o', '--output_dir', help='path to output dir', default='output_sgd_influence')
    parser.add_argument(
        '-e', '--n_epochs', help="epoch for SGD influence calculation(int) default=20", default=20, type=int)
    parser.add_argument(
        '-nt', '--n_trials', help='number or trials ', default=6, type=int)
    parser.add_argument(
        '-r', '--remove_n_list', help="list of n of samples to remove. ex: '-r 10 20' makes [10, 20]",
        type=int, nargs='+', default=[0, 1, 10, 40, 100, 200, 600, 1000, 3000, 5000, 10000, 30000])
    parser.add_argument(
        '-ds', '--dataset', choices=['cifar10', 'stl10', 'mnist'], default='cifar10')
    parser.add_argument(
        '-di', '--device-id', help='device id of gpu', default=0)
    parser.add_argument(
        '-ra', '--retrain_all', help='if True, retrain all.', action='store_true')
    parser.set_defaults(func=run)
    args = parser.parse_args()
    args.func(args)


if __name__ == '__main__':
    main()


# responsible_ai/eigencam/utils.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import cv2
import numpy as np
from PIL import Image

import base64
from io import BytesIO


from IPython.display import display, Javascript
from google.colab.output import eval_js


def overlay_images(base_img, overlay_img, overlay_coef=1.0):
    """
    Overlay two images.
    Parameters
    ----------
    base_img: ndarray
        2D array
    overlay_img: ndarray
        2D array
    overlay_coef: float
        mix rate of overlay_img to base_img (overlay_img: overlay_coef, base_img: 1). 
    Returns
    ----------
    ret_img: ndarray
        2D array of overlaid image
    """
    # resize
    _overlay_img = cv2.resize(
        overlay_img, (base_img.shape[1], base_img.shape[0]))
    # normalize
    _overlay_img = 255 * _overlay_img / _overlay_img.max()
    _overlay_img = _overlay_img.astype('uint8')
    # color adjust
    _overlay_img = cv2.applyColorMap(_overlay_img, cv2.COLORMAP_JET)
    base_img = cv2.cvtColor(base_img, cv2.COLOR_BGR2RGB)
    # overlay
    ret_img = _overlay_img * overlay_coef + base_img
    ret_img = 255 * ret_img / ret_img.max()
    ret_img = ret_img.astype('uint8')
    ret_img = cv2.cvtColor(ret_img, cv2.COLOR_BGR2RGB)
    return ret_img


def resize_image_for_yolo(img_orig, w=608, h=608):

    im_h, im_w, _ = img_orig.shape

    if (w * 1.0 / im_w) < (h * 1. / im_h):
        new_w = w
        new_h = int((im_h * w) / im_w)
    else:
        new_h = h
        new_w = int((im_w * h) / im_h)

    patch = cv2.resize(img_orig, (new_w, new_h)) / 255.
    img = np.ones((h, w, 3), np.float32) * 0.5
    # resize
    x0 = int((w - new_w) / 2)
    y0 = int((h - new_h) / 2)
    img[y0:y0 + new_h, x0:x0 + new_w] = patch

    return img, new_w, new_h


def decode_img_str(img_str):
    decimg = base64.b64decode(img_str.split(',')[1], validate=True)
    decimg = Image.open(BytesIO(decimg))
    decimg = np.array(decimg, dtype=np.uint8)
    decimg = cv2.cvtColor(decimg, cv2.COLOR_BGR2RGB)
    return decimg


def encode_img(img):
    _, encimg = cv2.imencode(".jpg", img, [int(cv2.IMWRITE_JPEG_QUALITY), 80])
    img_str = encimg.tostring()
    img_str = "data:image/jpeg;base64," + \
        base64.b64encode(img_str).decode('utf-8')
    return img_str


# responsible_ai/face_evaluation/face_evaluation.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import cv2


def extract_non_black(img):
    img_total = img[:, :, 0] + img[:, :, 1] + img[:, :, 2]
    non_black = np.where(img_total > 0)
    extract_img = []
    for i in range(len(non_black[0])):
        extract_img.append(img[non_black[0][i], non_black[1][i], :])
    extract_img = np.array(extract_img).reshape(1, -1, 3)
    return extract_img


def calc_ita(masked_imgs):
    itas = []
    for masked_img in masked_imgs:
        lab_img = cv2.cvtColor(masked_img, cv2.COLOR_RGB2Lab)
        l_img, a_img, b_img = cv2.split(lab_img)
        l_img = l_img * (100/255)
        b_img = b_img - 128
        ITA_img = (np.arctan((l_img-50)/b_img)*180)/np.pi
        a_hist, a_bins = np.histogram(ITA_img, bins=1000)
        ITA_each = a_bins[a_hist.argmax()]
        itas.append(ITA_each)
    return round(sum(itas)/len(itas), 2)


# responsible_ai/fairness_metrics/utils.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pandas as pd
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
import matplotlib.pyplot as plt
from IPython.display import Markdown, display
import opendatasets as od
from urllib.request import urlopen
from io import BytesIO
from zipfile import ZipFile
import sys


def display_log(log_dict):
    """
        Display and analyze the logs
    """

    confusion_matrix_log = \
        "Confusion matrices above show how the model performs on some test data. " \
        "We also print additional information (calculated from confusion matrices) to assess " \
        "the model fairness. For instance,\n" \
        "* The model predicted and gave favourable outcomes to " + str(log_dict['total_approvals']) + "" \
                                                                                                      " people. Among these individuals, " + str(
            log_dict['privileged_approvals']) + " belonged to " \
                                                "the privileged group and " + str(
            log_dict['unprivileged_approvals']) + " belonged to the unprivileged group.\n" \
                                                  "* Positive rate (PR) for privileged group is " + str(
            log_dict['privileged_PR']) + ", " \
                                         "and the PR for unprivileged group is " + str(
            log_dict['unprivileged_PR']) + " . " \
                                           "PR = TP /(TP+FP+TN+FN).\n" \
                                           "* True positive rate (TPR) for the privileged group is " + str(
            log_dict['privileged_TPR']) + ", " \
                                          "and TPR for the  unprivileged group is " + str(
            log_dict['unprivileged_TPR']) + " . " \
                                            "For the privileged group, the TPR = TP /(TP+FN).\n" \
                                            "* False positive rate (FPR) for the privileged group is " + str(
            log_dict['privileged_FPR']) + "," \
                                          " and the FPR for the unprivileged group is " + str(
            log_dict['unprivileged_FPR']) + " ." \
                                            "For the privileged group, TPR = FP /(FP+TN).\n\n" \
                                            "Let's check model fairness. In this tutorial, we explore three different types of metrics.\n"
    display(Markdown(confusion_matrix_log))

    # demographic parity
    demographic_parity_header = \
        "##### Demographic parity: \n" \
        "* As discussed in the Demographic parity section, model is fair" \
        " if all segments of the protected class receive equal proportion/rate of positive outcomes." \
        " We refer to this proportion as `Positive Rate` of the model."
    display(Markdown(demographic_parity_header))
    plot_fairness([log_dict['DPD']], "demographic parity")
    demographic_parity_log = \
        "* As seen in the above diagram, difference between positive rates of " \
        "unprivileged and privileged group is " + str(log_dict['DPD']) + " , " \
                                                                         "which means the " + str(
            ["privileged group gets " + str(round(abs(log_dict['DPD'] * 100), 3)) + ""
                                                                                    "% more positive outcome" if
             log_dict['DPD'] < 0 else "unprivileged group"
                                      " gets " + str(round(abs(log_dict['DPD'] * 100), 3)) + "% more positive outcome"][
                0]) + ".\n" \
                      "* Demographic parity difference (DPD) must be close to zero for the classifier to be fair.\n" \
                      "* Model fairness is " + str(log_dict['DPD']) + " , which means " \
        + str(["the model treats individuals from the privileged and"
               "unprivileged group fairly, and satisfies the demographic parity fairness criterion."
               if log_dict['DPD'] == 0 else
               "the model treats individuals from the privileged and"
               "unprivileged group with an acceptable bias. To further improve model fairness bias must be mitigated."
               if (-0.10 < log_dict['DPD'] < 0.10) else
               "the model is not satisfied demographic parity fairness criterion ."
               " To improve model fairness, we need to mitigate(minimize) the bias."][0])
    display(Markdown(demographic_parity_log))

    # disparate impact
    disparate_impact_header = \
        "##### Disparate Impact: \n" \
        "* As discussed in the Disparate Impact section, model is fair" \
        " if all segments of the protected class receive equal proportion/rate of positive outcomes." \
        " We refer to this proportion as `Positive Rate` of the model."
    display(Markdown(demographic_parity_header))
    plot_disparate_impact(log_dict['DI'], "Disparate Impact")
    disparate_impact_log = \
        "* As seen in the above diagram , the ratio of rate of favorable outcome for the" \
        " unprivileged group to that of the privileged group is " + str(log_dict['DI']) + "\n" \
                                                                                          "* Disparate Impact (DI) must be close to One for the classifier to be fair.\n" \
                                                                                          "* Model fairness is " + str(
            log_dict['DI']) + ", which means " \
        + str(["the model treats individuals from the privileged and "
               "unprivileged group fairly, and satisfies the disparate impact fairness criterion."
               if (log_dict['DI'] == 1) else
               "the model is not satisfied Disparate Impact fairness criterion ."
               " To improve model fairness , we need to mitigate (minimize) the bias."][0])

    display(Markdown(disparate_impact_log))

    # equal opportunity
    equal_opportunity_header = \
        "##### Equal opportunity : \n" \
        "* As discussed in the Equal opportunity section, model is fair, " \
        "if the proportion of people who should be selected by the model (`positives`) that are correctly " \
        "selected by the model is the same for each group. We refer to this proportion as the " \
        "true positive rate (TPR) of the model.\n"
    display(Markdown(equal_opportunity_header))
    plot_fairness([log_dict['EOD']], "Equal opportunity difference")
    equal_opportunity_log = \
        "* As seen in the above diagram the difference between positive rates of the " \
        "unprivileged and privileged group is " + str(log_dict['EOD']) + ", " \
                                                                         "which means the " + str(
            ["privileged group gets " + str(abs(log_dict['EOD'] * 100)) +
             "% more true positive rate. In other words, for privileged group,"
             "almost all people who should be approved are actually approved. "
             "For unprivileged group, if you should be approved, your chances of actually being "
             "approved are much lower" if log_dict['EOD'] < 0 else "unprivileged group gets"
                                                                   + str(abs(log_dict['EOD'] * 100)) +
                                                                   "% more true positive rate. In other words, for unprivileged group, "
                                                                   "almost all people who should be approved are actually approved. "
                                                                   "For privileged group, if you should be approved, your chances of actually "
                                                                   "being approved are much lower"][0]) + ".\n" \
                                                                                                          "* Equal opportunity difference (EOD) must be close to zero for the classifier to be fair.\n" \
                                                                                                          "* Model fairness is " + str(
            log_dict['EOD']) + ". which means " \
        + str(["the model treats individuals from the privileged and "
               "unprivileged group equally and satisfies the EOD."
               if (log_dict['EOD'] == 0) else
               "the model treats individuals from the privileged and"
               "unprivileged group with an acceptable bias. To further improve model fairness bias must be mitigated."
               if (-0.10 < log_dict['EOD'] < 0.10) else ""
                                                        "the model is not satisfied Equal opportunity fairness criterion ."
                                                        " To improve model fairness , we need to mitigate (minimize) the bias."][
                  0])
    display(Markdown(equal_opportunity_log))

    # equalized odds
    equalized_odds_header = \
        "##### Equalized odds : \n" \
        "* As discussed in the equalized section , model is fair, if the proportion of people correctly identify " \
        "the positive outcome at equal rates across groups , but also miss-classify the positive outcome " \
        "at equal rates across groups . we refer to this TPR and FPR of the model.\n"
    display(Markdown(equalized_odds_header))
    plot_fairness([log_dict['AAOD']], "Average abs odd difference")
    equalized_odds_log = \
        "* As seen in the above diagram , the average abs odd difference between privileged group" \
        " and unprivileged group is " + str(log_dict['AAOD']) + "\n" \
                                                                "* Average abs odds difference (AAOD) must be close to zero for the classifier to be fair.\n" \
                                                                "* Model fairness is " + str(
            log_dict['AAOD']) + ", which means " \
        + str(["the model is correctly identify the positive outcome, "
               "and miss-classify the positive outcome at equal rates across groups . "
               "and it satisfies the equalized odd fairness criterion."
               if (log_dict['AAOD'] == 0) else
               "the model is not satisfied Equalized odd fairness criterion ."
               " To improve model fairness , we need to mitigate (minimize) the bias."][0])

    display(Markdown(equalized_odds_log))


def load_german_data():
    """
    Load and preprocess german credit card dataset and returns the feature, target, and sensitive attribute.

    This dataset classifies people described by a set of attributes as good or bad credit risks.
    (For more info : https://archive.ics.uci.edu/ml/datasets/Statlog+%28German+Credit+Data%29)
    Returns:
        features (pandas.DataFrame): features of the german credit card dataset
        target (pandas.Series): target values of the german credit card dataset
        sensitive_attributes(pandas.DataFrame): sensitive attributes values of german credit card dataset
    """
    file_path = r'https://archive.ics.uci.edu/ml/machine-learning-databases/statlog/german/german.data'
    column_names = ['status', 'month', 'credit_history',
                    'purpose', 'credit_amount', 'savings', 'employment',
                    'investment_as_income_percentage', 'personal_status',
                    'other_debtors', 'residence_since', 'property', 'age',
                    'installment_plans', 'housing', 'number_of_credits',
                    'skill_level', 'people_liable_for', 'telephone',
                    'foreign_worker', 'credit']
    try:
        dataframe = pd.read_csv(
            file_path, sep=' ', header=None, names=column_names)
    except Exception as e:
        print("Error : ", e)
        sys.exit(1)
    # derive the gender attribute from personal_status (you can refer the german.doc)
    status_map = {'A91': 'male', 'A93': 'male', 'A94': 'male',
                  'A92': 'female', 'A95': 'female'}
    dataframe['gender'] = dataframe['personal_status'].replace(status_map)

    # sensitive attributes; we identify 'age' and 'gender' as sensitive attributes
    # privileged class for age and gender is male and age >=26 respectively.
    sensitive_attribs = ['age', 'gender']
    sensitive_attributes = (dataframe.loc[:, sensitive_attribs]
                            .assign(age=lambda df: (df['age'] >= 26).astype(int),
                                    gender=lambda df: (df['gender'] == 'male').astype(int)))

    # targets; 1 when someone makes good credit , otherwise 0
    target = (dataframe['credit'] == 1).astype(int)
    # features; note that the 'target' and sensitive attribute columns are dropped
    features = (dataframe
                .drop(columns=['credit', 'age', 'gender', 'personal_status'])
                .fillna('Unknown')
                .pipe(pd.get_dummies, drop_first=True))

    display(Markdown(
        f"features : {features.shape[0]} samples, {features.shape[1]} attributes"))
    display(Markdown(f"targets : {target.shape[0]} samples"))
    display(Markdown(
        f"sensitives attributes : {sensitive_attributes.shape[0]} samples, {sensitive_attributes.shape[1]} attributes"))

    return features, target, sensitive_attributes


def load_adult_data():
    """
    Load and preprocess Census Income dataset and returns the feature, target, and sensitive attribute.

    This dataset is used to predict whether income exceeds $50K/yr
    (For more info : https://archive.ics.uci.edu/ml/datasets/adult)

    Returns:
        features (pandas.DataFrame): features of the Census income dataset
        target (pandas.Series): target values of the Census income dataset
        sensitive_attributes(pandas.DataFrame): sensitive attributes values of Census income dataset
    """
    file_path = r'https://archive.ics.uci.edu/ml/machine-learning-databases/adult/adult.data'
    column_names = ['age', 'workclass', 'fnlwgt', 'education', 'education_num',
                    'marital_status', 'occupation', 'relationship', 'race', 'sex',
                    'capital_gain', 'capital_loss', 'hours_per_week', 'country', 'target']
    try:
        # Restrict races to White and Black
        input_data = (pd.read_csv(file_path, names=column_names,
                                  na_values="?", sep=r'\s*,\s*', engine='python').loc[
            lambda df: df['race'].isin(['White', 'Black'])]).dropna()
    except Exception as e:
        print("Error : ", e)
        sys.exit(1)

    # sensitive attributes; we identify 'race' and 'sex' as sensitive attributes
    # privileged class for race and sex is White and Male respectively
    sensitive_attribs = ['race', 'sex']
    sensitive_attributes = (input_data.loc[:, sensitive_attribs]
                            .assign(race=lambda df: (df['race'] == 'White').astype(int),
                                    sex=lambda df: (df['sex'] == 'Male').astype(int)))

    # targets; 1 when someone makes over 50k , otherwise 0
    target = (input_data['target'] == '>50K').astype(int)

    # features; note that the 'target' and sensitive attribute columns are dropped
    features = (input_data
                .drop(columns=['target', 'race', 'sex'])
                .pipe(pd.get_dummies, drop_first=True))

    display(Markdown(
        f"features : {features.shape[0]} samples, {features.shape[1]} attributes"))
    display(Markdown(f"targets : {target.shape[0]} samples"))
    display(Markdown(
        f"sensitives attributes : {sensitive_attributes.shape[0]} samples, {sensitive_attributes.shape[1]} attributes"))
    return features, target, sensitive_attributes


def load_compas_data():
    """
    Load and preprocess ProPublica Recidivism/COMPAS dataset and returns the feature, target, and sensitive attribute.
    This dataset is used to assess the likelihood that a criminal defendant will re-offend.
    (For more info : https://github.com/propublica/compas-analysis)

    Returns:
        features (pandas.DataFrame): features of the COMPAS dataset
        target (pandas.Series): target values of the COMPAS dataset
        sensitive_attributes(pandas.DataFrame): sensitive attributes values of COMPAS dataset
    """
    file_path = r'https://raw.githubusercontent.com/propublica/compas-analysis/master/compas-scores-two-years.csv'
    try:
        # Restrict races to African-American and Caucasian
        input_data = (pd.read_csv(file_path, index_col='id').loc[lambda df: df['race'].isin(
            ['African-American', 'Caucasian'])])
    except Exception as e:
        print("Error : ", e)
        sys.exit(1)

    input_data = input_data[['sex', 'age', 'age_cat', 'race',
                             'juv_fel_count', 'juv_misd_count', 'juv_other_count',
                             'priors_count', 'c_charge_degree', 'c_charge_desc',
                             'two_year_recid']].dropna()
    # sensitive attributes; we identify 'race' and 'sex' as sensitive attributes
    # privileged class for race and sex is Caucasian and Female respectively.
    sensitive_attribs = ['race', 'sex']
    sensitive_attributes = (input_data.loc[:, sensitive_attribs]
                            .assign(race=lambda df: (df['race'] == 'Caucasian').astype(int),
                                    sex=lambda df: (df['sex'] == 'Female').astype(int)))
    # in the dataset the label value 0 is considered for no-recidivism and 1 is considered for did-recidivism
    # for shake of understanding we will map the label value 0 for did-recidivism and 1 for no-recidivism
    # targets; 1 when someone makes has no-recidivism, otherwise 0
    target = (input_data['two_year_recid'] == 0).astype(int)

    # features; note that the 'target' and sensitive attribute columns are dropped
    features = (input_data
                .drop(columns=['two_year_recid', 'race', 'sex'])
                .pipe(pd.get_dummies, drop_first=True))
    display(Markdown(
        f"features : {features.shape[0]} samples, {features.shape[1]} attributes"))
    display(Markdown(f"targets : {target.shape[0]} samples"))
    display(Markdown(
        f"sensitives attributes : {sensitive_attributes.shape[0]} samples, {sensitive_attributes.shape[1]} attributes"))

    return features, target, sensitive_attributes


def load_bank_data():
    """
    Load and preprocess Bank marketing dataset and returns the feature, target, and sensitive attribute.
    The data is related with direct marketing campaigns (phone calls) of a Portuguese banking institution.
    The classification goal is to predict if the client will subscribe a term deposit (variable y).
    (for more info : https://archive.ics.uci.edu/ml/datasets/Bank+Marketing)

    Returns:
        features (pandas.DataFrame): features of the Bank marketing dataset
        target (pandas.Series): target values of the Bank marketing dataset
        sensitive_attributes(pandas.DataFrame): sensitive attributes values of Bank marketing dataset
    """

    try:
        http_response = urlopen(
            r'https://archive.ics.uci.edu/ml/machine-learning-databases/00222/bank-additional.zip')
        zipfile = ZipFile(BytesIO(http_response.read()))
        zipfile.extractall(path='.')
        file_path = "./bank-additional/bank-additional-full.csv"
        input_data = pd.read_csv(
            file_path, sep=';', na_values="unknown").dropna()
    except Exception as e:
        print("Error : ", e)
        sys.exit(1)

    # sensitive attributes; we identify "age" as sensitive attributes
    # privileged class age >=25.
    sensitive_attribs = ['age']
    sensitive_attributes = (input_data.loc[:, sensitive_attribs]
                            .assign(age=lambda df: (df['age'] >= 25).astype(int)))
    # targets; 1 if the client will subscribe a term deposit, otherwise 0
    target = (input_data['y'] == 'yes').astype(int)

    # features; note that the 'target' and sensitive attribute columns are dropped
    features = (input_data
                .drop(columns=['y', 'age'])
                .pipe(pd.get_dummies, drop_first=True))
    display(Markdown(
        f"features : {features.shape[0]} samples, {features.shape[1]} attributes"))
    display(Markdown(f"targets : {target.shape[0]} samples"))
    display(Markdown(
        f"sensitives attributes : {sensitive_attributes.shape[0]} samples, {sensitive_attributes.shape[1]} attributes"))

    return features, target, sensitive_attributes


def load_gender_discrimination_data():
    """
    Load Gender discrimination dataset and returns the feature, target, and sensitive attribute.

    Gender discrimination influence exists or not on payments and promotion,
    (for more info : https://www.kaggle.com/hjmjerry/gender-discrimination). To download this dataset
    please provide your Kaggle credentials.

    Returns:
        features (pandas.DataFrame): features of the Bank marketing dataset
        target (pandas.Series): target values of the Bank marketing dataset
        sensitive_attributes(pandas.DataFrame): sensitive attributes values of Bank marketing dataset
    """
    try:
        dataset_url = 'https://www.kaggle.com/hjmjerry/gender-discrimination'
        od.download(dataset_url)
        input_data = pd.read_csv(r'./gender-discrimination/Lawsuit.csv')
    except Exception as e:
        print("Error : ", e)
        sys.exit(1)
    # sensitive attributes; we identify "gender" as sensitive attributes
    # privileged class male.
    sensitive_attribs = ['Gender']
    sensitive_attributes = (input_data.loc[:, sensitive_attribs]
                            .assign(Gender=lambda df: (df['Gender'] == 1).astype(int)))
    # targets; Rank 1 if the faculty will get promote to full professor , otherwise 0
    target = (input_data['Rank'] == 3).astype(int)

    # features; note that the 'target' and sensitive attribute columns are dropped
    features = (input_data
                .drop(columns=['Rank', 'Gender'])
                .pipe(pd.get_dummies, drop_first=True))
    display(Markdown(
        f"features : {features.shape[0]} samples, {features.shape[1]} attributes"))
    display(Markdown(f"targets : {target.shape[0]} samples"))
    display(Markdown(
        f"sensitives attributes : {sensitive_attributes.shape[0]} samples, {sensitive_attributes.shape[1]} attributes"))
    return features, target, sensitive_attributes


def plot_confusion_matrix(y_true, y_pred, display_labels=["bad", "good"]):
    """
    Load Bank marketing dataset and returns the feature, target, and sensitive attribute

    Args:
        y_true: Ground truth (correct) target values.
        y_pred: Estimated targets as returned by a classifier.
        display_labels: Display labels for plot.

    Returns:
        cm: confusion matrix,
        disp: ConfusionMatrixDisplay
    """

    cm = confusion_matrix(y_true, y_pred)
    disp = ConfusionMatrixDisplay(
        confusion_matrix=cm, display_labels=display_labels)
    return cm, disp.plot(include_values=True, cmap=plt.cm.gist_ncar, xticks_rotation='horizontal', values_format='')


def plot_disparate_impact(fairness, metric="None"):
    """
    graphical visualization of disparate impact fairness

    Args:
        fairness: fairness of the ML model.
        metric: name of the fairness metric.

    """
    plt.figure(facecolor='#FFFFFF', figsize=(4, 4))
    plt.axhline(y=1.0, color='r', linestyle='-')
    plt.bar(["Original"], fairness, color="blue", width=2)
    plt.ylabel(metric)
    plt.title("Fairness of the model", fontsize=15)
    plt.text(0, fairness + 0.01, str(round(fairness, 3)), fontweight='bold',
             color='red', bbox=dict(facecolor='red', alpha=0.4))
    plt.show()


def plot_fairness(fairness, metric="None"):
    """
    graphical visualization of fairness of the model

    Args:
        fairness: fairness of the ML model.
        metric: name of the fairness metric.

    """
    plt.figure(facecolor='#FFFFFF', figsize=(4, 4))
    plt.ylim([-0.6, 0.6])
    plt.axhline(y=0.0, color='r', linestyle='-')
    plt.bar(["Original"], fairness, color="blue", width=2)
    plt.ylabel(metric)
    plt.title("Fairness of the model", fontsize=15)

    for index, value in enumerate(fairness):
        if value < 0:
            plt.text(index, value - 0.1,
                     str(round(value, 3)), fontweight='bold', color='red', bbox=dict(facecolor='red', alpha=0.4))
        else:
            plt.text(index, value + 0.1,
                     str(round(value, 3)), fontweight='bold', color='red', bbox=dict(facecolor='red', alpha=0.4))
    plt.show()


# responsible_ai/gan_data_debiased/args.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse


def get_args():
    """
    parse args
    """
    description = "Dataset Debiasing by GAN"
    parser = argparse.ArgumentParser(description)
    parser.add_argument("-d", "--device-id", type=int, default=0,
                        help="Device ID of the GPU for training")
    parser.add_argument("-c", "--context", type=str, default="cudnn",
                        help="Extension path: cpu or cudnn.")
    parser.add_argument("--type-config", "-t", type=str, default='float',
                        help='Type of computation. e.g. "float", "half".')

    parser.add_argument('--model_train', choices=['baseline', 'gan_debiased'],
                        type=str, default='baseline',
                        help="baseline model or gan_debiased model")
    parser.add_argument("--base-img-path", type=str,
                        default="data/celeba", help="original image path.")
    parser.add_argument('--attribute', type=int, default=1,
                        help="choose the target attribute to train "
                             "(e.g., Arched EyeBrows (1), Bushy Eyebrows(12), smilling (31),etc..")
    parser.add_argument('--protected_attribute', type=int, default=20,
                        help="choose the protected attributes (e.g., Male (20), Pale_Skin (26))")
    # as per author's citation, the baseline model is trained on the CelebA training set
    # with 162,770 images.
    parser.add_argument("--train-beg", type=int, default=0,
                        help="start the training samples")
    parser.add_argument("--valid-beg", type=int, default=162770,
                        help="start the validation sample")
    parser.add_argument("--test-beg", type=int, default=182610,
                        help="start the test sample")
    parser.add_argument("--batch-size", "-b", type=int, default=32,
                        help="Batch size.")
    parser.add_argument("--learning-rate", "-l", type=float, default=1e-4,
                        help="learning rate is a configurable hyper parameter "
                             "used in the training of neural network")
    parser.add_argument('--max_iter_base', type=int, default=5086,
                        help='max iteration to train the baseline model for an epoch')
    parser.add_argument('--max_iter_gan_debiased', type=int, default=15086,
                        help='max iteration to train the gan-debiased model for an epoch')
    parser.add_argument("--model-save-path", "-o",
                        type=str, default='./results',
                        help='Path where the model parameters saved.')
    # gan-based settings
    parser.add_argument('--num_images', type=int, default=175000,
                        help="Number of images to be generated.")
    parser.add_argument("--latent", type=int, default=512,
                        help="Number of latent variables.")
    parser.add_argument('--generate',
                        choices=[
                            'orig',
                            'flip',
                        ], default='orig', help="generate original or flipped images")

    parser.add_argument('--generator_model', type=str,
                        default=r'./result/example_0/Gen_phase_128_epoch_4.h5',
                        help="Model load path used in generation and validation")

    parser.add_argument('--fake_data_dir', type=str, default='data/fake_images',
                        help='Save the generated images by the PGGAN'
                             '(original images :fake_data_dir/AllGenImages'
                             'flipped images :fake_data_dir/"name of the attribute"')
    parser.add_argument('--record_latent_vector', type=str, default='./results/GAN_model',
                        help='Save the latent vectors')

    opt = vars(parser.parse_args())
    return opt


# responsible_ai/gan_data_debiased/classifier.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import pickle
from sklearn.metrics import average_precision_score
import numpy as np
import nnabla as nn
import nnabla.solvers as S
import nnabla.functions as F
from nnabla.models.imagenet import ResNet50
import nnabla.parametric_functions as PF
from utils import utils


def clf_resnet50(layer, n_classes=1, train=True):
    """
    This function uses ResNet-50 pretrained on ImageNet as the base architecture
    and replaces the linear layer in ResNet with two linear layers with the
    hidden layer of size 2,048. Dropout and ReLU are applied
    between these
    """
    layer_1 = F.relu(PF.affine(layer, 2048, name='classifier_1'))
    if train:
        layer_1 = F.dropout(layer_1, 0.5)
    out = PF.affine(layer_1, n_classes, name='classifier_2')
    return out


class attribute_classifier():
    """
    This class uses ResNet-50 pretrained on ImageNet as the base architecture
    and replaces the linear layer in ResNet with two linear layers with the
    hidden layer of size 2,048. Dropout and ReLU are applied
    between these
    """

    def __init__(self, batch_size=32, learning_rate=1e-4, max_iter=5086,
                 total_epochs=20, monitor_path=None,
                 val_weight=None, model_load_path=None):
        """
        Construct all the necessary attributes for the attribute classifier.
        Args:
            batch_size (int): number of samples contained in each generated batch
            learning_rate (float) : learning rate
            max_iter (int) : maximum iterations for an epoch
            total_epochs (int) : total epochs to train the model
            val_weight : sample weights
            monitor_path (str) : model parameter to be saved
            model_load_path (str) : load the model
        """
        self.batch_size = batch_size
        # Resnet 50
        # training graph
        model = ResNet50()
        self.input_image = nn.Variable((self.batch_size,) + model.input_shape)
        self.label = nn.Variable([self.batch_size, 1])
        # fine tuning
        pool = model(self.input_image, training=True, use_up_to='pool')
        self.clf = clf_resnet50(pool)
        self.clf.persistent = True
        # loss
        self.loss = F.mean(F.sigmoid_cross_entropy(self.clf, self.label))
        # hyper parameters
        self.solver = S.Adam(learning_rate)
        self.solver.set_parameters(nn.get_parameters())

        # validation graph
        self.x_v = nn.Variable((self.batch_size,) + model.input_shape)
        pool_v = model(self.x_v, training=False, use_up_to='pool')
        self.v_clf = clf_resnet50(pool_v, train=False)
        self.v_clf_out = F.sigmoid(self.v_clf)
        self.print_freq = 100
        self.validation_weight = val_weight
        # val params
        self.acc = 0.0
        self.total_epochs = total_epochs
        self.max_iter = max_iter
        self.monitor_path = monitor_path

        if model_load_path is not None:
            _ = nn.load_parameters(model_load_path)

    def train(self, train_loader, val_loader):
        """
        train the attribute classifier and save model with best accuracy on validation set
        Args:
             train_loader (NNabla data iterator) : specify training dataset
             val_loader (NNabla data iterator) : specify validation dataset
        """
        model_save = True
        for epoch in range(self.total_epochs):
            for i in range(self.max_iter):
                image, targets = train_loader.next()
                targets = targets[:, 0].reshape(-1, 1)
                self.input_image.d = image
                self.label.d = targets
                self.solver.zero_grad()
                self.loss.forward(clear_no_need_grad=True)
                self.loss.backward(clear_buffer=True)
                self.solver.update()

                if self.print_freq and (i % self.print_freq == 0):
                    print('Training epoch {}: [{}|{}], loss:{}'.format(
                        epoch, i + 1, self.max_iter, self.loss.d), flush=True)
            if model_save:
                acc, val_targets, val_scores = \
                    self.check_avg_precision(
                        val_loader, weights=self.validation_weight)
                if (acc > self.acc):
                    path_best_acc = os.path.join(self.monitor_path, "best")
                    if not os.path.exists(path_best_acc):
                        os.makedirs(path_best_acc)
                    path_best_acc_model = "{}/best_acc.h5".format(
                        path_best_acc)
                    nn.save_parameters(path_best_acc_model)
                    self.acc = acc
                    # computing model fairness
                    cal_thresh = utils.calibrated_threshold(val_targets[:, 0],
                                                            val_scores)
                    _, f1_thresh = utils.get_f1_threshold(val_targets[:, 0],
                                                          val_scores)
                    val_pred = np.where(val_scores > cal_thresh, 1, 0)
                    with open(path_best_acc + '/val_scores.pkl', 'wb+') as handle:
                        pickle.dump(val_scores, handle)
                    with open(path_best_acc + '/val_targets.pkl', 'wb+') as handle:
                        pickle.dump(val_targets, handle)
                    average_precision = utils.get_average_precision(val_targets[:, 0],
                                                                    val_scores)
                    deo = utils.get_difference_equality_opportunity(val_targets[:, 1],
                                                                    val_targets[:, 0], val_pred)
                    bias_amplification = utils.get_bias_amplification(val_targets[:, 1],
                                                                      val_targets[:, 0], val_pred)
                    kl_divergence = utils.get_kl_divergence(val_targets[:, 1],
                                                            val_targets[:, 0], val_scores)
                    val_results = {
                        'AP': average_precision,
                        'DEO': deo,
                        'BA': bias_amplification,
                        'KL': kl_divergence,
                        'f1_thresh': f1_thresh,
                        'cal_thresh': cal_thresh,
                        'epoch': epoch,
                        'accuracy': acc
                    }
                    with open(path_best_acc + '/val_results.pkl', 'wb+') as handle:
                        pickle.dump(val_results, handle)

    def check_avg_precision(self, loader, weights=None, print_out=True):
        """
        Compute average precision (AP) from prediction scores.
        Args:
            loader (NNabla data iterator) : specify the data iterator which to compute the AP
            weights : sample weights
            print_out (bool) : print the logs
        Return:
            acc (float) : accuracy of the model
            y_all (numpy.ndarray) : actual all target labels
            pred_all (numpy.ndarray) : predicted score

        """
        y_all, pred_all = self.get_scores(loader)
        acc = average_precision_score(
            y_all[:, 0], pred_all, sample_weight=weights)
        if print_out:
            print('Avg precision all = {}'.format(acc))
        return acc, y_all, pred_all

    def get_scores(self, loader):
        """
        get predicted scores of the  model
        Args:
            loader : specify the data loader which to predicted the scores
        Returns:
            y_all (numpy.ndarray) : actual all target labels
            pred_all (numpy.ndarray) : all predicted score

        """
        # only for validation and test set
        max_iter = int(loader.size/self.batch_size)
        rem_iter = 0
        if loader.size % self.batch_size != 0:
            rem_iter += 1
            rem_values = loader.size % self.batch_size
        y_all = []
        scores_all = []
        for _ in range(max_iter):
            image, targets = loader.next()
            self.x_v.d = image
            self.v_clf_out.forward(clear_buffer=True)
            y_all.append(targets)
            scores_all.append(self.v_clf_out.d.squeeze())
        for _ in range(rem_iter):
            loader._reset()
            image, targets = loader.next()
            self.x_v.d = image
            self.v_clf_out.forward(clear_buffer=True)
            y_all.append(targets[:rem_values])
            scores_all.append(self.v_clf_out.d.squeeze()[:rem_values])

        y_all = np.concatenate(y_all)
        pred_all = np.concatenate(scores_all)
        return y_all, pred_all


# responsible_ai/gan_data_debiased/compute_fairness.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pickle
import numpy as np
import nnabla as nn
from nnabla.ext_utils import get_extension_context
import args
import classifier as clf
import data_loader as dl
from utils import utils


def get_data_settings(opt):
    """
    Get the data settings
    Args:
        opt : variables that containing values for all of your options
    Returns:
        variables which you need to test & validate
    """
    data_params = {
        "train_beg": opt['train_beg'],
        "valid_beg": opt['valid_beg'],
        "test_beg": opt['test_beg'],
    }
    data_setting = {
        'path': opt['base_img_path'],
        'protected_attribute': opt['protected_attribute'],
        'attribute': opt['attribute'],
        'data_params': data_params,
        'batch_size': opt['batch_size']
    }
    opt['data_setting'] = data_setting
    return opt


if __name__ == "__main__":

    opt = args.get_args()
    opt = get_data_settings(opt)
    attr_list = utils.get_all_attr()
    ctx = get_extension_context(
        opt['context'], device_id=opt['device_id'], type_config=opt['type_config'])
    nn.set_default_context(ctx)
    batch_size = opt['data_setting']['batch_size']
    test = dl.actual_celeba_dataset(opt['data_setting'], batch_size,
                                    augment=False, split='test', shuffle=False)
    AC = clf.attribute_classifier(model_load_path="{}/{}/best/best_acc.h5".format(
        opt['model_save_path'], attr_list[opt['attribute']]))
    val_results = pickle.load(open(r'{}/{}/best/val_results.pkl'.format(
        opt['model_save_path'], attr_list[opt['attribute']]), 'rb'))
    cal_thresh = val_results['cal_thresh']
    f1_thresh = val_results['f1_thresh']
    test_targets, test_scores = AC.get_scores(test)
    test_pred = np.where(test_scores > cal_thresh, 1, 0)
    ap = utils.get_average_precision(test_targets[:, 0], test_scores)
    deo = utils.get_difference_equality_opportunity(test_targets[:, 1],
                                                    test_targets[:, 0], test_pred)
    ba = utils.get_bias_amplification(test_targets[:, 1],
                                      test_targets[:, 0], test_pred)
    kl = utils.get_kl_divergence(test_targets[:, 1],
                                 test_targets[:, 0], test_scores)

    test_results = {
        'AP': ap,
        'DEO': deo,
        'BA': ba,
        'KL': kl,
        'f1_thresh': f1_thresh,
        'cal_thresh': cal_thresh,
        'opt': opt
    }
    print('Test results: ')
    print('AP : {:.1f} +- {:.1f}', 100 * ap)
    print('DEO : {:.1f} +- {:.1f}', 100 * deo)
    print('BA : {:.1f} +- {:.1f}', 100 * ba)
    print('KL : {:.1f} +- {:.1f}', kl)

    with open(r'{}/{}/best/test_scores.pkl'.format(
            opt['model_save_path'], attr_list[opt['attribute']]), 'wb+') as handle:
        pickle.dump(test_scores, handle)

    with open(r'{}/{}/best/test_targets.pkl'.format(
            opt['model_save_path'], attr_list[opt['attribute']]), 'wb+') as handle:
        pickle.dump(test_targets, handle)

    with open(r'{}/{}/best/test_results.pkl'.format(
            opt['model_save_path'], attr_list[opt['attribute']]), 'wb+') as handle:
        pickle.dump(test_results, handle)


# responsible_ai/gan_data_debiased/crop_images.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import imageio
import numpy as np
from PIL import Image


def crop_images(input_path, output_path):
    """
    crop the aligned & cropped 178 * 218 images to 128 * 128
    Args:
        input_path (str) : path of image directory
        output_path (str) : path of cropped images to be saved
    """
    img_list = [f for f in os.listdir(
        input_path) if os.path.splitext(f)[1] == ".jpg"]
    if not os.path.isdir(output_path):
        os.mkdir(output_path)
    # to crop the aligned & cropped 178 * 218 images to 128 * 128
    cx = 121  # x-axis
    cy = 89  # y-axis
    c_pixels = 64  # align and center the images
    for _, item in enumerate(img_list):
        path = os.path.join(input_path, item)
        with open(path, 'rb') as image_path:
            img = Image.open(image_path)
            img = np.array(img.convert('RGB'))
            img = img[cx - c_pixels:cx + c_pixels, cy - c_pixels:cy + c_pixels]
            path = os.path.join(output_path, item)
            imageio.imwrite(path, img)


crop_images(r"\data\celeba\img_align_celeba", r"data\celeba\img_align_celeba")


# responsible_ai/gan_data_debiased/data_loader.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import pickle
import numpy as np
from PIL import Image
import albumentations as A
from nnabla.utils.data_iterator import data_iterator_simple
from nnabla.logger import logger


def data_iterator_celeba(img_path, attributes,
                         transform=None, batch_size=32, num_samples=-1, shuffle=True, rng=None):
    """
    create celebA data iterator
    Args:
        img_path(list) : list of image paths
        attributes (dict) : attribute list
        transform : transform the image(data augmentation)
        batch_size (int) :  number of samples contained in each generated batch
        num_samples (int) : number of samples taken in data loader
                            (if num_samples=-1, it will take all the images in the dataset)
        shuffle (bool) : shuffle the data
    Returns:
        simple data iterator
    """
    imgs = img_path
    attr = attributes
    if num_samples == -1:
        num_samples = len(imgs)
    else:
        logger.info(
            "Num. of data ({}) is used for debugging".format(num_samples))

    def load_func(i):
        pillow_image = Image.open(imgs[i])
        image = np.array(pillow_image)
        transformed_image = transform(image=image)['image'].transpose(2, 0, 1)
        return transformed_image, attr[imgs[i]]

    return data_iterator_simple(load_func, num_samples,
                                batch_size, shuffle=shuffle, rng=rng, with_file_cache=False)


def actual_celeba_dataset(data_settings, batch_size, augment=True, shuffle=True, split='train'):
    """
    Create actual celebA data loader

    Args:
        data_settings (dict) : provide the data settings
        batch_size (int) : number of samples contained in each generated batch
        augment (bool) : data augmentation (True for training)
        shuffle (bool) : shuffle the data
        split (string) : split the dataset depends on the split attribute(train, valid, and test)
    Returns:
        data loader
    """

    list_ids = []
    labels = {}

    # actual celebA dataset
    attribute = data_settings['attribute']
    protected_attribute = data_settings['protected_attribute']
    train_beg = data_settings['data_params']['train_beg']
    valid_beg = data_settings['data_params']['valid_beg']
    test_beg = data_settings['data_params']['test_beg']
    img_path = data_settings['path'] + '/img_align_celeba/'
    attr_path = data_settings['path'] + '/list_attr_celeba.txt'

    label_file = open(attr_path, 'r')
    label_file = label_file.readlines()
    # skipping the first two rows for header
    total_samples = len(label_file) - 2

    if split == 'train':
        number_samples = valid_beg - train_beg
        beg = train_beg
    elif split == 'valid':
        number_samples = test_beg - valid_beg
        beg = valid_beg
    elif split == 'test':
        number_samples = total_samples - test_beg
        beg = test_beg
    else:
        print('Error')
        return

    for i in range(beg + 2, beg + number_samples + 2):
        temp = label_file[i].strip().split()
        list_ids.append(img_path + temp[0])
        labels[img_path + temp[0]] = np.array(
            [int((int(temp[attribute + 1]) + 1) / 2),
             int((int(temp[protected_attribute + 1]) + 1) / 2)])
    # as per the author's citation, we have transformed the input image
    # (resize to 64 * 64, 256 * 256, 224 * 224)
    pre_process = [(64, 64), (256, 256), (224, 224)]
    mean_normalize = (0.485, 0.456, 0.406)
    std_normalize = (0.229, 0.224, 0.225)

    if augment:
        transform = A.Compose([
            A.Resize(pre_process[0][0], pre_process[0][1]),
            A.Resize(pre_process[1][0], pre_process[1][1]),
            A.RandomCrop(width=pre_process[2][0], height=pre_process[2][1]),
            A.HorizontalFlip(p=0.5),  # default p=0.5
            # normalize the image
            A.Normalize(mean=mean_normalize, std=std_normalize)
        ])
    else:
        transform = A.Compose([
            A.Resize(pre_process[0][0], pre_process[0][1]),
            A.Resize(pre_process[1][0], pre_process[1][1]),
            A.CenterCrop(width=pre_process[2][0], height=pre_process[2][1]),
            A.Normalize(mean=mean_normalize, std=std_normalize)
        ])
    loader = data_iterator_celeba(list_ids, labels,
                                  transform=transform, batch_size=batch_size, shuffle=shuffle)
    return loader


def fake_dataset_no_label(path, range1, batch_size=32, shuffle=False):
    """
    Create fake dataset with no label

    Args:
        path (str) : provide the data settings
        range1 (tuple) : range of generated images
        batch_size (int): number of samples contained in each generated batch
        shuffle (bool) : shuffle the data
    Returns:
        data loader
    """

    list_ids = []
    labels = {}

    for i in range(range1[0], range1[1]):
        list_ids.append(path + 'gen_'+str(i)+'.jpg')
        labels[path + 'gen_'+str(i)+'.jpg'] = -1

    # as per the author's citation, we have transformed the input image
    # (resize to 64 * 64, 256 * 256, 224 * 224)
    pre_process = [(64, 64), (256, 256), (224, 224)]
    mean_normalize = (0.485, 0.456, 0.406)
    std_normalize = (0.229, 0.224, 0.225)
    transform = A.Compose([
        A.Resize(pre_process[0][0], pre_process[0][1]),
        A.Resize(pre_process[1][0], pre_process[1][1]),
        A.CenterCrop(width=pre_process[2][0], height=pre_process[2][1]),
        A.Normalize(mean=mean_normalize, std=std_normalize)
    ])

    loader = data_iterator_celeba(list_ids, labels,
                                  transform=transform, batch_size=batch_size, shuffle=shuffle)

    return loader


def debiased_celeba_dataset(data_settings, batch_size, augment=True, split='train', shuffle=True):
    """
    Create augmentation data loader

    Args:
        data_settings (dict) : provide the data settings
        batch_size (int): number of samples contained in each generated batch
        augment (bool) : Data augmentation (True for training)
        shuffle (bool) : shuffle the data
        split (string) : split the dataset depends on the split attribute(train, valid, and test)
    Returns:
        data loader
    """

    list_ids = []
    labels = {}

    # real params configuration
    attribute = data_settings['real_params']['attribute']
    protected_attribute = data_settings['real_params']['protected_attribute']
    train_beg = data_settings['real_params']['data_params']['train_beg']
    valid_beg = data_settings['real_params']['data_params']['valid_beg']
    test_beg = data_settings['real_params']['data_params']['test_beg']
    img_path = data_settings['real_params']['path'] + '/img_align_celeba/'
    attr_path = data_settings['real_params']['path'] + '/list_attr_celeba.txt'

    label_file = open(attr_path, 'r')
    label_file = label_file.readlines()
    # skipping the first two rows for header
    total_samples = len(label_file) - 2

    if split == 'train':
        number_samples = valid_beg - train_beg
        beg = train_beg
    elif split == 'valid':
        number_samples = test_beg - valid_beg
        beg = valid_beg
    elif split == 'test':
        number_samples = total_samples - test_beg
        beg = test_beg
    else:
        print('Error')
        return

    for i in range(beg + 2, beg + number_samples + 2):
        temp = label_file[i].strip().split()
        list_ids.append(img_path + temp[0])
        labels[img_path + temp[0]] = np.array(
            [int((int(temp[attribute + 1]) + 1) / 2),
             int((int(temp[protected_attribute + 1]) + 1) / 2)])
    # as per the author's citation, we have transformed the input image
    # (resize to 64 * 64, 256 * 256, 224 * 224)
    pre_process = [(64, 64), (256, 256), (224, 224)]
    mean_normalize = (0.485, 0.456, 0.406)
    std_normalize = (0.229, 0.224, 0.225)
    if augment:
        transform = A.Compose([
            A.Resize(pre_process[0][0], pre_process[0][1]),
            A.Resize(pre_process[1][0], pre_process[1][1]),
            A.RandomCrop(width=pre_process[2][0], height=pre_process[2][1]),
            A.HorizontalFlip(p=0.5),
            A.Normalize(mean=mean_normalize, std=std_normalize)
        ])
    else:
        transform = A.Compose([
            A.Resize(pre_process[0][0], pre_process[0][1]),
            A.Resize(pre_process[1][0], pre_process[1][1]),
            A.CenterCrop(width=pre_process[2][0], height=pre_process[2][1]),
            A.Normalize(mean=mean_normalize, std=std_normalize)
        ])

    # generated image configuration
    gen_img_path = data_settings['gen_params']['generated_image_path']
    flipped_img_path = data_settings['gen_params']['flipped_images_path']
    attribute_labels = data_settings['gen_params']['label_path']
    domian_labels = data_settings['gen_params']['domain_path']
    gen_image_range = data_settings['gen_params']['flipped_image_range']
    flipped_image_range = data_settings['gen_params']['new_range']

    labeldata = pickle.load(open(attribute_labels, 'rb'))
    labeldata = labeldata[gen_image_range[0]:gen_image_range[1]]
    org_label_start = len(labeldata)
    labeldata = np.tile(labeldata, 2)

    domdata = pickle.load(open(domian_labels, 'rb'))
    domdata = domdata[gen_image_range[0]:gen_image_range[1]]
    domdata = np.concatenate([(1-domdata), domdata])

    for i in range(gen_image_range[0], gen_image_range[1]):
        list_ids.append(os.path.join(gen_img_path, 'gen_'+str(i)+'.jpg'))
        labels[os.path.join(gen_img_path, 'gen_'+str(i)+'.jpg')] = \
            np.array([labeldata[org_label_start], domdata[org_label_start]])
        org_label_start += 1

    for i in range(flipped_image_range[0], flipped_image_range[1]):
        list_ids.append(os.path.join(flipped_img_path, 'gen_'+str(i)+'.jpg'))
        labels[os.path.join(flipped_img_path, 'gen_'+str(i)+'.jpg')] =\
            np.array([labeldata[i], domdata[i]])

    loader = data_iterator_celeba(list_ids, labels,
                                  transform=transform, batch_size=batch_size, shuffle=shuffle)

    return loader


# responsible_ai/gan_data_debiased/eastimate_hyperplane.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pickle
import numpy as np
from sklearn import svm
import args
from utils import utils

if __name__ == "__main__":

    opt = args.get_args()
    attr_list = utils.get_all_attr()

    X = pickle.load(
        open(r"{}/latent_vectors.pkl".format(opt['record_latent_vector']), 'rb'))
    g = pickle.load(open("{}/all_{}_scores.pkl".format(opt['fake_data_dir'],
                                                       attr_list[opt['protected_attribute']]), 'rb'))
    a = pickle.load(open("{}/all_{}_scores.pkl".format(opt['fake_data_dir'],
                                                       attr_list[opt['attribute']]), 'rb'))

    # as per the author's citation, first 15000 fake images are used to learn the separating hyperplane.
    # 10000 synthetic images for training & 5000 synthetic images for validation, labeled with baseline model.
    no_samples_train = 10000
    no_samples_valid = 5000

    total_no_of_samples = len(X)  # total no of samples
    total_samples_train = no_samples_train + \
        no_samples_valid  # total no of samples to train
    no_of_flipped_images = total_no_of_samples - \
        total_samples_train  # no of flipped images to generate

    # Split the data into training and validation sets.
    X_train = X[:no_samples_train, :]
    g_train = g[:no_samples_train]
    a_train = a[:no_samples_train]

    X_val = X[no_samples_train:no_samples_valid, :]
    g_val = g[no_samples_train:no_samples_valid]
    a_val = a[no_samples_train:no_samples_valid]

    # train latent space attribute classifiers ha & hg to learn the attributes of x in the training set.
    # normalize both wt & wg .

    # Fit a protected attribute classifier in the latent space using a linear SVM (max_iter=500000)
    clf_g = svm.LinearSVC(max_iter=500000)  # hg
    clf_g.fit(X_train, g_train)
    # Normalize so that w_g has norm 1.
    clf_g_norm = np.linalg.norm(clf_g.coef_)
    clf_g.coef_ = clf_g.coef_ / (clf_g_norm)  # w_g
    clf_g.intercept_ = clf_g.intercept_ / clf_g_norm  # b_g

    # Fit a target attribute classifier in the latent space using a linear SVM (max_iter=500000)
    clf_a = svm.LinearSVC(max_iter=500000)  # ha
    clf_a.fit(X_train, a_train)
    # Normalize so that w_t has norm 1.
    clf_a_norm = np.linalg.norm(clf_a.coef_)
    clf_a.coef_ = clf_a.coef_ / (clf_a_norm)  # w_a
    clf_a.intercept_ = clf_a.intercept_ / clf_a_norm  # b_a

    # derive the  x'
    # x' = x - ((2 * (w_g^T x + b_g))/(1 - (w_g^T w_a)^2))(w_g - (w_g^T w_a)w_a)

    g_perp_a = clf_g.coef_ - (np.sum(clf_g.coef_ * clf_a.coef_)) * clf_a.coef_

    g_perp_a = g_perp_a / np.linalg.norm(g_perp_a)

    cos_theta = np.sum(clf_g.coef_ * clf_a.coef_)
    # cos(theta)^2 + sin(theta)^2 = 1
    sin_theta = np.sqrt(1 - cos_theta * cos_theta)

    # generate the flipped images for remaining 1,60,000 images
    X_all = np.zeros((no_of_flipped_images, X.shape[1]))

    # For every x, find x' with flipped protected attribute
    for j in range(total_samples_train, total_no_of_samples):
        x = X[j]
        dist = np.sum(clf_g.coef_ * x) + clf_g.intercept_
        X_all[j - total_samples_train] = x - \
            ((2 * dist) / sin_theta) * g_perp_a

    with open(r"{}/latent_vectors_{}.pkl".format(opt['record_latent_vector'],
                                                 attr_list[opt['attribute']]), 'wb+') as handle:
        pickle.dump(X_all, handle)


# responsible_ai/gan_data_debiased/generate_images.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pickle
from os import path, mkdir
import sys
import os
import numpy as np
import nnabla as nn
from nnabla.utils.image_utils import imsave
from nnabla.ext_utils import get_extension_context
import args
from utils import utils
sys.path.append(os.path.abspath('../../image-generation/pggan/'))
from functions import pixel_wise_feature_vector_normalization
from helpers import load_gen


def convert_images_to_uint8(images, drange=[-1, 1]):
    """
        convert float32 -> uint8
    """
    if isinstance(images, nn.Variable):
        images = images.d
    if isinstance(images, nn.NdArray):
        images = images.data

    scale = 255 / (drange[1] - drange[0])
    images = images * scale + (0.5 - drange[0] * scale)
    return np.uint8(np.clip(images, 0, 255))


def generate_images(gen, num_images, n_latent=512,
                    hyper_sphere=True, save_dir=None, latent_vector=None):
    """
    generate the images
    Args:
        gen : load generator
        num_images (int) : number of images to generate
        n_latent (int) : 512-D latent space trained on the CelebA
        hyper_sphere (bool) : default True
        save_dir (str) : directory to save the images
        latent_vector (str) : path to save the latent vectors(.pkl file)
    """

    if not path.isdir(save_dir):
        mkdir(save_dir)
    z_data = np.random.randn(num_images, n_latent, 1, 1)
    # Saving latent vectors
    with open(latent_vector, 'wb+') as f:
        pickle.dump(z_data.reshape((num_images, n_latent)), f)
    z = nn.Variable.from_numpy_array(z_data)
    z = pixel_wise_feature_vector_normalization(z) if hyper_sphere else z
    batch_size = 64
    iterations = int(num_images/batch_size)
    if num_images % batch_size != 0:
        iterations += 1
    count = 0
    for ell in range(iterations):
        y = gen(z[ell * batch_size:(ell + 1) * batch_size], test=True)
        images = convert_images_to_uint8(y, drange=[-1, 1])
        for i in range(images.shape[0]):
            imsave(save_dir+'/gen_'+str(count)+'.jpg',
                   images[i], channel_first=True)
            count += 1
    print("images are generated")


def generate_flipped_images(gen, latent_vector, hyper_sphere=True, save_dir=None):
    """
    generate flipped images
    Args:
        gen : generator
        latent_vector(numpy.ndarray) : latent_vector
        hyper_sphere (bool) : default True
        save_dir (str) : directory to save the images
    """
    if not path.isdir(save_dir):
        mkdir(save_dir)
    z_data = np.reshape(
        latent_vector, (latent_vector.shape[0], latent_vector.shape[1], 1, 1))
    z = nn.Variable.from_numpy_array(z_data)
    z = pixel_wise_feature_vector_normalization(z) if hyper_sphere else z
    batch_size = 64  # we have taken batch size of 64
    num_images = latent_vector.shape[0]
    iterations = int(num_images / batch_size)
    if num_images % batch_size != 0:
        iterations += 1
    count = 0
    for ell in range(iterations):
        y = gen(z[ell * batch_size:(ell + 1) * batch_size], test=True)
        images = convert_images_to_uint8(y, drange=[-1, 1])
        for i in range(images.shape[0]):
            imsave(save_dir + '/gen_' + str(count) +
                   '.jpg', images[i], channel_first=True)
            count += 1

    print("all paired images generated")


if __name__ == "__main__":
    # args
    opt = args.get_args()
    experiment = opt['generate']
    num_images = opt['num_images']
    attr_list = utils.get_all_attr()
    # Context
    ctx = get_extension_context(
        opt['context'], device_id=opt['device_id'], type_config=opt['type_config'])
    nn.set_default_context(ctx)
    nn.set_auto_forward(True)

    # Generate config
    model_load_path = opt['generator_model']
    use_bn = False
    last_act = 'tanh'
    use_wscale = True
    use_he_backward = False
    # Load generator
    gen = load_gen(model_load_path, use_bn=use_bn, last_act=last_act,
                   use_wscale=use_wscale, use_he_backward=use_he_backward)
    if experiment == 'orig':
        save_dir = "{}/AllGenImages".format(opt["fake_data_dir"])
        latent_vector = r"{}/latent_vectors.pkl".format(
            opt['record_latent_vector'])
        generate_images(gen, num_images,
                        save_dir=save_dir, latent_vector=latent_vector)

    if experiment == 'flip':
        save_dir = "{}/{}/".format(
            opt["fake_data_dir"], attr_list[opt['attribute']])
        latent = pickle.load(open(r"{}/latent_vectors_{}.pkl".format(
            opt['record_latent_vector'], attr_list[opt['attribute']]), 'rb'))
        generate_flipped_images(gen, latent, save_dir=save_dir)


# responsible_ai/gan_data_debiased/hallucinate_labels.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pickle
import numpy as np
import nnabla as nn
from nnabla.ext_utils import get_extension_context
import args
import classifier as clf
import data_loader as dl
from utils import utils

if __name__ == "__main__":

    opt = args.get_args()
    attr_list = utils.get_all_attr()
    ctx = get_extension_context(
        opt['context'], device_id=opt['device_id'], type_config=opt['type_config'])
    nn.set_default_context(ctx)
    generated_image_range = (0, 175000)  # total generated images 175000
    loader = dl.fake_dataset_no_label("{}/AllGenImages/".format(opt["fake_data_dir"]),
                                      generated_image_range,
                                      batch_size=opt['batch_size'],
                                      shuffle=False)

    AC = clf.attribute_classifier(model_load_path="{}/{}/best/best_acc.h5".format(
        opt['model_save_path'], attr_list[opt['attribute']]))
    _, scores = AC.get_scores(loader)
    threshold = pickle.load(open("{}/{}/best/val_results.pkl".format(
        opt['model_save_path'], attr_list[opt['attribute']]), 'rb'))['f1_thresh']
    scores = np.where(scores > threshold, 1.0, 0.0)
    with open("{}/all_{}_scores.pkl".format(opt['fake_data_dir'],
                                            attr_list[opt['attribute']]), 'wb+') as handle:
        pickle.dump(scores, handle)


# responsible_ai/gan_data_debiased/main.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys
from nnabla.ext_utils import get_extension_context
import nnabla as nn
import args
import data_loader as dl
import classifier as clf
from utils import utils


def model_train_setting(opt):
    """
    Get the model train settings
    Args:
        opt : variables that containing values for all of your options
    Returns:
        variables which you need to train
    """
    attr_list = utils.get_all_attr()
    if opt['model_train'] == 'baseline':
        data_params = {
            "train_beg": opt['train_beg'],
            "valid_beg": opt['valid_beg'],
            "test_beg": opt['test_beg'],
        }
        data_setting = {
            'path': opt['base_img_path'],
            'protected_attribute': opt['protected_attribute'],
            'attribute': opt['attribute'],
            'data_params': data_params,
            'batch_size': opt['batch_size'],
            'learning_rate': opt['learning_rate'],
            'max_iter': opt['max_iter_base']
        }
        opt['data_setting'] = data_setting

    if opt['model_train'] == 'gan_debiased':
        data_params = {
            "train_beg": opt['train_beg'],
            "valid_beg": opt['valid_beg'],
            "test_beg": opt['test_beg'],
        }
        real_params = {
            'path': opt['base_img_path'],
            'attribute': opt['attribute'],
            'protected_attribute': opt['protected_attribute'],
            'data_params': data_params
        }
        generated_images = "{}/AllGenImages".format(opt["fake_data_dir"])
        flipped_images = "{}/{}/".format(opt["fake_data_dir"],
                                         attr_list[opt['attribute']])
        label_score = "{}/all_{}_scores.pkl".format(opt['fake_data_dir'],
                                                    attr_list[opt['attribute']])
        domain_score = "{}/all_{}_scores.pkl".format(opt['fake_data_dir'],
                                                     attr_list[opt['protected_attribute']])

        generated_params = {
            'generated_image_path': generated_images,
            'flipped_images_path': flipped_images,
            'label_path': label_score,
            'domain_path': domain_score,
            # flipped the images from 15000 to 175000
            'flipped_image_range': (15000, 175000),
            'orig_label_range': (160000, 320000),  # original label range
            'new_range': (0, 160000),  # new images
        }
        data_setting = {
            'real_params': real_params,
            'gen_params': generated_params,
            'batch_size': opt['batch_size'],
            'learning_rate': opt['learning_rate'],
            'max_iter': opt['max_iter_gan_debiased']
        }
        opt['data_setting'] = data_setting

    return opt


def main():
    """
    main method
    """

    opt = args.get_args()
    opt = model_train_setting(opt)
    ctx = get_extension_context(
        opt['context'], device_id=opt['device_id'], type_config=opt['type_config'])

    nn.set_default_context(ctx)
    # model configurations
    batch_size = opt['data_setting']['batch_size']
    learning_rate = opt['data_setting']['learning_rate']
    max_iter = opt['data_setting']['max_iter']
    if (opt["model_train"] == 'baseline'):
        train = dl.actual_celeba_dataset(opt['data_setting'], batch_size,
                                         augment=True, split='train', shuffle=True)
        val = dl.actual_celeba_dataset(opt['data_setting'], batch_size,
                                       augment=False, split='valid', shuffle=False)
        val_weight = None

    elif (opt["model_train"] == 'gan_debiased'):
        train = dl.debiased_celeba_dataset(opt['data_setting'], batch_size,
                                           augment=True, split='train', shuffle=True)
        val = dl.actual_celeba_dataset(opt['data_setting']['real_params'], batch_size,
                                       augment=False, split='valid', shuffle=False)
        val_weight = utils.compute_class_weight(val)
    else:
        print("please provide proper argument")
        sys.exit(0)

    attr_list = utils.get_all_attr()
    if not os.path.exists(opt['model_save_path']):
        os.makedirs(opt['model_save_path'])
    monitor_path = os.path.join(
        opt['model_save_path'], attr_list[opt['attribute']])
    if not os.path.exists(monitor_path):
        os.makedirs(monitor_path)

    attribute_classifier_model = clf.attribute_classifier(batch_size=batch_size,
                                                          learning_rate=learning_rate,
                                                          max_iter=max_iter,
                                                          monitor_path=monitor_path,
                                                          val_weight=val_weight)
    attribute_classifier_model.train(train, val)


if __name__ == '__main__':
    main()


# responsible_ai/gender_bias_mitigation/utils.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import sys
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

GERMAN_CREDIT_DATSET_COLUMN_NAMES = ['status', 'month', 'credit_history',
                                     'purpose', 'credit_amount', 'savings', 'employment',
                                     'investment_as_income_percentage', 'personal_status',
                                     'other_debtors', 'residence_since', 'property', 'age',
                                     'installment_plans', 'housing', 'number_of_credits',
                                     'skill_level', 'people_liable_for', 'telephone',
                                     'foreign_worker', 'credit']


def load_dataset(filepath, column_names=GERMAN_CREDIT_DATSET_COLUMN_NAMES):
    """
    Loads and returns dataset loaded as pandas dataframe
    """
    try:
        dataframe = pd.read_csv(
            filepath, sep=' ', header=None, names=column_names)
    except IOError as err:
        print("IOError: {}".format(err))
        sys.exit(1)
    return dataframe


def preprocess_dataset(dataframe, label_name, protected_attribute_names=["sex"],
                       privileged_classes=[["male"]], favorable_class=1, categorical_features=None,
                       features_to_drop=None):
    """
    Dataset specific preprocessing

    Args:
        dataframe: DataFrame to perform standard preprocessing
        label_name: Name of the target variable
        protected_attribute_names: list of names corresponding to protected attribute
        privileged_classes: Each element is a list of values -> privileged classes
        favorable_class: Label values which are considered favorable
        categorical_features: to be expanded into one-hot vectors
        features_to_drop: column names to drop

    Returns:
        pandas.DataFrame: pre-processed dataframe
    """
    # Create a one-hot encoding of the categorical variables.
    categorical_features = sorted(set(categorical_features) - set(features_to_drop),
                                  key=dataframe.columns.get_loc)
    dataframe = pd.get_dummies(
        dataframe, columns=categorical_features, prefix_sep='=')

    # Map protected attributes to privileged/unprivileged
    for attr, vals in zip(protected_attribute_names, privileged_classes):
        privileged_values = [1.]
        unprivileged_values = [0.]
        if callable(vals):
            dataframe[attr] = dataframe[attr].apply(vals)
        else:
            priv = np.logical_or.reduce(
                np.equal.outer(vals, dataframe[attr].to_numpy()))
            dataframe.loc[priv, attr] = privileged_values[0]
            dataframe.loc[~priv, attr] = unprivileged_values[0]

    # map favourable class
    if callable(favorable_class):
        dataframe[label_name] = dataframe[label_name].apply(favorable_class)
    else:
        dataframe[label_name] = (dataframe[label_name]
                                 == favorable_class).astype(int)

    # Drop unrequested columns
    dataframe = dataframe[sorted(set(dataframe.columns.tolist()) - set(features_to_drop),
                                 key=dataframe.columns.get_loc)]

    return dataframe


def train_model(clf, x_train, y_train, **kwargs):
    """
    Train the model on given training data
    Args:
        clf: various ML models
        x_train: Training data
        y_train: Training targets
        kwargs: dictionary that maps each keyword to the value that we pass alongside it
    Returns:
        classifier model
    """
    # Instantiate the classification model
    if "sample_weight" in kwargs:
        clf.fit(x_train, y_train.ravel(),
                sample_weight=kwargs["sample_weight"])
    else:
        clf.fit(x_train, y_train)
    return clf


def plot_fairness_comp(original, mitigated, metric="SPD"):
    """
    Args:
        original: original bias of the dataset
        mitigated: mitigated bias of the dataset
        metric: name of the fairness metric

    Returns:
        plot the fairness comparison

    """
    plt.figure(facecolor='#FFFFFF', figsize=(4, 4))
    plt.bar(["Original", "Mitigated"], [
            original, mitigated], color=["blue", "green"])
    plt.ylabel(metric)
    plt.title("Before vs After Bias Mitigation", fontsize=15)
    y = [original, mitigated]
    for index, value in enumerate(y):
        if value < 0:
            plt.text(index, value - 0.001,
                     str(round(value, 3)), fontweight='bold', color='red',
                     bbox=dict(facecolor='red', alpha=0.4))
        else:
            plt.text(index, value + 0.001,
                     str(round(value, 3)), fontweight='bold', color='red',
                     bbox=dict(facecolor='red', alpha=0.4))
    plt.grid(None, axis="y")
    plt.show()


# responsible_ai/gradcam/gradcam.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import cv2
import numpy as np


def gradcam(middle_layer):
    """
    Calculate GradCAM.

    Parameters
    ----------
    middle_layer: nn.Variable
        The layer of interest to apply GradCAM


    Returns
    ----------
    heatmap: ndarray
        2D array of same size as width and height of middle_layer
    """
    conv_layer_output = middle_layer.d
    conv_layer_grad = middle_layer.g
    pooled_grad = conv_layer_grad.mean(axis=(0, 2, 3), keepdims=True)
    heatmap = pooled_grad * conv_layer_output
    heatmap = np.maximum(heatmap, 0)    # ReLU
    heatmap = heatmap.mean(axis=(0, 1))
    max_v, min_v = np.max(heatmap), np.min(heatmap)
    if max_v != min_v:
        heatmap = (heatmap - min_v) / (max_v - min_v)
    return heatmap


def overlay_images(base_img, overlay_img, overlay_coef=1.0):
    """
    Overlay two images.

    Parameters
    ----------
    base_img: ndarray
        2D array
    overlay_img: ndarray
        2D array
    overlay_coef: float
        mix rate of overlay_img to base_img (overlay_img: overlay_coef, base_img: 1). 

    Returns
    ----------
    ret_img: ndarray
        2D array of overlaid image
    """
    # resize
    _overlay_img = cv2.resize(
        overlay_img, (base_img.shape[1], base_img.shape[0]))
    # normalize
    _overlay_img = 255 * _overlay_img / _overlay_img.max()
    _overlay_img = _overlay_img.astype('uint8')
    # color adjust
    _overlay_img = cv2.applyColorMap(_overlay_img, cv2.COLORMAP_JET)
    base_img = cv2.cvtColor(base_img, cv2.COLOR_BGR2RGB)
    # overlay
    ret_img = _overlay_img * overlay_coef + base_img
    ret_img = 255 * ret_img / ret_img.max()
    ret_img = ret_img.astype('uint8')
    ret_img = cv2.cvtColor(ret_img, cv2.COLOR_BGR2RGB)
    return ret_img


# responsible_ai/kernel_shap/calculate.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import nnabla as nn
from scipy.special import binom
import copy
import itertools


class Ridge:
    def __init__(self, alpha=0, fit_intercept=True):
        self.alpha = alpha
        self.fit_intercept = fit_intercept

    def fit(self, X: np.ndarray, y: np.ndarray, weights=None):
        y = y.reshape(-1, 1)
        if X.shape[0] != y.shape[0]:
            raise Exception(
                f"Number of X and y rows don't match ({X.shape[0]} != {y.shape[0]})")

        if self.fit_intercept:
            X = np.concatenate([np.ones((X.shape[0], 1)), X], 1)

        if weights is None:
            tmp = X
        else:
            tmp = np.transpose(np.transpose(X) * np.transpose(weights))
        etmp_dot = np.dot(np.transpose(tmp), X)

        if self.alpha != 0:
            etmp_dot = etmp_dot + self.alpha * np.eye(etmp_dot.shape[0])

        try:
            tmp2 = np.linalg.inv(etmp_dot)
        except np.linalg.LinAlgError:
            tmp2 = np.linalg.pinv(etmp_dot)

        self.w = np.dot(tmp2, np.dot(np.transpose(tmp), y))

    def predict(self, X: np.ndarray):
        if self.fit_intercept:
            X = np.concatenate([np.ones((X.shape[0], 1)), X], 1)
        return np.dot(X, self.w)


class KernelSHAP:
    def __init__(self, data, model, X, alpha, nsamples='auto'):
        self.data = data
        self.model = model
        self.X = X
        self.alpha = alpha
        self.nsamples = nsamples

    def shap_values(self):

        out = self.model(self.data)
        out.forward()

        self.train_samples = self.data.d.shape[0]
        self.num_features = self.data.d.shape[1]
        self.num_classes = out.d.shape[1]

        self.weights = np.ones(self.train_samples)
        self.weights /= np.sum(self.weights)
        self.nsamples_run = 0
        groups = [np.array([i]) for i in range(self.num_features)]
        self.fnull = np.sum((out.d.T * self.weights).T, 0)
        expected_value = self.fnull

        explanations = []

        for i in range(len(self.X.d)):
            instance = self.X[i:i + 1, :]
            explanations.append(self.explain(instance, groups))
        s = explanations[0].shape

        outs = [np.zeros((self.X.d.shape[0], s[0])) for j in range(s[1])]
        for i in range(self.X.d.shape[0]):
            for j in range(s[1]):
                outs[j][i] = explanations[i][:, j]
        if len(self.X.d) == 1:
            outs = [sp[0] for sp in outs]
        return outs, expected_value

    def explain(self, instance, groups):
        self.nsamples_added = 0
        varying_inds = self.varying_groups(instance.d, groups)
        varying_feature_groups = [groups[i] for i in varying_inds]
        M = len(varying_feature_groups)
        varying_feature_groups = np.array(varying_feature_groups)
        varying_feature_groups = varying_feature_groups.flatten()

        model_out = self.model(instance)
        model_out.forward()
        fx = model_out.d[0]
        fx = np.array(fx)

        if M == 0:
            phi = np.zeros((self.num_features, self.num_classes))

        elif M == 1:
            phi = np.zeros((self.num_features, self.num_classes))
            diff = fx - self.fnull
            for d in range(self.num_classes):
                phi[varying_inds[0], d] = diff[d]

        else:
            if self.nsamples == "auto":
                self.nsamples = 2 * M + 2**11

            max_samples = 2 ** 30
            if M <= 30:
                max_samples = 2 ** M - 2
                if self.nsamples > max_samples:
                    self.nsamples = max_samples

            self.synth_data = np.tile(self.data.d, (self.nsamples, 1))
            self.mask_matrix = np.zeros((self.nsamples, M))
            self.kernel_weights = np.zeros(self.nsamples)
            self.y = np.zeros(
                (self.nsamples * self.train_samples, self.num_classes))
            self.ey = np.zeros((self.nsamples, self.num_classes))

            num_subset_sizes = np.int(np.ceil((M - 1) / 2.0))
            num_paired_subset_sizes = np.int(np.floor((M - 1) / 2.0))
            weight_vector = np.array([(M - 1.0) / (i * (M - i))
                                     for i in range(1, num_subset_sizes + 1)])
            weight_vector[:num_paired_subset_sizes] *= 2
            weight_vector /= np.sum(weight_vector)

            num_full_subsets = 0
            num_samples_left = self.nsamples
            group_inds = np.arange(M, dtype='int64')
            mask = np.zeros(M)
            remaining_weight_vector = copy.copy(weight_vector)

            for subset_size in range(1, num_subset_sizes + 1):
                nsubsets = binom(M, subset_size)
                if subset_size <= num_paired_subset_sizes:
                    nsubsets *= 2

                if num_samples_left * remaining_weight_vector[subset_size - 1] / nsubsets >= 1.0 - 1e-8:
                    num_full_subsets += 1
                    num_samples_left -= nsubsets

                    if remaining_weight_vector[subset_size - 1] < 1.0:
                        remaining_weight_vector /= (1 -
                                                    remaining_weight_vector[subset_size - 1])

                    w = weight_vector[subset_size - 1] / binom(M, subset_size)
                    if subset_size <= num_paired_subset_sizes:
                        w /= 2.0
                    for inds in itertools.combinations(group_inds, subset_size):
                        mask[:] = 0.0
                        mask[np.array(inds, dtype='int64')] = 1.0
                        self.addsample(varying_feature_groups,
                                       instance.d, mask, w)
                        if subset_size <= num_paired_subset_sizes:
                            mask[:] = np.abs(mask - 1)
                            self.addsample(varying_feature_groups,
                                           instance.d, mask, w)
                else:
                    break

            self.run()

            phi = np.zeros((self.num_features, self.num_classes))
            for d in range(self.num_classes):
                vphi = self.solve(d, M, fx)
                phi[varying_inds, d] = vphi

        return phi

    def varying_groups(self, x, groups):
        varying = np.zeros(self.num_features)
        for i in range(0, self.num_features):
            inds = groups[i]
            x_group = x[0, inds]
            num_mismatches = np.sum(np.frompyfunc(
                self.not_equal, 2, 1)(x_group, self.data.d[:, inds]))
            varying[i] = num_mismatches > 0
        varying_indices = np.nonzero(varying)[0]
        return varying_indices

    def not_equal(self, i, j):
        if isinstance(i, str) or isinstance(j, str):
            return 0 if i == j else 1
        return 0 if np.isclose(i, j, equal_nan=True) else 1

    def addsample(self, varying_feature_groups, x, m, w):
        offset = self.nsamples_added * self.train_samples

        mask = m == 1.0
        groups = varying_feature_groups[mask]
        if len(groups.shape) == 2:
            for group in groups:
                self.synth_data[offset:offset +
                                self.train_samples, group] = x[0, group]
        else:
            evaluation_data = x[0, groups]
            self.synth_data[offset:offset +
                            self.train_samples, groups] = evaluation_data
        self.mask_matrix[self.nsamples_added, :] = m
        self.kernel_weights[self.nsamples_added] = w
        self.nsamples_added += 1

    def run(self):
        num_to_run = self.nsamples_added * self.train_samples - \
            self.nsamples_run * self.train_samples
        data_run = self.synth_data[self.nsamples_run *
                                   self.train_samples:self.nsamples_added*self.train_samples, :]

        data_run_nn = nn.Variable(data_run.shape)
        data_run_nn.d = data_run
        modelout = self.model(data_run_nn)
        modelout.forward()
        modelout = modelout.d
        self.y[self.nsamples_run * self.train_samples:self.nsamples_added *
               self.train_samples, :] = np.reshape(modelout, (num_to_run, self.num_classes))

        for i in range(self.nsamples_run, self.nsamples_added):
            ey_val = np.zeros(self.num_classes)
            for j in range(0, self.train_samples):
                ey_val += self.y[i * self.train_samples +
                                 j, :] * self.weights[j]
            self.ey[i, :] = ey_val

    def solve(self, dim, M, fx):
        ey_adj = self.ey[:, dim] - self.fnull[dim]
        nonzero_inds = np.arange(M)

        if len(nonzero_inds) == 0:
            return np.zeros(M)

        ey_adj2 = ey_adj - self.mask_matrix[:, nonzero_inds[-1]] * (
                    fx[dim] - self.fnull[dim])
        etmp = np.transpose(np.transpose(
            self.mask_matrix[:, nonzero_inds[:-1]]) - self.mask_matrix[:, nonzero_inds[-1]])

        model = Ridge(alpha=self.alpha, fit_intercept=False)
        model.fit(etmp, ey_adj2, self.kernel_weights)
        w = model.w.reshape(-1)

        phi = np.zeros(M)
        phi[nonzero_inds[:-1]] = w
        phi[nonzero_inds[-1]] = (fx[dim] - self.fnull[dim]) - sum(w)

        for i in range(M):
            if np.abs(phi[i]) < 1e-10:
                phi[i] = 0

        return phi

    def calculate_shap(self):
        if len(self.X.d.shape) == 1:
            self.X = self.X.reshape((1, len(self.X.d)))
        values, expected_values = self.shap_values()
        return values, expected_values


# responsible_ai/kernel_shap/kernel_shap.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import numpy as np
import csv

import nnabla as nn

import nnabla.functions as F
import nnabla.parametric_functions as PF

from calculate import KernelSHAP
from visualize import visualize
from train import train_model


def func(args):
    # Load csv

    with open(args.input, 'r') as f:
        reader = csv.reader(f)
        _ = next(reader)
        sample = np.array([[float(r) for r in row] for row in reader])[:, :-1]
    if len(sample.shape) == 1:
        sample = sample.reshape((1, sample.shape[0]))

    with open(args.train, 'r') as f:
        reader = csv.reader(f)
        feature_names = next(reader)[:-1]
        train = np.array([[float(r) for r in row] for row in reader])
    train, train_labels = train[:, :-1], (train[:, -1]).astype(int)
    train_labels = train_labels.reshape((len(train_labels), 1))

    # prepare variable
    data = nn.Variable(train.shape)
    X = nn.Variable(sample.shape)
    labels = nn.Variable(train_labels.shape)

    data.d = train
    X.d = sample
    labels.d = train_labels
    data.need_grad = True

    # define model
    def model(x):
        h = PF.affine(x, (32,), name='Affine')
        h = F.relu(h, True)
        h = PF.affine(h, (args.output_size,), name='Affine_2')
        h = F.softmax(h)
        return h

    # train model
    train_model(model, data, labels)

    # calculate and visualize SHAP
    Kernelshap = KernelSHAP(data, model, X, alpha=0)
    shap_values, expected_values = Kernelshap.calculate_shap()
    fig = visualize(expected_values, shap_values, sample,
                    feature_names, args.class_index, args.index)
    fig.savefig(args.output)


def main():
    parser = argparse.ArgumentParser(
        description='SHAP(tabular)\n'
                    '\n'
                    'A Unified Approach to Interpreting Model Predictions'
                    'Scott Lundberg, Su-In Lee'
                    'https://arxiv.org/abs/1705.07874\n'
                    '', formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument(
        '-i', '--input', help='path to input csv file (csv)', required=True)
    parser.add_argument(
        '-t', '--train', help='path to training dataset csv file (csv)', required=True)
    parser.add_argument(
        '-os', '--output_size', help='output_size', required=True, type=int)
    parser.add_argument(
        '-i2', '--index', help='index to be explained (int), default=0', required=True, default=0, type=int)
    parser.add_argument(
        '-a', '--alpha', help='alpha of Ridge, default=0', required=True, default=0, type=float)
    parser.add_argument(
        '-c2', '--class_index', help='class index (int), default=0', required=True, default=0, type=int)
    parser.add_argument(
        '-o', '--output', help='path to output image, default=shap_tabular.png', required=True, default='shap_tabular.png')
    parser.set_defaults(func=func)
    args = parser.parse_args()

    args.func(args)


if __name__ == '__main__':
    main()


# responsible_ai/kernel_shap/train.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import nnabla as nn
import nnabla.functions as F
import nnabla.solvers as S


def train_model(model, data, labels):
    model(data)
    solver = S.Adam(alpha=0.001, beta1=0.9, beta2=0.999, eps=1e-08)
    solver.set_parameters(nn.get_parameters())
    for i in range(1500):
        out = model(data)
        loss = F.categorical_cross_entropy(out, labels)
        loss.forward()
        solver.zero_grad()
        loss.backward()
        solver.update()


# responsible_ai/kernel_shap/visualize.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np

import matplotlib.pyplot as plt
from matplotlib import lines
from matplotlib.font_manager import FontProperties
from matplotlib.path import Path
from matplotlib.patches import PathPatch
import matplotlib


def plot_additive_plot(data, figsize, text_rotation=0, min_perc=0.05):

    negative_features, total_neg, positive_features, total_pos = format_data(
        data)

    base_value = data['base_value']
    out_value = data['out_value']
    offset_text = (np.abs(total_neg) + np.abs(total_pos)) * 0.04

    fig, ax = plt.subplots(figsize=figsize)

    update_limits(ax, total_pos, positive_features, total_neg,
                  negative_features, base_value)

    separators = (ax.get_xlim()[1] - ax.get_xlim()[0]) / 200

    rectangle_list, separator_list = plot_bars(out_value, negative_features, 'negative',
                                               separators)
    for i in rectangle_list:
        ax.add_patch(i)

    for i in separator_list:
        ax.add_patch(i)

    rectangle_list, separator_list = plot_bars(out_value, positive_features, 'positive',
                                               separators)
    for i in rectangle_list:
        ax.add_patch(i)

    for i in separator_list:
        ax.add_patch(i)

    total_effect = np.abs(total_neg) + total_pos
    fig, ax = plot_labels(fig, ax, out_value, negative_features, 'negative',
                          offset_text, total_effect, min_perc=min_perc, text_rotation=text_rotation)

    fig, ax = plot_labels(fig, ax, out_value, positive_features, 'positive',
                          offset_text, total_effect, min_perc=min_perc, text_rotation=text_rotation)

    plot_higher_lower_element(out_value, offset_text)

    plot_base_element(base_value, ax)

    out_names = data['out_names'][0]
    plot_output_element(out_names, out_value, ax)

    if data['link'] == 'logit':
        plt.xscale('logit')
        ax.xaxis.set_major_formatter(matplotlib.ticker.ScalarFormatter())
        ax.ticklabel_format(style='plain')
        plt.show()
    return fig


def format_data(data):
    negative_features = np.array([[data['features'][x]['effect'],
                                  data['features'][x]['value'],
                                  data['feature_names'][x]]
                                  for x in data['features'].keys() if data['features'][x]['effect'] < 0])

    negative_features = np.array(
        sorted(negative_features, key=lambda x: float(x[0]), reverse=False))

    positive_features = np.array([[data['features'][x]['effect'],
                                   data['features'][x]['value'],
                                   data['feature_names'][x]]
                                 for x in data['features'].keys() if data['features'][x]['effect'] >= 0])
    positive_features = np.array(
        sorted(positive_features, key=lambda x: float(x[0]), reverse=True))

    if data['link'] == 'identity':
        def convert_func(x): return x
    elif data['link'] == 'logit':
        def convert_func(x): return 1 / (1 + np.exp(-x))
    else:
        assert False, "ERROR: Unrecognized link function: " + str(data['link'])

    neg_val = data['out_value']
    for i in negative_features:
        val = float(i[0])
        neg_val = neg_val + np.abs(val)
        i[0] = convert_func(neg_val)
    if len(negative_features) > 0:
        total_neg = np.max(negative_features[:, 0].astype(float)) - \
                    np.min(negative_features[:, 0].astype(float))
    else:
        total_neg = 0

    pos_val = data['out_value']
    for i in positive_features:
        val = float(i[0])
        pos_val = pos_val - np.abs(val)
        i[0] = convert_func(pos_val)

    if len(positive_features) > 0:
        total_pos = np.max(positive_features[:, 0].astype(float)) - \
                    np.min(positive_features[:, 0].astype(float))
    else:
        total_pos = 0

    data['out_value'] = convert_func(data['out_value'])
    data['base_value'] = convert_func(data['base_value'])

    return negative_features, total_neg, positive_features, total_pos


def update_limits(ax, total_pos, pos_features, total_neg,
                  neg_features, base_value):
    ax.set_ylim(-0.5, 0.15)
    padding = np.max([np.abs(total_pos) * 0.2,
                      np.abs(total_neg) * 0.2])

    if len(pos_features) > 0:
        min_x = min(np.min(pos_features[:, 0].astype(
            float)), base_value) - padding
    else:
        min_x = 0
    if len(neg_features) > 0:
        max_x = max(np.max(neg_features[:, 0].astype(
            float)), base_value) + padding
    else:
        max_x = 0
    ax.set_xlim(min_x, max_x)

    plt.tick_params(top=True, bottom=False, left=False, right=False, labelleft=False,
                    labeltop=True, labelbottom=False)
    plt.locator_params(axis='x', nbins=12)

    for key, spine in zip(plt.gca().spines.keys(), plt.gca().spines.values()):
        if key != 'top':
            spine.set_visible(False)


def plot_bars(out_value, features, feature_type, separators, bar=0.1):
    rectangle_list = []
    separator_list = []

    pre_val = out_value
    for index, features in zip(range(len(features)), features):
        if feature_type == 'positive':
            left_bound = float(features[0])
            right_bound = pre_val
            pre_val = left_bound

            separator_indent = np.abs(separators)
            separator_pos = left_bound
            colors = ['#FF0D57', '#FFC3D5']
        else:
            left_bound = pre_val
            right_bound = float(features[0])
            pre_val = right_bound

            separator_indent = - np.abs(separators)
            separator_pos = right_bound
            colors = ['#1E88E5', '#D1E6FA']

        if index == 0:
            if feature_type == 'positive':
                points_rectangle = [[left_bound, 0],
                                    [right_bound, 0],
                                    [right_bound, bar],
                                    [left_bound, bar],
                                    [left_bound + separator_indent, (bar / 2)]
                                    ]
            else:
                points_rectangle = [[right_bound, 0],
                                    [left_bound, 0],
                                    [left_bound, bar],
                                    [right_bound, bar],
                                    [right_bound + separator_indent, (bar / 2)]
                                    ]

        else:
            points_rectangle = [[left_bound, 0],
                                [right_bound, 0],
                                [right_bound + separator_indent *
                                    0.90, (bar / 2)],
                                [right_bound, bar],
                                [left_bound, bar],
                                [left_bound + separator_indent * 0.90, (bar / 2)]]

        line = plt.Polygon(points_rectangle, closed=True, fill=True,
                           facecolor=colors[0], linewidth=0)
        rectangle_list += [line]

        points_separator = [[separator_pos, 0],
                            [separator_pos + separator_indent, (bar / 2)],
                            [separator_pos, bar]]

        line = plt.Polygon(points_separator, closed=None, fill=None,
                           edgecolor=colors[1], lw=3)
        separator_list += [line]

    return rectangle_list, separator_list


def plot_labels(fig, ax, out_value, features, feature_type, offset_text, total_effect=0, min_perc=0.05, text_rotation=0):
    start_text = out_value
    pre_val = out_value

    if feature_type == 'positive':
        colors = ['#FF0D57', '#FFC3D5']
        alignement = 'right'
        sign = 1
    else:
        colors = ['#1E88E5', '#D1E6FA']
        alignement = 'left'
        sign = -1

    if feature_type == 'positive':
        x, y = np.array([[pre_val, pre_val], [0, -0.18]])
        line = lines.Line2D(x, y, lw=1., alpha=0.5, color=colors[0])
        line.set_clip_on(False)
        ax.add_line(line)
        start_text = pre_val

    box_end = out_value
    val = out_value
    for feature in features:
        feature_contribution = np.abs(
            float(feature[0]) - pre_val) / np.abs(total_effect)
        if feature_contribution < min_perc:
            break

        val = float(feature[0])

        if feature[1] == "":
            text = feature[2]
        else:
            text = feature[2] + ' = ' + feature[1]

        if text_rotation != 0:
            va = 'top'
        else:
            va = 'baseline'

        textout_val = plt.text(start_text - sign * offset_text,
                               -0.15, text,
                               fontsize=12, color=colors[0],
                               horizontalalignment=alignement,
                               va=va,
                               rotation=text_rotation)
        textout_val.set_bbox(dict(facecolor='none', edgecolor='none'))

        fig.canvas.draw()
        box_size = textout_val.get_bbox_patch(
        ).get_extents().transformed(ax.transData.inverted())
        if feature_type == 'positive':
            box_end_ = box_size.get_points()[0][0]
        else:
            box_end_ = box_size.get_points()[1][0]

        if box_end_ > ax.get_xlim()[1]:
            textout_val.remove()
            break

        if (sign * box_end_) > (sign * val):
            x, y = np.array([[val, val], [0, -0.18]])
            line = lines.Line2D(x, y, lw=1., alpha=0.5, color=colors[0])
            line.set_clip_on(False)
            ax.add_line(line)
            start_text = val
            box_end = val

        else:
            box_end = box_end_ - sign * offset_text
            x, y = np.array([[val, box_end, box_end],
                             [0, -0.08, -0.18]])
            line = lines.Line2D(x, y, lw=1., alpha=0.5, color=colors[0])
            line.set_clip_on(False)
            ax.add_line(line)
            start_text = box_end

        pre_val = float(feature[0])

    extent = [out_value, box_end, 0, -0.31]
    path = [[out_value, 0], [pre_val, 0], [box_end, -0.08],
            [box_end, -0.2], [out_value, -0.2],
            [out_value, 0]]

    path = Path(path)
    patch = PathPatch(path, facecolor='none', edgecolor='none')
    ax.add_patch(patch)

    lower_lim, upper_lim = ax.get_xlim()
    if (box_end < lower_lim):
        ax.set_xlim(box_end, upper_lim)

    if (box_end > upper_lim):
        ax.set_xlim(lower_lim, box_end)

    if feature_type == 'positive':
        colors = np.array([(255, 13, 87), (255, 255, 255)]) / 255.
    else:
        colors = np.array([(30, 136, 229), (255, 255, 255)]) / 255.

    cm = matplotlib.colors.LinearSegmentedColormap.from_list('cm', colors)

    _, Z2 = np.meshgrid(np.linspace(0, 10), np.linspace(-10, 10))
    im = plt.imshow(Z2, interpolation='quadric', cmap=cm,
                    vmax=0.01, alpha=0.3,
                    origin='lower', extent=extent,
                    clip_path=patch, clip_on=True, aspect='auto')
    im.set_clip_path(patch)

    return fig, ax


def plot_output_element(out_name, out_value, ax):
    x, y = np.array([[out_value, out_value], [0, 0.24]])
    line = lines.Line2D(x, y, lw=2., color='#F2F2F2')
    line.set_clip_on(False)
    ax.add_line(line)

    font0 = FontProperties()
    font = font0.copy()
    font.set_weight('bold')
    textout_val = plt.text(out_value, 0.25, '{0:.2f}'.format(out_value),
                           fontproperties=font,
                           fontsize=14,
                           horizontalalignment='center')
    textout_val.set_bbox(dict(facecolor='white', edgecolor='white'))

    textout_val = plt.text(out_value, 0.33, out_name,
                           fontsize=12, alpha=0.5,
                           horizontalalignment='center')
    textout_val.set_bbox(dict(facecolor='white', edgecolor='white'))


def plot_base_element(base_value, ax):
    x, y = np.array([[base_value, base_value], [0.13, 0.25]])
    line = lines.Line2D(x, y, lw=2., color='#F2F2F2')
    line.set_clip_on(False)
    ax.add_line(line)

    textout_val = plt.text(base_value, 0.33, 'base value',
                           fontsize=12, alpha=0.5,
                           horizontalalignment='center')
    textout_val.set_bbox(dict(facecolor='white', edgecolor='white'))


def plot_higher_lower_element(value, text):
    plt.text(value - text, 0.405, 'higher',
             fontsize=13, color='#FF0D57',
             horizontalalignment='right')

    plt.text(value + text, 0.405, 'lower',
             fontsize=13, color='#1E88E5',
             horizontalalignment='left')

    plt.text(value, 0.4, r'$\leftarrow$',
             fontsize=13, color='#1E88E5',
             horizontalalignment='center')

    plt.text(value, 0.425, r'$\rightarrow$',
             fontsize=13, color='#FF0D57',
             horizontalalignment='center')


def visualize(expected_values, values, X, feature_names=None, out_index=0, index=0):
    if feature_names is None:
        feature_names = ['Feature ' + str(i)
                         for i in range(values[0].shape[1])]

    expected_value = expected_values[out_index]
    if len(values[0].shape) == 1:
        shap_value = values[out_index]
        feature = X
    else:
        shap_value = values[out_index][index, :]
        feature = X[index]
    out_value = np.sum(shap_value) + expected_value

    features = {}
    for i in filter(lambda j: shap_value[j] != 0, range(len(shap_value))):
        features[i] = {
            "effect": shap_value[i],
            "value": feature[i]
        }

    data_config = {
        "out_names": "f[x]",
        "base_value": expected_value,
        "out_value": out_value,
        "link": 'identity',
        "feature_names": feature_names,
        "features": features
    }

    fig = plot_additive_plot(data_config, figsize=(
        20, 3), text_rotation=0, min_perc=0.05)

    return fig


# responsible_ai/massage_data/utils.py

# Copyright 2023 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import sys
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import accuracy_score, confusion_matrix


def get_num_pos_neg(X, target_attribute, target_class_cond, sensitive_attribute, sensitive_class_cond):
    """

    Args:
        X (pandas DataFrame): input dataset
        target_attribute: target attribute of the dataset
        target_class_cond (int): target class label (1/0)
        sensitive_attribute (string): sensitive attribute of the dataset
        sensitive_class_cond (int): sensitive class condition (1/0)

    Returns:
        number of instances in the dataset where the target attribute satisfies the given sensitive attribute condition.

    """
    return len(X[(X[sensitive_attribute] == sensitive_class_cond) & (X[target_attribute] == target_class_cond)])


def get_num_instances(X, attribute, condition):
    """
    get the number of instances
    Args:
        X: input dataset
        attribute (string): attribute for which the number of instances needs to be counted
        condition (int): condition that the attribute must satisfy

    Returns:
        number of instances in the dataset where the attribute satisfies the given condition.

    """
    return len(X[(X[attribute] == condition)])


def get_base_rate(X, target_attribute, target_class_cond, sensitive_attribute, sensitive_class_cond):
    """
    Compute the base rate : Pr(Y = 1) = P/(P+N)
    Args:
        X (pandas DataFrame): input dataset
        target_attribute (string): target attribute of the dataset
        target_class_cond (int): target class label (1/0)
        sensitive_attribute (string): sensitive attribute of the dataset
        sensitive_class_cond (int): sensitive class condition (1/0)
    """
    return (get_num_pos_neg(X, target_attribute, target_class_cond, sensitive_attribute, sensitive_class_cond)
            / get_num_instances(X, sensitive_attribute, sensitive_class_cond))


def statistical_parity_difference(X, target_attribute, target_class_cond,
                                  sensitive_attribute, privileged_class, unprivileged_class):
    """
    Compute difference in the metric between unprivileged and privileged groups.
    Args:
        X (pandas DataFrame): input dataset
        target_attribute (string): target attribute of the dataset
        target_class_cond (int): target class label
        sensitive_attribute (string): sensitive attribute of the dataset
        privileged_class (int): privileged class label of the sensitive attribute
        unprivileged_class (int): unprivileged class label of the sensitive attribute
    Returns: Disparate Impact

    """

    positive_privileged = get_base_rate(
        X, target_attribute, target_class_cond, sensitive_attribute, privileged_class)
    positive_unprivileged = get_base_rate(X, target_attribute, target_class_cond, sensitive_attribute,
                                          unprivileged_class)

    return positive_unprivileged - positive_privileged


def disparate_impact(X, target_attribute, target_class_cond,
                     sensitive_attribute, privileged_class, unprivileged_class):
    """
    Compute the ratio of the metric between unprivileged and privileged groups.

    Args:
        X (pandas DataFrame): input dataset
        target_attribute (string): target attribute of the dataset
        target_class_cond(int): target class label
        sensitive_attribute (string): sensitive attribute of the dataset
        privileged_class (int): privileged class label of the sensitive attribute
        unprivileged_class (int): unprivileged class label of the sensitive attribute

    Returns: Disparate Impact

    """

    positive_privileged = get_base_rate(
        X, target_attribute, target_class_cond, sensitive_attribute, privileged_class)
    positive_unprivileged = get_base_rate(X, target_attribute, target_class_cond, sensitive_attribute,
                                          unprivileged_class)

    return positive_unprivileged / positive_privileged


def get_model_fairness(y, privileged_group, preds):
    """
    Get the model fairness

    Args:
        y: True labels for the data
        privileged_group: privileged group for which we want to check fairness
        preds: predicted probabilities for the positive class

    Returns:
        demographic parity(dpd),disparate_impact(di),
        equal opportunity difference(eod),
        equalised odds(AAOD)

    """
    y_unprivileged, preds_unprivileged = y[privileged_group == False], preds[
        privileged_group == False]
    y_privileged, preds_privileged = y[privileged_group], preds[privileged_group]
    cm_unprivileged = confusion_matrix(y_unprivileged, preds_unprivileged)
    cm_privileged = confusion_matrix(y_privileged, preds_privileged)

    unprivileged_PR = (
        cm_unprivileged[1, 1] + cm_unprivileged[0, 1]) / cm_unprivileged.sum()
    privileged_PR = (cm_privileged[1, 1] +
                     cm_privileged[0, 1]) / cm_privileged.sum()
    dpd = unprivileged_PR - privileged_PR
    di = unprivileged_PR / privileged_PR
    unprivileged_TPR = cm_unprivileged[1, 1] / cm_unprivileged[1].sum()
    privileged_TPR = cm_privileged[1, 1] / cm_privileged[1].sum()
    eod = unprivileged_TPR - privileged_TPR

    unprivileged_FPR = cm_unprivileged[0, 1] / cm_unprivileged[0].sum()
    privileged_FPR = cm_privileged[0, 1] / cm_privileged[0].sum()
    aaod = 0.5 * (
            abs(unprivileged_FPR - privileged_FPR) + abs(unprivileged_TPR - privileged_TPR))
    return dpd, di, eod, aaod


def train_model(clf, x_train, y_train, **kwargs):
    """
    Train the model on given training data
    Args:
        clf: various ML models
        x_train: Training data
        y_train: Training targets
        kwargs: dictionary that maps each keyword to the value that we pass alongside it
    Returns:
        classifier model
    """
    # Instantiate the classification model
    if "sample_weight" in kwargs:
        clf.fit(x_train, y_train.ravel(),
                sample_weight=kwargs["sample_weight"])
    else:
        clf.fit(x_train, y_train)
    return clf


def plot_fairness_comp(original, mitigated, metric="DPD"):
    """
    Args:
        original: original bias of the dataset
        mitigated: mitigated bias of the dataset
        metric: name of the fairness metric

    Returns:
        plot the fairness comparison

    """
    plt.figure(facecolor='#FFFFFF', figsize=(4, 4))
    plt.axhline(y=0.0, color='r', linestyle='-')
    plt.bar(["Original", "Mitigated"], [
            original, mitigated], color=["blue", "green"])
    plt.ylabel(metric)
    plt.title("Before vs After Bias Mitigation", fontsize=15)
    y = [original, mitigated]
    for index, value in enumerate(y):
        if value < 0:
            plt.text(index, value - 0.01,
                     str(round(value, 3)), fontweight='bold', color='red',
                     bbox=dict(facecolor='red', alpha=0.4))
        else:
            plt.text(index, value + 0.01,
                     str(round(value, 3)), fontweight='bold', color='red',
                     bbox=dict(facecolor='red', alpha=0.4))
    plt.grid(None, axis="y")
    plt.show()


# responsible_ai/prejudice_remover_regularizer/utils.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pandas as pd
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
import matplotlib.pyplot as plt
from IPython.display import Markdown, display
import sys
import numpy as np


def load_adult_dataset():
    """
    Loads Census Income Dataset; 
    target variable is whether income exceeds $50K/yr
    (For more info : https://archive.ics.uci.edu/ml/datasets/adult)

    Returns:
        features (pandas.DataFrame): features of the Census income dataset
        target (pandas.Series): target values of the Census income dataset
        sensitive_attributes(pandas.DataFrame): sensitive attributes values of Census income dataset
    """
    file_path = r'https://archive.ics.uci.edu/ml/machine-learning-databases/adult/adult.data'
    column_names = ['age', 'workclass', 'fnlwgt', 'education', 'education_num',
                    'marital_status', 'occupation', 'relationship', 'race', 'sex',
                    'capital_gain', 'capital_loss', 'hours_per_week', 'country', 'target']
    try:
        # Restrict races to White and Black
        input_data = (pd.read_csv(file_path, names=column_names,
                                  na_values="?", sep=r'\s*,\s*', engine='python').loc[lambda df: df['race'].isin(['White', 'Black'])]).dropna()
    except Exception as e:
        print("Error : ", e)
        sys.exit(1)

    # sensitive attributes; we identify 'race' and 'sex' as sensitive attributes
    # privileged class for race and sex is White and Male respectively
    sensitive_attribute_list = ['race', 'sex']
    sensitive_attributes = (input_data.loc[:, sensitive_attribute_list]
                            .assign(race=lambda df: (df['race'] == 'White').astype(int),
                                    sex=lambda df: (df['sex'] == 'Male').astype(int)))

    # targets; 1 when someone makes over 50k , 0 otherwise
    target = (input_data['target'] == '>50K').astype(int)

    # features; note that the 'target' and sensitive attribute columns are dropped
    features = (input_data
                .drop(columns=['target', 'race', 'sex', 'fnlwgt'])
                .pipe(pd.get_dummies, drop_first=True))

    display(Markdown(
        f"features : {features.shape[0]} samples, {features.shape[1]} attributes"))
    display(Markdown(f"targets : {target.shape[0]} samples"))
    display(Markdown(
        f"sensitives attributes : {sensitive_attributes.shape[0]} samples, {sensitive_attributes.shape[1]} attributes"))
    return features, target, sensitive_attributes


def CVS(output_f, output_m):
    """
    Calders and Verwer defined a discrimination score,
    by subtracting the conditional probability of the positive class given a sensitive value 
    from that given a non-sensitive value.
    Args:
        output_f (list): output of unprivileged class
        output_m (list) : output of privileged class
    Returns:
        CV Score (float)
    """
    yf_pred = (output_f >= 0.5)
    ym_pred = (output_m >= 0.5)
    corr_f = np.sum(yf_pred == True)
    corr_m = np.sum(ym_pred == True)
    P_y1_s1 = corr_f / output_f.shape[0]
    P_y1_s0 = corr_m / output_m.shape[0]
    CV_score = np.abs(P_y1_s0 - P_y1_s1)
    return round(CV_score.item(), 4)
    print('Calder-Verwer discrimination score: %.4f' % (CV_score.item()))


def fairness_plot(x, x_name, y_left, y_left_name, y_right, y_right_name):
    """
    Graphical Comparison of fairness & performance with respect to
    Prejudice Regularizer 
    """
    fig, ax1 = plt.subplots(figsize=(10, 7))
    ax1.plot(x, y_left, color='b')
    ax1.set_xlabel(x_name, fontsize=16, fontweight='bold')
    ax1.set_ylabel(y_left_name, color='b', fontsize=16, fontweight='bold')
    ax1.xaxis.set_tick_params(labelsize=14)
    ax1.yaxis.set_tick_params(labelsize=14)
    ax1.set_ylim(min(y_left)-1, max(y_left)+1)

    ax2 = ax1.twinx()
    ax2.plot(x, y_right, color='r')
    ax2.set_ylabel(y_right_name, color='r', fontsize=16, fontweight='bold')
    ax2.set_ylim(0, max(y_right)+0.1)
    ax2.yaxis.set_tick_params(labelsize=14)


# responsible_ai/prejudice_remover_regularizer_images/args.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse


def get_args():
    """
     parse args
     """
    description = "Prejudice Remover Regularizer"
    parser = argparse.ArgumentParser(description)
    parser.add_argument("-d", "--device-id", type=int, default=0,
                        help="device ID of the GPU for training")
    parser.add_argument("-c", "--context", type=str, default="cudnn",
                        help="cpu or cudnn context for training")
    parser.add_argument("--type-config", "-t", type=str, default='float',
                        help='computation type, e.g. "float", "half"')
    parser.add_argument("--celeba_image_train_dir", type=str,
                        default="./data/train", help="path of training image directory")
    parser.add_argument("--celeba_image_valid_dir", type=str,
                        default="./data/valid", help="path of validation image directory")
    parser.add_argument("--attr_path", type=str, default=r'./data/celeba/list_attr_celeba.txt',
                        help="celebA attribute file path and name (ex: ./data/list_attr_celeba.txt)")
    parser.add_argument('--attribute', type=str, default='Attractive',
                        help="choose the target attribute to train "
                             "(e.g., Attractive, Arched EyeBrows, Bushy Eyebrows, smilling ,etc.")
    parser.add_argument('--protected_attribute', type=str, default='Male',
                        help="choose the protected attributes (e.g., Male , Pale_Skin)")
    parser.add_argument("--batch-size", "-b", type=int, default=32,
                        help="batch size")
    parser.add_argument("--learning-rate", "-l", type=float, default=1e-4,
                        help="learning rate is a critical hyperparameter in most learning algorithms")
    parser.add_argument('--max_iter', type=int, default=5086,
                        help='max iterations per epoch')
    parser.add_argument('--total_epochs', type=int, default=20,
                        help='total no of epochs to train the model')

    # Prejudice Remover Regularizer tuning parametr
    parser.add_argument('--eta', type=int, default=0,
                        help='parameter for prejudice remover')

    # save model
    parser.add_argument("--model-save-path", "-o",
                        type=str, default='./results', help='path where the model parameters are saved')
    parser.add_argument("--model-load-path",
                        type=str, default=None, help='path where to load the .h5 model')

    opt = vars(parser.parse_args())

    return opt


# responsible_ai/prejudice_remover_regularizer_images/classifier.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import pickle

import numpy as np
import nnabla as nn
import nnabla.parametric_functions as PF
import nnabla.solvers as S
import nnabla.functions as F

from nnabla.models.imagenet import ResNet50
from sklearn.metrics import average_precision_score
from utils import utils


def clf_resnet50(layer, n_classes=1, train=True):
    """This function uses ResNet-50 pretrained on ImageNet as the base architecture
    and replaces the linear layer in ResNet with two linear layers with the
    hidden layer of size 2,048. Dropout and ReLU are applied between these.
    """
    layer_1 = F.relu(PF.affine(layer, 2048, name='classifier_1'))
    if train:
        layer_1 = F.dropout(layer_1, 0.5)
    out = PF.affine(layer_1, n_classes, name='classifier_2')
    return out


class AttributeClassifier:
    """This class uses ResNet-50 pretrained on ImageNet as the base architecture
    and replaces the linear layer in ResNet with two hidden layers of size 2,048.
    Dropout and ReLU are applied.
    """

    def __init__(self, batch_size=32, learning_rate=1e-4, max_iter=5086,
                 total_epochs=20, eta=0, monitor_path=None,
                 validation_weight=None, model_load_path=None):
        """
        Construct all the necessary attributes for the attribute classifier.
        Args:
            batch_size (int): number of samples contained in each generated batch
            learning_rate (float) : learning rate
            max_iter (int) : maximum iterations for an epoch
            total_epochs (int) : total epochs to train the model
            eta (float): parameter for prejudice remover(eta)
            validation_weight : sample weights
            monitor_path (str) : model parameter to be saved
            model_load_path (str) : load the model
        """
        self.batch_size = batch_size
        self.eta = eta

        # ResNet-50 training graph
        model = ResNet50()
        self.input_image = nn.Variable((self.batch_size,) + model.input_shape)
        self.label = nn.Variable([self.batch_size, 1])
        # fine tuning
        pool = model(self.input_image, training=True, use_up_to='pool')
        self.clf = clf_resnet50(pool, n_classes=1)
        self.clf.persistent = True
        self.clf_sigm = F.sigmoid(self.clf)
        # loss
        self.loss = F.mean(F.sigmoid_cross_entropy(self.clf, self.label))
        self.loss.persistent = True
        # hyper parameters
        self.solver = S.Adam(learning_rate)
        self.solver.weight_decay(1e-05)
        self.solver.set_parameters(nn.get_parameters())

        # validation graph
        self.x_v = nn.Variable((self.batch_size,) + model.input_shape)
        pool_v = model(self.x_v, training=False, use_up_to='pool')
        self.v_clf = clf_resnet50(pool_v, train=False)
        self.v_clf_out = F.sigmoid(self.v_clf)
        self.validation_weight = validation_weight
        # val params
        self.acc = 0.0
        self.total_epochs = total_epochs
        self.max_iter = max_iter
        self.monitor_path = monitor_path
        if model_load_path is not None:
            _ = nn.load_parameters(model_load_path)

    def pr_loss(self, output_f, output_m, eta):
        """Prejudice Remover Regularizer
        Args:
            output_f (NNabla Variable): output of unprivileged class
            output_m (NNabla Variable) : output of privileged class
            eta : parameter for prejudice remover(eta)
        Returns:
            Prejudice Index
        """
        N_female = F.constant(output_f.shape[0])
        N_male = F.constant(output_m.shape[0])
        # male sample, #female sample
        Dxisi = F.stack(N_male, N_female, axis=0)
        y_pred_female = F.sum(output_f)
        y_pred_male = F.sum(output_m)
        P_ys = F.stack(y_pred_male, y_pred_female, axis=0) / Dxisi
        P = F.concatenate(output_f, output_m, axis=0)
        P_y = F.sum(P) / (output_f.shape[0]+output_m.shape[0])
        P_s1y1 = F.log(P_ys[1]) - F.log(P_y)
        P_s1y0 = F.log(1-P_ys[1]) - F.log(1-P_y)
        P_s0y1 = F.log(P_ys[0]) - F.log(P_y)
        P_s0y0 = F.log(1-P_ys[0]) - F.log(1-P_y)
        P_s1y1 = F.reshape(P_s1y1, (P_s1y1.d.size,))
        P_s1y0 = F.reshape(P_s1y0, (P_s1y0.d.size,))
        P_s0y1 = F.reshape(P_s0y1, (P_s0y1.d.size,))
        P_s0y0 = F.reshape(P_s0y0, (P_s0y0.d.size,))
        # PI
        PI_s1y1 = output_f * P_s1y1
        PI_s1y0 = (1 - output_f) * P_s1y0
        PI_s0y1 = output_m * P_s0y1
        PI_s0y0 = (1 - output_m) * P_s0y0
        PI = F.sum(PI_s1y1) + F.sum(PI_s1y0) + F.sum(PI_s0y1) + F.sum(PI_s0y0)
        PI = eta * PI
        return PI

    def train(self, train_loader, val_loader):
        """Train the attribute classifier and save model with best accuracy on validation set
        Args:
             train_loader (NNabla data iterator) : specify training dataset
             val_loader (NNabla data iterator) : specify validation dataset
        """

        model_save = True
        print_freq = 100
        for epoch in range(self.total_epochs):
            for i in range(self.max_iter):
                image, targets = train_loader.next()
                protected_variable = targets[:, 1].reshape(-1, 1)
                targets = targets[:, 0].reshape(-1, 1)
                self.input_image.d = image
                self.label.d = targets
                self.solver.zero_grad()
                self.clf_sigm.forward(clear_no_need_grad=True)
                self.output_f = self.clf_sigm[protected_variable == 0]
                self.output_m = self.clf_sigm[protected_variable == 1]
                self.pi_Loss = self.pr_loss(
                    self.output_f, self.output_m, self.eta)
                self.t_loss = self.loss + self.pi_Loss
                self.t_loss.forward(clear_no_need_grad=True)
                self.t_loss.backward(clear_buffer=True)
                self.solver.update()

                if print_freq and (i % print_freq == 0):
                    print(f'Training epoch {epoch}: [{i+1}|{self.max_iter}], Total loss:{self.t_loss.d} Loss : {self.loss.d}',
                          flush=True)
            if model_save:
                acc, val_targets, val_scores = self.check_avg_precision(
                    val_loader)
                if acc > self.acc:
                    self.acc = acc
                    if not os.path.exists(self.monitor_path):
                        os.makedirs(self.monitor_path)
                    path_best_acc_model = f"{self.monitor_path}/best_acc.h5"
                    nn.save_parameters(path_best_acc_model)

                    # computing model fairness
                    cal_thresh = utils.calibrated_threshold(val_targets[:, 0],
                                                            val_scores)
                    _, f1_thresh = utils.get_f1_threshold(val_targets[:, 0],
                                                          val_scores)
                    val_pred = np.where(val_scores > cal_thresh, 1, 0)
                    with open(self.monitor_path + '/val_scores.pkl', 'wb+') as handle:
                        pickle.dump(val_scores, handle)
                    with open(self.monitor_path + '/val_targets.pkl', 'wb+') as handle:
                        pickle.dump(val_targets, handle)
                    average_precision = utils.get_average_precision(val_targets[:, 0],
                                                                    val_scores)
                    deo = utils.get_diff_in_equal_opportunity(val_targets[:, 1],
                                                              val_targets[:, 0], val_pred)
                    bias_amplification = utils.get_bias_amplification(val_targets[:, 1],
                                                                      val_targets[:, 0], val_pred)
                    outf = val_scores[val_targets[:, 1] == 0]
                    outm = val_scores[val_targets[:, 1] == 1]
                    cv_score = utils.get_cvs(outf, outm, cal_thresh)
                    val_results = {
                        'AP': average_precision,
                        'CV_score': cv_score,
                        'DEO': deo,
                        'BA': bias_amplification,
                        'f1_thresh': f1_thresh,
                        'cal_thresh': cal_thresh,
                        'accuracy': acc,
                    }
                    with open(self.monitor_path + '/val_results.pkl', 'wb+') as handle:
                        pickle.dump(val_results, handle)

    def check_avg_precision(self, loader, weights=None, print_out=True):
        """Compute average precision (AP) from prediction scores.
        Args:
            loader (NNabla data iterator) : specify the data iterator which to compute the AP
            weights : sample weights
            print_out (bool) : print the logs
        Return:
            acc (float) : accuracy of the model
            y_all (numpy.ndarray) : actual all target labels
            pred_all (numpy.ndarray) : predicted score
        """
        y_all, pred_all = self.get_scores(loader)
        acc = average_precision_score(
            y_all[:, 0], pred_all, sample_weight=weights)
        if print_out:
            print(f'Avg precision all = {acc}')
        return acc, y_all, pred_all

    def get_scores(self, loader):
        """
        get predicted scores of the  model
        Args:
            loader : specify the data loader which to predicted the scores
        Returns:
            y_all (numpy.ndarray) : actual all target labels
            pred_all (numpy.ndarray) : all predicted score
        """
        # only for validation and test set
        max_iter = int(loader.size/self.batch_size)
        rem_iter = 0
        if loader.size % self.batch_size != 0:
            rem_iter += 1
            rem_values = loader.size % self.batch_size
        y_all = []
        scores_all = []
        for _ in range(max_iter):
            image, targets = loader.next()
            self.x_v.d = image
            self.v_clf_out.forward(clear_buffer=True)
            y_all.append(targets)
            scores_all.append(self.v_clf_out.d.squeeze())
        for _ in range(rem_iter):
            loader._reset()
            image, targets = loader.next()
            self.x_v.d = image
            self.v_clf_out.forward(clear_buffer=True)
            y_all.append(targets[:rem_values])
            scores_all.append(self.v_clf_out.d.squeeze()[:rem_values])

        y_all = np.concatenate(y_all)
        pred_all = np.concatenate(scores_all)
        return y_all, pred_all


# responsible_ai/prejudice_remover_regularizer_images/data_loader.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import numpy as np
import albumentations as A

from PIL import Image
from nnabla.utils.data_iterator import data_iterator_simple


def data_iterator_celeba(img_path, attr_path, batch_size,
                         target_attribute='Attractive', protected_attribute='Male',
                         num_samples=-1, augment=False, shuffle=False, rng=None):
    """
    create celebA data iterator
    Args:
        img_path (str) : image path directory
        attr_path (str) : celebA attribute file path (ex: list_attr_celeba.txt)
        batch_size (int) :  number of samples contained in each generated batch
        target_attribute (str) : target attribute (ex: Arched EyeBrows,
                                Bushy Eyebrows, smiling,etc..)
        protected_attribute (str): protected attribute (ex: Male, Pale_Skin)
        num_samples (int) : number of samples taken in data loader
                            (if num_samples=-1, it will take all the images in the dataset)
        augment (bool) : data augmentation (True for training)
        shuffle (bool) : shuffle the data (True /False)
        rng : None
    Returns:
        simple data iterator
    """

    imgs = []
    for file in sorted(os.listdir(img_path), key=lambda x: int(x.split(".")[0])):
        imgs.append(os.path.join(img_path, file))
    with open(attr_path, 'r') as f:
        lines = f.readlines()

    attr_list = lines[1].strip().split()
    attr_idx_dict = {attr: i for i, attr in enumerate(attr_list)}
    labels_dict = {}
    for line in lines[2:]:
        line = line.strip().split()
        key = line[0]
        attr = line[1:]
        labels_dict[key] = np.array([int((int(attr[attr_idx_dict[target_attribute]]) + 1) / 2),
                                     int((int(attr[attr_idx_dict[protected_attribute]]) + 1) / 2)])

    # as per the author's citation, we have transformed the input image
    # (resize to , 256 * 256, 224 * 224)
    pre_process = [(256, 256), (224, 224)]
    mean_normalize = (0.485, 0.456, 0.406)
    std_normalize = (0.229, 0.224, 0.225)

    if augment:
        transform = A.Compose([
            A.Resize(pre_process[0][0], pre_process[0][1]),
            A.RandomCrop(width=pre_process[1][0], height=pre_process[1][1]),
            A.HorizontalFlip(p=0.5),
            A.Normalize(mean=mean_normalize, std=std_normalize)
        ])
    else:
        transform = A.Compose([
            A.Resize(pre_process[0][0], pre_process[0][1]),
            A.CenterCrop(width=pre_process[1][0], height=pre_process[1][1]),
            A.Normalize(mean=mean_normalize, std=std_normalize)
        ])
    if num_samples == -1:
        num_samples = len(imgs)
    else:
        print(f"Num. of data ({num_samples}) is used for debugging")

    def load_func(i):
        img = Image.open(imgs[i])
        img = np.array(img.convert('RGB'))
        # transform
        transformed_image = transform(image=img)['image'].transpose(2, 0, 1)
        return transformed_image, labels_dict[os.path.basename(imgs[i])]

    return data_iterator_simple(load_func, num_samples, batch_size,
                                shuffle=shuffle, rng=rng, with_file_cache=False)


# responsible_ai/prejudice_remover_regularizer_images/main.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import nnabla as nn

from nnabla.ext_utils import get_extension_context
import args
import data_loader as di
import classifier as cl


def main():
    """Main method
    """
    opt = args.get_args()
    ctx = get_extension_context(
        opt['context'], device_id=opt['device_id'], type_config=opt['type_config'])
    nn.set_default_context(ctx)

    # model configurations
    batch_size = opt['batch_size']
    learning_rate = opt['learning_rate']
    max_iter = opt['max_iter']
    total_epochs = opt['total_epochs']
    monitor_path = opt['model_save_path']
    pr_eta = opt['eta']
    validation_weight = None
    rng = np.random.RandomState(1)
    train_loader = di.data_iterator_celeba(
        opt['celeba_image_train_dir'], opt['attr_path'], batch_size,
        target_attribute=opt['attribute'], protected_attribute=opt['protected_attribute'],
        augment=True, shuffle=True, rng=rng)
    valid_loader = di.data_iterator_celeba(
        opt['celeba_image_valid_dir'], opt['attr_path'], batch_size,
        target_attribute=opt['attribute'], protected_attribute=opt['protected_attribute'],
        rng=rng)

    clf = cl.AttributeClassifier(batch_size=batch_size, total_epochs=total_epochs,
                                 learning_rate=learning_rate,
                                 eta=pr_eta,
                                 max_iter=max_iter,
                                 validation_weight=validation_weight,
                                 monitor_path=monitor_path)
    clf.train(train_loader, valid_loader)


if __name__ == '__main__':
    main()


# responsible_ai/prejudice_remover_regularizer_images/split_dataset.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import argparse
import shutil


def parse_args():
    """
    parse args
    """
    description = "Split the CelebA dataset "
    parser = argparse.ArgumentParser(description)
    parser.add_argument("--img_path", type=str, default=r'./data/celeba/images',
                        help="Source image path directory")
    parser.add_argument("--attr_path", type=str, default=r'./data/celeba/list_attr_celeba.txt',
                        help="celebA attribute file path (ex: list_attr_celeba.txt)")
    parser.add_argument("--out_dir", type=str, default=r'./train',
                        help='Path where the split data to be saved')
    parser.add_argument("--split", type=str, default='train',
                        help='split the dataset depends on the '
                             'split attribute(train, valid, and test)')
    opt = vars(parser.parse_args())

    return opt


def split_celeba_dataset(img_path, attr_path, out_dir, split="test"):
    """
    split the celebA dataset
    Args:
        img_path (str) : image path directory
        attr_path (str) : celebA attribute file path (ex: list_attr_celeba.txt)
        out_dir (str) : Path where the split data to be saved
        split (string) : split the dataset depends on the split attribute(train, valid, and test)
    """
    # split dataset like authors
    train_beg = 0  # train starts from
    valid_beg = 162770  # valid starts from
    test_beg = 182610  # test starts from

    with open(attr_path, 'r') as label_file:
        label_file = label_file.readlines()

    # skipping the first two rows for header
    total_samples = len(label_file) - 2
    if split == 'train':
        number_samples = valid_beg - train_beg
        beg = train_beg

    elif split == 'valid':
        number_samples = test_beg - valid_beg
        beg = valid_beg

    elif split == 'test':
        number_samples = total_samples - test_beg
        beg = test_beg
    else:
        print('Error')
        return
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)

    for i in range(beg + 2, beg + number_samples + 2):
        temp = label_file[i].strip().split()
        src_dir = os.path.join(img_path, temp[0])
        dst_dir = os.path.join(out_dir, temp[0])
        shutil.copy(src_dir, dst_dir)
    print("splitting completed")


if __name__ == '__main__':
    opt = parse_args()
    split_celeba_dataset(opt['img_path'], opt['attr_path'],
                         opt['out_dir'], opt['split'])


# responsible_ai/rejection_option_based_classification/utils.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pandas as pd
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
import matplotlib.pyplot as plt
from IPython.display import Markdown, display
import sys
from sklearn import metrics
import numpy as np


def load_adult_data():
    """
    Load and preprocess Census Income dataset

    This dataset is used to predict whether income exceeds $50K/yr
    (For more info : https://archive.ics.uci.edu/ml/datasets/adult)

    Returns:
        features (pandas.DataFrame): features of the Census income dataset
        target (pandas.Series): target values of the Census income dataset
        sensitive_attributes(pandas.DataFrame): sensitive attributes values of Census income dataset
    """
    file_path = r'https://archive.ics.uci.edu/ml/machine-learning-databases/adult/adult.data'
    column_names = ['age', 'workclass', 'fnlwgt', 'education', 'education_num',
                    'marital_status', 'occupation', 'relationship', 'race', 'sex',
                    'capital_gain', 'capital_loss', 'hours_per_week', 'country', 'target']
    try:
        # # Restrict races to White and Black
        input_data = (pd.read_csv(file_path, names=column_names,
                                  na_values="?", sep=r'\s*,\s*', engine='python').loc[lambda df: df['race'].isin(['White', 'Black'])]).dropna()
    except Exception as e:
        print("Error : ", e)
        sys.exit(1)

    # sensitive attributes; we identify 'race' and 'sex' as sensitive attributes
    # privileged class for race and sex is White and Male respectively
    sensitive_attribute_names = ['race', 'sex']
    sensitive_attributes = (input_data.loc[:, sensitive_attribute_names].assign(race=lambda df: (df['race'] == 'White').astype(int),
                                                                                sex=lambda df: (df['sex'] == 'Male').astype(int)))

    # targets; 1 when someone makes over 50k , otherwise 0
    target = (input_data['target'] == '>50K').astype(int)

    # features; note that the 'target' and sensitive attribute columns are dropped
    features = (input_data
                .drop(columns=['target', 'race', 'sex', 'fnlwgt'])
                .fillna('Unknown')
                .pipe(pd.get_dummies, drop_first=True))

    display(Markdown(
        f"features : {features.shape[0]} samples, {features.shape[1]} attributes"))
    display(Markdown(f"targets : {target.shape[0]} samples"))
    display(Markdown(
        f"sensitives attributes : {sensitive_attributes.shape[0]} samples, {sensitive_attributes.shape[1]} attributes"))
    display(
        Markdown(f"sensitives attributes names : {sensitive_attribute_names} "))

    return features, target, sensitive_attributes


def get_fairness(y, privileged_group, preds):
    """
    Compute multiple fairness metrics for the classifier:
    1. Demographic parity
    2. Equal opportunity
    3. Equalized odds
    Args:
        y : True data (or target, ground truth)
        privileged_group : list of privileged group values
        preds : data predicted (calculated, output) by your model
    Returns:
        demographic_parity_difference (float): Demographic parity
        equal_opportunity_difference (float): Equal opportunity
        average_abs_odds_difference(float): Equalized odds
    """

    y_unprivileged, preds_unprivileged = y[privileged_group ==
                                           False], preds[privileged_group == False]
    y_privileged, preds_privileged, = y[privileged_group], preds[privileged_group]
    cm_unprivileged = confusion_matrix(y_unprivileged, preds_unprivileged)
    cm_privileged = confusion_matrix(y_privileged, preds_privileged)
    unprivileged_PR = (
        cm_unprivileged[1, 1] + cm_unprivileged[0, 1]) / cm_unprivileged.sum()
    privileged_PR = (cm_privileged[1, 1] +
                     cm_privileged[0, 1]) / cm_privileged.sum()

    # compute demographic parity
    demographic_parity_difference = unprivileged_PR - privileged_PR
    unprivileged_TPR = cm_unprivileged[1, 1] / cm_unprivileged[1].sum()
    privileged_TPR = cm_privileged[1, 1] / cm_privileged[1].sum()

    # compute equal opportunity
    equal_opportunity_difference = unprivileged_TPR - privileged_TPR
    unprivileged_FPR = cm_unprivileged[0, 1] / cm_unprivileged[0].sum()
    privileged_FPR = cm_privileged[0, 1] / cm_privileged[0].sum()

    # compute Equalized odds
    average_abs_odds_difference = 0.5 * \
        (abs(unprivileged_FPR - privileged_FPR) +
         abs(unprivileged_TPR - privileged_TPR))

    return demographic_parity_difference, equal_opportunity_difference, average_abs_odds_difference


def plot_fairness(fairness, ax, metric="DPD", title="fairness metric", bar_x_axis="Original"):
    """
    plot single fairness metric
    """

    ax.set_ylim([-0.6, 0.6])
    ax.axhline(y=0.0, color='r', linestyle='-')
    ax.bar([bar_x_axis], fairness, color="blue", width=2)
    ax.set_ylabel(metric)
    ax.set_title(title, fontsize=10)

    for index, value in enumerate(fairness):
        if value < 0:
            ax.text(index, value - 0.1, str(round(value, 3)), fontweight='bold',
                    color='red', bbox=dict(facecolor='red', alpha=0.4))
        else:
            ax.text(index, value + 0.1, str(round(value, 3)), fontweight='bold',
                    color='red', bbox=dict(facecolor='red', alpha=0.4))


def plot_fairness_multi(DPD, EOD, AAOD, accuracy, bar_x_axis="original"):
    """
    plot three fairness metrics
    """
    fig, axes = plt.subplots(1, 3, figsize=(10, 4), sharey=True)
    fig.suptitle("Model Fairness", fontsize=16)
    plot_fairness([DPD], ax=axes[0], metric="DPD",
                  title="Demographic Parity", bar_x_axis=bar_x_axis)
    plot_fairness([EOD], ax=axes[1], metric="EOD",
                  title="Equal Opportunity", bar_x_axis=bar_x_axis)
    plot_fairness([AAOD], ax=axes[2], metric="AAOD",
                  title="Equalized Odds", bar_x_axis=bar_x_axis)
    fig.text(0.92, 0.65, '\n'.join(
        ["Accuracy of the model:", f"- accuracy : {accuracy:.3f}"]), fontsize='15')
    plt.show()


# responsible_ai/rejection_option_based_classification_images/args.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse


def get_args():
    """
     parse args
     """
    description = "Reject Option-based Classification"
    parser = argparse.ArgumentParser(description)
    parser.add_argument("-d", "--device-id", type=int, default=0,
                        help="device ID of the GPU for training")
    parser.add_argument("-c", "--context", type=str, default="cudnn",
                        help="cpu or cudnn context for training")
    parser.add_argument("--type-config", "-t", type=str, default='float',
                        help='computation type, e.g. "float", "half"')
    parser.add_argument("--celeba_image_train_dir", type=str,
                        default="./data/train", help="path of training image directory")
    parser.add_argument("--celeba_image_valid_dir", type=str,
                        default="./data/valid", help="path of validation image directory")
    parser.add_argument("--celeba_image_test_dir", type=str,
                        default="./data/test", help="path of test image directory")
    parser.add_argument("--attr_path", type=str, default=r'./data/celeba/list_attr_celeba.txt',
                        help="celebA attribute file path and name "
                             "(ex: ./data/list_attr_celeba.txt)")
    parser.add_argument('--attribute', type=str, default='Attractive',
                        help="choose the target attribute to train "
                             "(e.g., Attractive, Arched EyeBrows, Bushy Eyebrows, smilling ,etc.")
    parser.add_argument('--protected_attribute', type=str, default='Male',
                        help="choose the protected attributes (e.g., Male , Pale_Skin)")
    parser.add_argument("--batch-size", "-b", type=int, default=32,
                        help="batch size")
    parser.add_argument("--learning-rate", "-l", type=float, default=1e-4,
                        help="learning rate is a critical hyper parameter "
                             "in most learning algorithms")
    parser.add_argument('--max_iter', type=int, default=5086,
                        help='max iterations per epoch')
    parser.add_argument('--total_epochs', type=int, default=20,
                        help='total no of epochs to train the model')
    # ROC
    parser.add_argument("--optimization_metric", type=str,
                        default="DPD", help="name of the metric to use for "
                                            "optimization - (Demographic parity difference (DPD), "
                                            "Absolute average odds difference (AAOD), "
                                            "Equal opportunity difference (EOD))")
    parser.add_argument("--metric_upper_bound", type=float,
                        default=0.10, help="upper bound constraint on the metric value")
    parser.add_argument("--metric_lower_bound", type=float,
                        default=0.0, help="lower bound constraint on the metric value")
    # save model
    parser.add_argument("--model-save-path", "-o",
                        type=str, default='./results',
                        help='path where the model parameters are saved')
    parser.add_argument("--model-load-path",
                        type=str, default=None, help='path where to load the .h5 model')

    opt = vars(parser.parse_args())

    return opt


# responsible_ai/rejection_option_based_classification_images/classifier.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import pickle

import numpy as np
import nnabla as nn
import nnabla.parametric_functions as PF
import nnabla.solvers as S
import nnabla.functions as F

from nnabla.models.imagenet import ResNet50
from sklearn.metrics import average_precision_score
from utils import utils


def clf_resnet50(layer, n_classes=1, train=True):
    """This function uses ResNet-50 pretrained on ImageNet as the base architecture
    and replaces the linear layer in ResNet with two linear layers with the
    hidden layer of size 2,048. Dropout and ReLU are applied between these.
    """
    layer_1 = F.relu(PF.affine(layer, 2048, name='classifier_1'))
    if train:
        layer_1 = F.dropout(layer_1, 0.5)
    out = PF.affine(layer_1, n_classes, name='classifier_2')
    return out


class AttributeClassifier:
    """This class uses ResNet-50 pretrained on ImageNet as the base architecture
    and replaces the linear layer in ResNet with two hidden layers of size 2,048.
    Dropout and ReLU are applied.
    """

    def __init__(self, batch_size=32, learning_rate=1e-4, max_iter=5086,
                 total_epochs=20, monitor_path=None,
                 validation_weight=None, model_load_path=None):
        """
        Construct all the necessary attributes for the attribute classifier.
        Args:
            batch_size (int): number of samples contained in each generated batch
            learning_rate (float) : learning rate
            max_iter (int) : maximum iterations for an epoch
            total_epochs (int) : total epochs to train the model
            validation_weight : sample weights
            monitor_path (str) : model parameter to be saved
            model_load_path (str) : load the model
        """
        self.batch_size = batch_size

        # ResNet-50 training graph
        model = ResNet50()
        self.input_image = nn.Variable((self.batch_size,) + model.input_shape)
        self.label = nn.Variable([self.batch_size, 1])
        # fine tuning
        pool = model(self.input_image, training=True, use_up_to='pool')
        self.clf = clf_resnet50(pool, n_classes=1)
        self.clf.persistent = True
        self.clf_sigm = F.sigmoid(self.clf)
        # loss
        self.loss = F.mean(F.sigmoid_cross_entropy(self.clf, self.label))
        self.loss.persistent = True
        # hyper parameters
        self.solver = S.Adam(learning_rate)
        self.solver.set_parameters(nn.get_parameters())

        # validation graph
        self.x_v = nn.Variable((self.batch_size,) + model.input_shape)
        pool_v = model(self.x_v, training=False, use_up_to='pool')
        self.v_clf = clf_resnet50(pool_v, train=False)
        self.v_clf_out = F.sigmoid(self.v_clf)
        self.validation_weight = validation_weight
        # val params
        self.acc = 0.0
        self.total_epochs = total_epochs
        self.max_iter = max_iter
        self.monitor_path = monitor_path
        if model_load_path is not None:
            _ = nn.load_parameters(model_load_path)

    def train(self, train_loader, val_loader):
        """Train the attribute classifier and save model with best accuracy on validation set
        Args:
             train_loader (NNabla data iterator) : specify training dataset
             val_loader (NNabla data iterator) : specify validation dataset
        """

        model_save = True
        print_freq = 100
        for epoch in range(self.total_epochs):
            for i in range(self.max_iter):
                image, targets = train_loader.next()
                targets = targets[:, 0].reshape(-1, 1)
                self.input_image.d = image
                self.label.d = targets
                self.solver.zero_grad()
                self.loss.forward(clear_no_need_grad=True)
                self.loss.backward(clear_buffer=True)
                self.solver.update()

                if print_freq and (i % print_freq == 0):
                    print(f'Training epoch {epoch}: [{i+1}|{self.max_iter}], Loss : {self.loss.d}',
                          flush=True)
            if model_save:
                acc, val_targets, val_scores = self.check_avg_precision(
                    val_loader)
                if acc > self.acc:
                    self.acc = acc
                    if not os.path.exists(self.monitor_path):
                        os.makedirs(self.monitor_path, exist_ok=True)
                    path_best_acc_model = f"{self.monitor_path}/best_acc.h5"
                    nn.save_parameters(path_best_acc_model)

                    # computing model fairness
                    accuracy, clf_thresh = \
                        utils.get_best_classification_thershold(val_scores,
                                                                val_targets[:, 0])
                    val_pred = np.where(val_scores > clf_thresh, 1, 0)
                    with open(self.monitor_path + '/val_scores.pkl', 'wb+') as handle:
                        pickle.dump(val_scores, handle)
                    with open(self.monitor_path + '/val_targets.pkl', 'wb+') as handle:
                        pickle.dump(val_targets, handle)
                    average_precision = \
                        utils.get_average_precision(
                            val_targets[:, 0], val_scores)
                    dpd, eod, aaod = \
                        utils.get_fairness(
                            val_targets[:, 0], val_targets[:, 1] == 1, val_pred)
                    val_results = {
                        'AP': average_precision,
                        'Statistical parity difference': dpd,
                        'Equal opportunity difference': eod,
                        'Absolute average odds difference': aaod,
                        'cal_thresh': clf_thresh,
                        'accuracy': accuracy,
                    }
                    with open(self.monitor_path + '/val_results.pkl', 'wb+') as handle:
                        pickle.dump(val_results, handle)

    def check_avg_precision(self, loader, weights=None, print_out=True):
        """Compute average precision (AP) from prediction scores.
        Args:
            loader (NNabla data iterator) : specify the data iterator which to compute the AP
            weights : sample weights
            print_out (bool) : print the logs
        Return:
            acc (float) : accuracy of the model
            y_all (numpy.ndarray) : actual all target labels
            pred_all (numpy.ndarray) : predicted score
        """
        y_all, pred_all = self.get_scores(loader)
        acc = average_precision_score(
            y_all[:, 0], pred_all, sample_weight=weights)
        if print_out:
            print(f'Avg precision all = {acc}')
        return acc, y_all, pred_all

    def get_scores(self, loader):
        """
        get predicted scores of the  model
        Args:
            loader : specify the data loader which to predicted the scores
        Returns:
            y_all (numpy.ndarray) : actual all target labels
            pred_all (numpy.ndarray) : all predicted score
        """
        # only for validation and test set
        max_iter = int(loader.size/self.batch_size)
        rem_iter = 0
        if loader.size % self.batch_size != 0:
            rem_iter += 1
            rem_values = loader.size % self.batch_size
        y_all = []
        scores_all = []
        for _ in range(max_iter):
            image, targets = loader.next()
            self.x_v.d = image
            self.v_clf_out.forward(clear_buffer=True)
            y_all.append(targets)
            scores_all.append(self.v_clf_out.d.squeeze())
        for _ in range(rem_iter):
            loader._reset()
            image, targets = loader.next()
            self.x_v.d = image
            self.v_clf_out.forward(clear_buffer=True)
            y_all.append(targets[:rem_values])
            scores_all.append(self.v_clf_out.d.squeeze()[:rem_values])

        y_all = np.concatenate(y_all)
        pred_all = np.concatenate(scores_all)
        return y_all, pred_all


# responsible_ai/rejection_option_based_classification_images/compute_fairness.py

import os
import pickle
import numpy as np
from sklearn import metrics
from nnabla.ext_utils import get_extension_context
import nnabla as nn
import args
import classifier as clf
import data_loader as di
from utils import utils

if __name__ == '__main__':

    opt = args.get_args()
    ctx = get_extension_context(
        opt['context'], device_id=opt['device_id'], type_config=opt['type_config'])
    nn.set_default_context(ctx)
    batch_size = opt['batch_size']
    model_to_load = os.path.join(
        opt['model_save_path'], opt['attribute'], 'best_acc.h5')
    val_result_file = os.path.join(
        opt['model_save_path'], opt['attribute'], 'val_results.pkl')
    if not (os.path.exists(model_to_load) and os.path.exists(val_result_file)):
        print(f'Provided model path : {model_to_load}')
        print(f'Provided val result file : {val_result_file}')
        raise ('Provided model save path is not proper')
    rng = np.random.RandomState(1)
    # load test data
    test_loader = di.data_iterator_celeba(
        opt['celeba_image_test_dir'], opt['attr_path'], batch_size,
        target_attribute=opt['attribute'], protected_attribute=opt['protected_attribute'],
        rng=rng)
    nn.clear_parameters()
    attribute_classifier_model = clf.AttributeClassifier(
        model_load_path=model_to_load)
    cal_thresh = pickle.load(open(val_result_file, 'rb'))['cal_thresh']
    test_targets, test_scores = attribute_classifier_model.get_scores(
        test_loader)
    test_pred = np.where(test_scores > cal_thresh, 1, 0)
    accuracy = metrics.accuracy_score(test_targets[:, 0], test_pred)
    dpd, eod, aaod = utils.get_fairness(
        test_targets[:, 0], test_targets[:, 1] == 1, test_pred)
    utils.plot_fairness_multi(dpd, eod, aaod, accuracy, bar_x_axis="original")
    print(f"Before applying ROC(Test set) : \n\
                                    Demographic Parity Difference (DPD) : {round(dpd * 100, 2)} \n\
                                    Equal Opportunity Difference (EOD) : {round(eod * 100, 2)}\n\
                                    Absolute Average Odd Difference (AAOD): {round(aaod * 100, 2)}\n\
                                    Accuracy : {round(accuracy * 100, 2)}")


# responsible_ai/rejection_option_based_classification_images/data_loader.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import numpy as np
import albumentations as A

from PIL import Image
from nnabla.utils.data_iterator import data_iterator_simple


def data_iterator_celeba(img_path, attr_path, batch_size,
                         target_attribute='Attractive', protected_attribute='Male',
                         num_samples=-1, augment=False, shuffle=False, rng=None):
    """
    create celebA data iterator
    Args:
        img_path (str) : image path directory
        attr_path (str) : celebA attribute file path (ex: list_attr_celeba.txt)
        batch_size (int) :  number of samples contained in each generated batch
        target_attribute (str) : target attribute (ex: Arched EyeBrows,
                                Bushy Eyebrows, smiling,etc..)
        protected_attribute (str): protected attribute (ex: Male, Pale_Skin)
        num_samples (int) : number of samples taken in data loader
                            (if num_samples=-1, it will take all the images in the dataset)
        augment (bool) : data augmentation (True for training)
        shuffle (bool) : shuffle the data (True /False)
        rng : None
    Returns:
        simple data iterator
    """

    imgs = []
    for file in sorted(os.listdir(img_path), key=lambda x: int(x.split(".")[0])):
        imgs.append(os.path.join(img_path, file))
    with open(attr_path, 'r') as f:
        lines = f.readlines()

    attr_list = lines[1].strip().split()
    attr_idx_dict = {attr: i for i, attr in enumerate(attr_list)}
    labels_dict = {}
    for line in lines[2:]:
        line = line.strip().split()
        key = line[0]
        attr = line[1:]
        labels_dict[key] = np.array([int((int(attr[attr_idx_dict[target_attribute]]) + 1) / 2),
                                     int((int(attr[attr_idx_dict[protected_attribute]]) + 1) / 2)])

    # as per the author's citation, we have transformed the input image
    # (resize to , 256 * 256, 224 * 224)
    pre_process = [(256, 256), (224, 224)]
    mean_normalize = (0.485, 0.456, 0.406)
    std_normalize = (0.229, 0.224, 0.225)

    if augment:
        transform = A.Compose([
            A.Resize(pre_process[0][0], pre_process[0][1]),
            A.RandomCrop(width=pre_process[1][0], height=pre_process[1][1]),
            A.HorizontalFlip(p=0.5),
            A.Normalize(mean=mean_normalize, std=std_normalize)
        ])
    else:
        transform = A.Compose([
            A.Resize(pre_process[0][0], pre_process[0][1]),
            A.CenterCrop(width=pre_process[1][0], height=pre_process[1][1]),
            A.Normalize(mean=mean_normalize, std=std_normalize)
        ])
    if num_samples == -1:
        num_samples = len(imgs)
    else:
        print(f"Num. of data ({num_samples}) is used for debugging")

    def load_func(i):
        img = Image.open(imgs[i])
        img = np.array(img.convert('RGB'))
        # transform
        transformed_image = transform(image=img)['image'].transpose(2, 0, 1)
        return transformed_image, labels_dict[os.path.basename(imgs[i])]

    return data_iterator_simple(load_func, num_samples, batch_size,
                                shuffle=shuffle, rng=rng, with_file_cache=False)


# responsible_ai/rejection_option_based_classification_images/roc.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
from warnings import warn
import numpy as np
from sklearn import metrics
import nnabla as nn
from nnabla.ext_utils import get_extension_context
import classifier as clf
import data_loader as di
from utils import utils
import args


def reject_option_classification(y_true, y_predicted_score,
                                 privileged_group, metric_name="DPD",
                                 metric_upper_bound=0.10, metric_lower_bound=0.05):
    """
    Reject option classification is a postprocessing technique that swaps outcomes
    between privileged and underprivileged groups near the decision boundary.
    Args:
        y_true (numpy.ndarray) : ground truth (correct) target values.
        y_predicted_score (numpy.ndarray) : estimated probability predictions (targets scores)
                                            as returned by a classifier.
        privileged_group (numpy.ndarray): list of privileged group values.
        metric_name (str) : name of the metric to use for the optimization
                            (DPD(Demographic parity difference,
                            AAOD (Absolute average odds difference),
                            EOD (Equal opportunity difference))
        metric_upper_bound (float) : upper bound of constraint on the metric value
        metric_lower_bound (float) : upper bound of constraint on the metric value

    Returns:
        ROC_margin (float): critical region boundary,
        classification_threshold (float) : optimal classification threshold
    """

    low_classification_threshold = 0.01  # smallest classification threshold
    high_classification_threshold = 0.99  # highest classification threshold
    # number of classification threshold b/w low class threshold and high class threshold
    number_classification_threshold = 100
    # number of relevant ROC margins to be used in the optimization search
    number_ROC_margin = 50

    fair_metric_array = np.zeros(
        number_classification_threshold * number_ROC_margin)
    balanced_accuracy_array = np.zeros_like(fair_metric_array)
    ROC_margin_array = np.zeros_like(fair_metric_array)
    classification_threshold_array = np.zeros_like(fair_metric_array)
    count = 0
    # Iterate through class thresholds
    for class_thresh in np.linspace(low_classification_threshold,
                                    high_classification_threshold,
                                    number_classification_threshold):

        classification_threshold = class_thresh
        if class_thresh <= 0.5:
            low_ROC_margin = 0.0
            high_ROC_margin = class_thresh
        else:
            low_ROC_margin = 0.0
            high_ROC_margin = (1.0 - class_thresh)

        # Iterate through ROC margins
        for ROC_margin in np.linspace(
                low_ROC_margin,
                high_ROC_margin,
                number_ROC_margin):
            ROC_margin = ROC_margin
            # Predict using the current threshold and margin
            y_pred = predict(y_predicted_score, classification_threshold,
                             ROC_margin, privileged_group)
            acc = metrics.accuracy_score(y_true, y_pred)
            dpd, eod, aaod = utils.get_fairness(
                y_true, privileged_group, y_pred)
            ROC_margin_array[count] = ROC_margin
            classification_threshold_array[count] = classification_threshold
            balanced_accuracy_array[count] = acc
            if metric_name == "DPD":
                fair_metric_array[count] = dpd
            elif metric_name == "AAOD":
                fair_metric_array[count] = aaod
            elif metric_name == "EOD":
                fair_metric_array[count] = eod
            count += 1

    rel_inds = np.logical_and(fair_metric_array >= metric_lower_bound,
                              fair_metric_array <= metric_upper_bound)
    if any(rel_inds):
        best_ind = np.where(balanced_accuracy_array[rel_inds]
                            == np.max(balanced_accuracy_array[rel_inds]))[0][0]
    else:
        warn("Unable to satisy fairness constraints")
        rel_inds = np.ones(len(fair_metric_array), dtype=bool)
        best_ind = np.where(fair_metric_array[rel_inds]
                            == np.min(fair_metric_array[rel_inds]))[0][0]

    ROC_margin = ROC_margin_array[rel_inds][best_ind]
    classification_threshold = classification_threshold_array[rel_inds][best_ind]

    return ROC_margin, classification_threshold


def predict(y_predicted_score, classification_threshold, ROC_margin, privileged_group):
    """
    Obtain fair predictions with ROC method
    Args:
        y_predicted_score (numpy.ndarray): estimated probability predictions
                                           (targets score) as returned by a classifier.
        classification_threshold (float) : optimal classification threshold
        ROC_margin (float) : critical region boundary
        privileged_group (numpy.ndarray) : list of privileged group values.
    Returns:
        y_pred numpy.ndarray: predictions using ROC method.
    """

    y_pred = np.where(y_predicted_score > classification_threshold, 1, 0)

    # Indices of critical region around the classification boundary
    critical_region_indices = np.logical_and(
        y_predicted_score <= classification_threshold + ROC_margin,
        y_predicted_score > classification_threshold - ROC_margin)

    # Indices of privileged and unprivileged groups
    cond_priv = privileged_group
    cond_unpriv = privileged_group == False
    y_pred[np.logical_and(critical_region_indices,
                          cond_priv)] = 0
    y_pred[np.logical_and(critical_region_indices,
                          cond_unpriv)] = 1

    return y_pred


if __name__ == '__main__':

    opt = args.get_args()

    ctx = get_extension_context(
        opt['context'], device_id=opt['device_id'], type_config=opt['type_config'])
    nn.set_default_context(ctx)

    batch_size = opt['batch_size']
    metric_name = opt['optimization_metric']
    metric_upper_bound = opt['metric_upper_bound']
    metric_lower_bound = opt['metric_lower_bound']
    model_to_load = os.path.join(
        opt['model_save_path'], opt['attribute'], 'best_acc.h5')
    if not os.path.exists(model_to_load):
        print(f'Provided model path : {model_to_load}')
        raise ('Provided model save path is not proper')

    nn.clear_parameters()
    attribute_classifier_model = clf.AttributeClassifier(
        model_load_path=model_to_load)
    rng = np.random.RandomState(1)
    valid_loader = di.data_iterator_celeba(
        opt['celeba_image_valid_dir'], opt['attr_path'], batch_size,
        target_attribute=opt['attribute'], protected_attribute=opt['protected_attribute'],
        rng=rng)
    test_loader = di.data_iterator_celeba(
        opt['celeba_image_test_dir'], opt['attr_path'], batch_size,
        target_attribute=opt['attribute'], protected_attribute=opt['protected_attribute'],
        rng=rng)

    val_targets, val_scores = attribute_classifier_model.get_scores(
        valid_loader)
    test_targets, test_scores = attribute_classifier_model.get_scores(
        test_loader)

    ROC_margin, classification_threshold = \
        reject_option_classification(val_targets[:, 0], val_scores, val_targets[:, 1] == 1,
                                     metric_name=metric_name, metric_upper_bound=metric_upper_bound,
                                     metric_lower_bound=metric_lower_bound)
    pred = predict(val_scores, classification_threshold,
                   ROC_margin, val_targets[:, 1] == 1)
    accuracy = metrics.accuracy_score(val_targets[:, 0], pred)
    dpd, eod, aaod = utils.get_fairness(
        val_targets[:, 0], val_targets[:, 1] == 1, pred)
    utils.plot_fairness_multi(dpd, eod, aaod, accuracy, bar_x_axis="ROC")
    print(f"After applying ROC (Validation set) : \n\
                                    Demographic Parity Difference (DPD) : {round(dpd * 100, 2)} \n\
                                    Equal Opportunity Difference (EOD) : {round(eod * 100, 2)}\n\
                                    Absolute Average Odd Difference (AAOD): {round(aaod * 100, 2)}\n\
                                    Accuracy : {round(accuracy * 100, 2)}")

    pred = predict(test_scores, classification_threshold,
                   ROC_margin, test_targets[:, 1] == 1)
    accuracy = metrics.accuracy_score(test_targets[:, 0], pred)
    dpd, eod, aaod = utils.get_fairness(
        test_targets[:, 0], test_targets[:, 1] == 1, pred)
    utils.plot_fairness_multi(dpd, eod, aaod, accuracy, bar_x_axis="ROC")
    print(f"After applying ROC (Test set) : \n\
                                    Demographic Parity Difference (DPD) : {round(dpd * 100, 2)} \n\
                                    Equal Opportunity Difference (EOD) : {round(eod * 100, 2)}\n\
                                    Absolute Average Odd Difference (AAOD): {round(aaod * 100, 2)}\n\
                                    Accuracy : {round(accuracy * 100, 2)}")


# responsible_ai/rejection_option_based_classification_images/split_dataset.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import argparse
import shutil


def parse_args():
    """
    parse args
    """
    description = "Split the CelebA dataset "
    parser = argparse.ArgumentParser(description)
    parser.add_argument("--img_path", type=str, default=r'./data/celeba/images',
                        help="Source image path directory")
    parser.add_argument("--attr_path", type=str, default=r'./data/celeba/list_attr_celeba.txt',
                        help="celebA attribute file path (ex: list_attr_celeba.txt)")
    parser.add_argument("--out_dir", type=str, default=r'./train',
                        help='Path where the split data to be saved')
    parser.add_argument("--split", type=str, default='train',
                        help='split the dataset depends on the '
                             'split attribute(train, valid, and test)')
    opt = vars(parser.parse_args())

    return opt


def split_celeba_dataset(img_path, attr_path, out_dir, split="test"):
    """
    split the celebA dataset
    Args:
        img_path (str) : image path directory
        attr_path (str) : celebA attribute file path (ex: list_attr_celeba.txt)
        out_dir (str) : Path where the split data to be saved
        split (string) : split the dataset depends on the split attribute(train, valid, and test)
    """
    # split dataset like authors
    train_beg = 0  # train starts from
    valid_beg = 162770  # valid starts from
    test_beg = 182610  # test starts from

    with open(attr_path, 'r') as label_file:
        label_file = label_file.readlines()

    # skipping the first two rows for header
    total_samples = len(label_file) - 2
    if split == 'train':
        number_samples = valid_beg - train_beg
        beg = train_beg

    elif split == 'valid':
        number_samples = test_beg - valid_beg
        beg = valid_beg

    elif split == 'test':
        number_samples = total_samples - test_beg
        beg = test_beg
    else:
        print('Error')
        return
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)

    for i in range(beg + 2, beg + number_samples + 2):
        temp = label_file[i].strip().split()
        src_dir = os.path.join(img_path, temp[0])
        dst_dir = os.path.join(out_dir, temp[0])
        shutil.copy(src_dir, dst_dir)
    print("splitting completed")


if __name__ == '__main__':
    opt = parse_args()
    split_celeba_dataset(opt['img_path'], opt['attr_path'],
                         opt['out_dir'], opt['split'])


# responsible_ai/rejection_option_based_classification_images/train.py

# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import numpy as np
import nnabla as nn
from nnabla.ext_utils import get_extension_context
import args
import data_loader as di
import classifier as cl


def train():
    """
    Train the Attribute classifier
    """
    opt = args.get_args()
    ctx = get_extension_context(
        opt['context'], device_id=opt['device_id'], type_config=opt['type_config'])
    nn.set_default_context(ctx)

    # model configurations
    batch_size = opt['batch_size']
    learning_rate = opt['learning_rate']
    max_iter = opt['max_iter']
    total_epochs = opt['total_epochs']
    validation_weight = None
    rng = np.random.RandomState(1)
    monitor_path = os.path.join(opt['model_save_path'], opt['attribute'])
    train_loader = di.data_iterator_celeba(
        opt['celeba_image_train_dir'], opt['attr_path'], batch_size,
        target_attribute=opt['attribute'], protected_attribute=opt['protected_attribute'],
        augment=True, shuffle=True, rng=rng)
    valid_loader = di.data_iterator_celeba(
        opt['celeba_image_valid_dir'], opt['attr_path'], batch_size,
        target_attribute=opt['attribute'], protected_attribute=opt['protected_attribute'],
        rng=rng)

    clf = cl.AttributeClassifier(batch_size=batch_size, total_epochs=total_epochs,
                                 learning_rate=learning_rate,
                                 max_iter=max_iter,
                                 validation_weight=validation_weight,
                                 monitor_path=monitor_path)
    clf.train(train_loader, valid_loader)


if __name__ == '__main__':
    train()


# responsible_ai/representer_point/args.py

# Copyright (c) 2021 Sony Group Corporation. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


def get_args():
    import argparse
    from distutils.util import strtobool

    parser = argparse.ArgumentParser(description="Pretraining Model")
    parser.add_argument("--model_save_interval", type=int, default=30)
    parser.add_argument(
        "--shuffle_label",
        type=strtobool,
        default=True,
        help="whether shuffle trainig label or not",
    )
    parser.add_argument(
        "--checkpoint",
        type=str,
        default=None,
        help="if resume training, you put your weight",
    )

    parser.add_argument(
        "--device_id",
        "-d",
        type=str,
        default="0",
        help="Device ID the training run on. This is only valid if you specify `-c cudnn`.",
    )

    parser.add_argument(
        "--type_config",
        "-t",
        type=str,
        default="float",
        help='Type of computation. e.g. "float", "half".',
    )

    parser.add_argument(
        "--context",
        "-c",
        type=str,
        default="cudnn",
        help="Extension path. ex) cpu, cudnn.",
    )

    parser.add_argument(
        "--train_batch_size",
        type=int,
        default=100,
        help="Batch size for training (per replica).",
    )

    parser.add_argument("--val-iter", type=int, default=100)
    parser.add_argument(
        "--val_batch_size",
        type=int,
        default=100,
        help="Batch size for eval. (per replica)",
    )

    parser.add_argument(
        "--train_epochs", type=int, default=270, help="Number of epochs to train for."
    )

    parser.add_argument(
        "--seed", "-s", help="random seed number default=0", default=0, type=int
    )

    args = parser.parse_args()

    return args


# responsible_ai/representer_point/cifar10_load.py

# Copyright (c) 2021 Sony Group Corporation. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Provide data iterator for CIFAR10 examples.
"""
import numpy as np
import tarfile
import random
from nnabla.utils.data_iterator import data_iterator
from nnabla.logger import logger
from nnabla.utils.data_source import DataSource
from nnabla.utils.data_source_loader import download


class Cifar10NumpySource(DataSource):
    def __init__(self, x, y, y_shuffled=False, shuffle=False, rng=None):
        super(Cifar10NumpySource, self).__init__(shuffle=shuffle, rng=rng)
        self._images = x
        self._labels = y
        if y_shuffled is not False:
            self._shuffle_labels = y_shuffled
        else:
            self._shuffle_labels = False
        self._shuffle = shuffle
        self._size = self._labels.size
        if self._shuffle_labels is not False:
            self._variables = ("x", "y", "shuffle")
        else:
            self._variables = ("x", "y")
        if rng is None:
            rng = np.random.RandomState(313)
        self.rng = rng
        self.reset()

    def _get_data(self, position):
        image = self._images[self._indexes[position]]
        label = self._labels[self._indexes[position]]

        if self._shuffle_labels is not False:
            shuffle = self._shuffle_labels[self._indexes[position]]
            return (image, label, shuffle)
        else:

            return (image, label)

    def reset(self):
        if self._shuffle:
            self._indexes = self.rng.permutation(self._size)
        else:
            self._indexes = np.arange(self._size)
        super(Cifar10NumpySource, self).reset()

    @property
    def images(self):
        """Get copy of whole data with a shape of (N, 1, H, W)."""
        return self._images.copy()

    @property
    def labels(self):
        """Get copy of whole label with a shape of (N, 1)."""
        return self._labels.copy()

    @property
    def shuffled_labels(self):
        """Get copy of whole label with a shape of (N, 1)."""
        if self._shuffle_labels:
            return self._shuffle_labels.copy()


class Cifar10DataSource(DataSource):
    """
    Get data directly from cifar10 dataset from Internet(yann.lecun.com).
    """

    def __init__(
        self,
        train=True,
        shuffle=False,
        rng=None,
        label_shuffle=False,
        label_shuffle_rate=0.1,
    ):
        print("label_shuffled: ", label_shuffle)
        super(Cifar10DataSource, self).__init__(shuffle=shuffle, rng=rng)

        self._train = train
        self._label_shuffle = label_shuffle
        self._shuffle = shuffle

        data_uri = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"
        logger.info("Getting labeled data from {}.".format(data_uri))
        r = download(data_uri)  # file object returned
        with tarfile.open(fileobj=r, mode="r:gz") as fpin:
            # Training data
            if train:
                images = []
                labels = []
                for member in fpin.getmembers():
                    if "data_batch" not in member.name:
                        continue
                    fp = fpin.extractfile(member)
                    data = np.load(fp, encoding="bytes", allow_pickle=True)
                    images.append(data[b"data"])
                    labels.append(data[b"labels"])
                self._size = 50000
                self._images = np.concatenate(
                    images).reshape(self._size, 3, 32, 32)

                self._labels = np.concatenate(labels).reshape(-1, 1)
                self.raw_label = self._labels.copy()

                if label_shuffle:
                    self.shuffle_rate = label_shuffle_rate
                    self.label_shuffle()
                    print(f"{self.shuffle_rate*100}% of data was shuffled ")
                    print(len(np.where(self._labels != self.raw_label)[0]))
            # Validation data
            else:
                for member in fpin.getmembers():
                    if "test_batch" not in member.name:
                        continue
                    fp = fpin.extractfile(member)
                    data = np.load(fp, encoding="bytes", allow_pickle=True)
                    images = data[b"data"]
                    labels = data[b"labels"]
                self._size = 10000
                self._images = images.reshape(self._size, 3, 32, 32)
                self._labels = np.array(labels).reshape(-1, 1)
                self.raw_label = self._labels.copy()
        r.close()
        logger.info("Getting labeled data from {}.".format(data_uri))

        self._size = self._labels.size
        if self._label_shuffle is not False:
            self._variables = ("x", "y", "shuffle")
        else:
            self._variables = ("x", "y", "y")
        if rng is None:
            rng = np.random.RandomState(313)
        self.rng = rng
        self.reset()

    def _get_data(self, position):
        image = self._images[self._indexes[position]]
        label = self.raw_label[self._indexes[position]]

        if self._label_shuffle:
            shuffle = self._labels[self._indexes[position]]
            return (image, label, shuffle)
        else:
            return (image, label, label)

    def label_shuffle(self):
        num_cls = int(np.max(self._labels)) + 1
        raw_label = self._labels.copy()
        extract_num = int(len(self._labels) * self.shuffle_rate // 10)
        for i in range(num_cls):
            extract_ind = np.where(raw_label == i)[0]
            labels = [j for j in range(num_cls)]
            labels.remove(i)  # candidate of shuffle label
            artificial_label = [
                labels[int(i) % (num_cls - 1)] for i in range(int(extract_num))
            ]
            artificial_label = np.array(
                random.sample(artificial_label, len(artificial_label))
            ).astype("float32")
            convert_label = np.array([i for _ in range(len(extract_ind))])
            convert_label[-extract_num:] = artificial_label

            self._labels[extract_ind] = convert_label.reshape(-1, 1)

    def reset(self):
        if self._shuffle:
            self._indexes = self.rng.permutation(self._size)
        else:
            self._indexes = np.arange(self._size)
        super(Cifar10DataSource, self).reset()

    @property
    def _raw_label(self):
        return self.raw_label.copy()

    @property
    def images(self):
        """Get copy of whole data with a shape of (N, 1, H, W)."""
        return self._images.copy()

    @property
    def labels(self):
        """Get copy of whole label with a shape of (N, 1)."""
        return self._labels.copy()


def data_source_cifar10(train=True, rng=None, shuffle=True, label_shuffle=False):
    """
    Provide DataIterator with :py:class:`Cifar10DataSource`
    with_memory_cache and with_file_cache option's default value is all False,
    because :py:class:`Cifar10DataSource` is able to store all data into memory.
    """
    return Cifar10DataSource(
        train=train, shuffle=shuffle, rng=rng, label_shuffle=label_shuffle
    )


def data_iterator_cifar10(
    batch_size,
    train=True,
    rng=None,
    shuffle=True,
    with_memory_cache=False,
    with_file_cache=False,
):
    """
    Provide DataIterator with :py:class:`Cifar10DataSource`
    with_memory_cache and with_file_cache option's default value is all False,
    because :py:class:`Cifar10DataSource` is able to store all data into memory.

    """
    return data_iterator(
        Cifar10DataSource(train=train, shuffle=shuffle, rng=rng),
        batch_size,
        rng,
        with_memory_cache,
        with_file_cache,
    )


# responsible_ai/representer_point/compute_score.py

# Copyright (c) 2021 Sony Group Corporation. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os
import random
from distutils.util import strtobool

import h5py
import nnabla as nn
import nnabla.functions as F
import nnabla.parametric_functions as PF
import nnabla.solvers as S
import numpy as np
from nnabla.ext_utils import get_extension_context
from cifar10_load import data_source_cifar10


def label_shuffle(labels_, shuffle_rate, seed):
    random.seed(seed)
    num_cls = int(np.max(labels_)) + 1
    raw_label = labels_.copy()
    extract_num = int(len(labels_) * shuffle_rate // 10)
    for i in range(num_cls):
        extract_ind = np.where(raw_label == i)[0]
        labels = [j for j in range(num_cls)]
        labels.remove(i)  # candidate of shuffle label
        artificial_label = [
            labels[int(i) % (num_cls - 1)] for i in range(int(extract_num))
        ]
        artificial_label = np.array(
            random.sample(artificial_label, len(artificial_label))
        ).astype("float32")
        convert_label = np.array([i for _ in range(len(extract_ind))])
        convert_label[-extract_num:] = artificial_label

        labels_[extract_ind] = convert_label.reshape(-1, 1)

    return raw_label, labels_


def load_cifar_labels(shuffle=False):

    train_data_source = data_source_cifar10(
        train=True, shuffle=True, label_shuffle=shuffle
    )
    test_data_source = data_source_cifar10(
        train=False, shuffle=False, label_shuffle=False
    )
    y_train = train_data_source.labels
    y_test = test_data_source.labels

    if shuffle:
        y_raw, y_train = label_shuffle(y_train, shuffle_rate=0.1, seed=0)

    return y_train, y_test


def loss_function(pred, y_output, N):
    # calculate cross entropy loss
    Phi = F.sum(-y_output * F.log_softmax(pred, axis=1))
    # calculate l2 norm of affine layer
    l2 = 0
    for param in nn.get_parameters().values():
        l2 += F.sum(param ** 2)
    loss = l2 * args.lmbd + Phi / N

    return loss, Phi, l2


def backtracking_line_search(grad_norm, x, y, loss, N, val, l2, threshold=1e-10):
    t = 10.0
    beta = 0.5
    params = nn.get_parameters().values()
    p_data_org_list = [F.identity(p.data) for p in params]
    p_grad_org_list = [F.identity(p.grad) for p in params]

    while True:
        for p, p_data_org, p_grad_org in zip(params, p_data_org_list, p_grad_org_list):
            p.data.copy_from(p_data_org - t * p_grad_org)

        loss.forward()
        if t < threshold:
            print("t too small")
            break
        if (loss.d - val + t * grad_norm.data ** 2 / 2) >= 0:
            t = beta * t
        else:
            break

    return params


def calculate_alpha(
    parameter_list,
    X,
    Y,
    Y_label,
    feature_valid,
    solver,
    output_valid,
    pred,
    loss,
    phi,
    l2,
):

    min_loss = 10000.0
    feature_valid.d = X
    output_valid.d = Y
    for epoch in range(args.epoch):
        phi_loss = 0

        loss.forward()
        solver.zero_grad()
        loss.backward()
        phi_loss = phi.d / len(X)
        temp_W = parameter_list
        grad_loss = F.add_n(
            *[F.mean(F.abs(p.grad)) for p in nn.get_parameters().values()]
        )
        grad_norm = F.add_n(*[F.norm(p.grad)
                            for p in nn.get_parameters().values()])

        if grad_loss.data < min_loss:
            if epoch == 0:
                init_grad = grad_loss.data
            min_loss = grad_loss.data
            best_W = temp_W
            if min_loss < init_grad / 200:
                print("stopping criteria reached in epoch :{}".format(epoch))
                break
        parameter_list = backtracking_line_search(
            grad_norm, X, Y, loss, len(X), loss.d, l2
        )
        if epoch % 100 == 0:
            print(
                "Epoch:{:4d}\tloss:{}\tphi_loss:{}\tl2(lmbd):{}\tgrad:{}".format(
                    epoch, loss.d, phi_loss, args.lmbd * l2.d, grad_loss.data
                )
            )

    for weight, param in zip(nn.get_parameters().values(), best_W):
        weight.data.copy_from(param.data)

    softmax_value = F.softmax(pred)
    softmax_value.forward()
    # derivative of softmax cross entropy
    weight_matrix = softmax_value.d - Y
    weight_matrix = np.divide(weight_matrix, (-2.0 * args.lmbd * len(Y)))

    np.save(os.path.join(data_dir, "weight_matrix.npy"), weight_matrix)

    # computer alpha
    alpha = []
    for ind, label in enumerate(Y_label.reshape(-1)):
        alpha.append(float(weight_matrix[ind, int(label)]))
    alpha = np.abs(np.array(alpha))
    np.save(os.path.join(data_dir, "alpha_vgg_nnabla_score.npy"), alpha)

    # calculate correlation
    w = np.matmul(X.T, weight_matrix)
    temp = np.matmul(X, w)
    softmax_value = F.softmax(nn.Variable.from_numpy_array(temp))
    softmax_value.forward()
    y_p = softmax_value.d

    print(
        "L1 difference between ground truth prediction and prediction by representer theorem decomposition"
    )
    print(np.mean(np.abs(Y - y_p)))

    from scipy.stats.stats import pearsonr

    print(
        "pearson correlation between ground truth  prediction and prediciton by representer theorem"
    )
    corr, _ = pearsonr(Y.reshape(-1), y_p.reshape(-1))
    print(corr)


def main():
    extension_module = args.context
    ctx = get_extension_context(
        extension_module, device_id=args.device_id, type_config=args.type_config
    )
    nn.set_default_context(ctx)

    with h5py.File(os.path.join(data_dir, "info.h5"), "r") as hf:
        Y_label = hf["label"][:]
        X_feature = hf["feature"]["train"][:]
        Y_train = hf["output"]["train"][:]
        parameter_list = [hf["param"]["weight"][:], hf["param"]["bias"][:]]

    feature_shape = X_feature.shape

    # computation graph
    feature_valid = nn.Variable((feature_shape[0], feature_shape[1]))
    output_valid = nn.Variable((feature_shape[0], 10))
    with nn.parameter_scope("final_layer"):
        pred = PF.affine(feature_valid, 10)
    loss, phi, l2 = loss_function(pred, output_valid, len(X_feature))

    # parameter initialized
    for weight, param in zip(nn.get_parameters().values(), parameter_list):
        weight.data.copy_from(nn.NdArray.from_numpy_array(param))

    solver = S.Sgd(lr=1.0)
    solver.set_parameters(nn.get_parameters())
    for ind, (name, param) in enumerate(nn.get_parameters().items()):
        param.grad.zero()
        param.need_grad = True

    calculate_alpha(
        parameter_list,
        X_feature,
        Y_train,
        Y_label,
        feature_valid,
        solver,
        output_valid,
        pred,
        loss,
        phi,
        l2,
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--lmbd", type=float, default=0.003)
    parser.add_argument("--epoch", type=int, default=3000)
    parser.add_argument("--shuffle_label", type=strtobool)
    parser.add_argument(
        "--context",
        "-c",
        type=str,
        default="cudnn",
        help="Extension path. ex) cpu, cudnn.",
    )
    parser.add_argument(
        "--device_id",
        "-d",
        type=str,
        default="0",
        help="Device ID the training run on. This is only valid if you specify `-c cudnn`.",
    )
    parser.add_argument(
        "--type_config",
        "-t",
        type=str,
        default="float",
        help='Type of computation. e.g. "float", "half".',
    )

    args = parser.parse_args()
    input_dir = (
        "./data/input/shuffle" if args.shuffle_label else "./data/input/no_shuffle"
    )
    data_dir = "./data/info/shuffle" if args.shuffle_label else "./data/info/no_shuffle"

    print(args)
    main()


# responsible_ai/representer_point/evaluate_correlation.py

# Copyright (c) 2021 Sony Group Corporation. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import matplotlib.pyplot as plt
import h5py
import numpy as np
import nnabla as nn
import nnabla.functions as F

from scipy.stats.stats import pearsonr

with h5py.File("./data/info/no_shuffle/info.h5", "r") as hf:
    train_feature = hf["feature"]["train"][:]
    test_feature = hf["feature"]["test"][:]
    train_output = hf["output"]["train"][:]
    test_output = hf["output"]["test"][:]
feature_dict = {"train": train_feature, "test": test_feature}
output_dict = {"train": train_output, "test": test_output}

weight_matrix = np.load("./data/info/no_shuffle/weight_matrix.npy")

predicted_dict = {}

for phase, feature in feature_dict.items():
    w = np.matmul(train_feature.T, weight_matrix)
    temp = np.matmul(feature, w)
    softmax = F.softmax(nn.Variable.from_numpy_array(temp))
    softmax.forward()
    predicted_dict[phase] = softmax.d


np.random.seed(401)
train_rand_idx = np.random.choice(50000, 2000, replace=False)
np.random.seed(4)
test_rand_idx = np.random.choice(10000, 1000, replace=False)
rand_idx_dict = {"train": train_rand_idx, "test": test_rand_idx}


fontsize = 13
fig, ax = plt.subplots(1, 2)
for ind, phase in enumerate(["train", "test"]):

    output = output_dict[phase]
    pred = predicted_dict[phase]
    print(pred.shape)
    print(output.shape)
    corr, _ = pearsonr(pred.reshape(-1), output.reshape(-1))
    ax[ind].scatter(
        pred[rand_idx_dict[phase]], output[rand_idx_dict[phase]], marker="."
    )
    ax[ind].set_title("Correlation: %f" % corr)
    ax[ind].set_xlabel(f"Predicted {phase} Output", fontsize=fontsize)
    ax[ind].set_ylabel(f"Actual {phase} Output", fontsize=fontsize)
    ax[ind].set_aspect("equal")

plt.tight_layout()
plt.savefig("./data/info/no_shuffle/correlation.png", dpi=200)


# responsible_ai/representer_point/generate_feature.py

# Copyright (c) 2021 Sony Group Corporation. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import functools
import os
from distutils.util import strtobool

import h5py
import nnabla as nn
import nnabla.functions as F
import numpy as np
from nnabla.ext_utils import get_extension_context
from nnabla.utils.data_iterator import data_iterator


from cifar10_load import Cifar10NumpySource
from model import vgg16_prediction

bs_valid = 100
CHECKPOINTS_PATH_FORMAT = "params_270.h5"


def ensure_dir(dir):
    if not os.path.exists(dir):
        os.makedirs(dir)


def feature_extractor(dataloaders_dict, image_valid, pred_hidden, pred):
    phase_feature, phase_preds = {}, {}
    for phase, dataloader in dataloaders_dict.items():
        extracted_feature, preds = [], []
        iteration = int(dataloader._size / bs_valid)
        for j in range(iteration):
            image, label = dataloader.next()
            image_valid.d = image
            pred.forward(clear_buffer=True)
            extracted_feature.append(pred_hidden.d.reshape(len(image), -1))
            preds.append(pred.d)
        concat = np.concatenate(extracted_feature, 0)
        preds = np.concatenate(preds, 0)
        phase_feature[phase] = concat
        phase_preds[phase] = preds
    return phase_feature, phase_preds


def main():
    extension_module = args.context
    ctx = get_extension_context(
        extension_module, device_id=args.device_id, type_config=args.type_config
    )
    nn.set_default_context(ctx)
    input_dir = (
        "./data/input/shuffle" if args.shuffle_label else "./data/input/no_shuffle"
    )

    X_train = np.load(os.path.join(input_dir, "x_train.npy"))
    X_val = np.load(os.path.join(input_dir, "x_val.npy"))
    if args.shuffle_label:
        print("## label shuffled")
        Y_train = np.load(os.path.join(input_dir, "y_shuffle_train.npy"))
    else:
        Y_train = np.load(os.path.join(input_dir, "y_train.npy"))
    Y_val = np.load(os.path.join(input_dir, "y_val.npy"))

    prediction = functools.partial(vgg16_prediction)

    test = True
    image_valid = nn.Variable((bs_valid, 3, 32, 32))
    output, pred_hidden = prediction(image_valid, test)
    pred_hidden.persistent = True
    pred = F.softmax(output, 1)

    save_dir = "shuffle" if args.shuffle_label else "no_shuffle"
    weight_dir = os.path.join("./tmp.monitor", save_dir)
    nn.load_parameters(os.path.join(weight_dir, CHECKPOINTS_PATH_FORMAT))

    affine_weight = []
    last_layer = [name for name in nn.get_parameters().keys()
                  ][-1].split("/")[0]
    for name, param in nn.get_parameters().items():
        if last_layer in name:
            affine_weight.append(param.d)

    data_save_dir = (
        "./data/info/shuffle" if args.shuffle_label else "./data/info/no_shuffle"
    )
    ensure_dir(data_save_dir)

    data_source_train = Cifar10NumpySource(X_train, Y_train)
    train_loader = data_iterator(
        data_source_train, bs_valid, None, False, False)
    data_source_val = Cifar10NumpySource(X_val, Y_val)
    val_loader = data_iterator(data_source_val, bs_valid, None, False, False)
    dataloaders_dict = {"train": train_loader, "test": val_loader}

    phase_feature, phase_output = feature_extractor(
        dataloaders_dict, image_valid, pred_hidden, pred
    )

    with h5py.File(os.path.join(data_save_dir, "info.h5"), "w") as hf:
        hf.create_dataset("label", data=Y_train)
        hf.create_group("param")
        for name, param in zip(["weight", "bias"], affine_weight):
            hf["param"].create_dataset(name, data=param)
        hf.create_group("feature")
        for phase, feature in phase_feature.items():
            print(feature.shape)
            hf["feature"].create_dataset(phase, data=feature)
        hf.create_group("output")
        for phase, output in phase_output.items():
            hf["output"].create_dataset(phase, data=output)

    print("done")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Feature generator")
    parser.add_argument("--shuffle_label", type=strtobool)
    parser.add_argument("--device_id", "-d", type=str, default="0")
    parser.add_argument("--type_config", "-t", type=str, default="float")
    parser.add_argument("--context", "-c", type=str, default="cudnn")
    args = parser.parse_args()
    main()


# responsible_ai/representer_point/mislabel_detection.py

# Copyright (c) 2021 Sony Group Corporation. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import

import matplotlib.pyplot as plt
import numpy as np

score = np.load("./data/info/shuffle/alpha_vgg_nnabla_score.npy")
y_raw = np.load("./data/input/shuffle/y_train.npy")
y_shuffle = np.load("./data/input/shuffle/y_shuffle_train.npy")

sort_ind = np.argsort(score)[::-1]
y_raw_sort = y_raw[sort_ind]
y_shuffle_sort = y_shuffle[sort_ind]

mislabel_length = len(np.where(y_raw != y_shuffle)[0])
print("mislabel: ", mislabel_length)
extracted = np.array([])
count = 0
for raw, shuffle in zip(y_raw_sort, y_shuffle_sort):
    if raw != shuffle:
        count += 1
    extracted = np.append(extracted, count)

detection_rate = extracted / mislabel_length
fraction_rate = [i / len(y_raw) for i in range(len(y_raw))]


fontsize = 16
plt.figure()
plt.plot(fraction_rate, detection_rate)
plt.xlabel("fraction of data checked", fontsize=fontsize)
plt.grid()
plt.ylabel("fraction of detected mislabel", fontsize=fontsize)
plt.savefig("./data/info/shuffle/mislabel_vgg_nnabla_detect.png")


# responsible_ai/representer_point/model.py

# Copyright (c) 2021 Sony Group Corporation. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import nnabla.functions as F
import nnabla.parametric_functions as PF


def categorical_error(pred, label):
    # TODO: Use F.top_n_error
    """
    Compute categorical error given score vectors and labels as
    numpy.ndarray.
    """
    pred_label = pred.argmax(1)
    return (pred_label != label.flat).mean()


def loss_function(pred, label):
    loss = F.mean(F.softmax_cross_entropy(pred, label))
    return loss


def vgg16_prediction(image, test=False, ncls=10, seed=0):
    # Preprocess
    if not test:
        image = F.image_augmentation(
            image, min_scale=0.8, max_scale=1.2, flip_lr=True, seed=seed
        )
        image.need_grad = False
    # Convolution layers
    h = PF.convolution(
        image, 64, (3, 3), pad=(1, 1), stride=(1, 1), name="block1_conv1"
    )
    h = F.relu(h)
    h = PF.convolution(h, 64, (3, 3), pad=(
        1, 1), stride=(1, 1), name="block1_conv2")
    h = F.relu(h)
    h = F.max_pooling(h, (2, 2), stride=(2, 2))

    h = PF.convolution(h, 128, (3, 3), pad=(
        1, 1), stride=(1, 1), name="block2_conv1")
    h = F.relu(h)
    h = PF.convolution(h, 128, (3, 3), pad=(
        1, 1), stride=(1, 1), name="block2_conv2")
    h = F.relu(h)
    h = F.max_pooling(h, (2, 2), stride=(2, 2))

    h = PF.convolution(h, 256, (3, 3), pad=(
        1, 1), stride=(1, 1), name="block3_conv1")
    h = F.relu(h)
    h = PF.convolution(h, 256, (3, 3), pad=(
        1, 1), stride=(1, 1), name="block3_conv2")
    h = F.relu(h)
    h = PF.convolution(h, 256, (3, 3), pad=(
        1, 1), stride=(1, 1), name="block3_conv3")
    h = F.relu(h)
    h = F.max_pooling(h, (2, 2), stride=(2, 2))

    h = PF.convolution(h, 512, (3, 3), pad=(
        1, 1), stride=(1, 1), name="block4_conv1")
    h = F.relu(h)
    h = PF.convolution(h, 512, (3, 3), pad=(
        1, 1), stride=(1, 1), name="block4_conv2")
    h = F.relu(h)
    h = PF.convolution(h, 512, (3, 3), pad=(
        1, 1), stride=(1, 1), name="block4_conv3")
    h = F.relu(h)
    h = F.max_pooling(h, (2, 2), stride=(2, 2))

    h = PF.convolution(h, 512, (3, 3), pad=(
        1, 1), stride=(1, 1), name="block5_conv1")
    h = F.relu(h)
    h = PF.convolution(h, 512, (3, 3), pad=(
        1, 1), stride=(1, 1), name="block5_conv2")
    h = F.relu(h)
    h = PF.convolution(h, 512, (3, 3), pad=(
        1, 1), stride=(1, 1), name="block5_conv3")
    h = F.relu(h)
    hidden = F.max_pooling(h, (2, 2), stride=(2, 2))

    # Fully-Connected layers

    h = PF.affine(hidden, 4096, name="fc1")
    h = F.relu(h)
    h = PF.affine(h, 4096, name="fc2")
    h = F.relu(h)
    pred = PF.affine(h, ncls, name="fc3")

    return pred, h


# responsible_ai/representer_point/train.py

# Copyright (c) 2021 Sony Group Corporation. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from __future__ import absolute_import

import functools
import math
import os
import sys

import nnabla as nn
import nnabla.solvers as S
import nnabla.utils.save as save
import numpy as np
from nnabla.ext_utils import get_extension_context
from nnabla.utils.data_iterator import data_iterator

from args import get_args
from cifar10_load import data_source_cifar10
from model import loss_function, vgg16_prediction

sys.path.append("../../utils/")
from neu.checkpoint_util import load_checkpoint, save_checkpoint
from neu.learning_rate_scheduler import create_learning_rate_scheduler
from neu.save_nnp import save_nnp
from neu.yaml_wrapper import read_yaml


def set_seed(seed=200):
    np.random.seed(seed)
    rng = np.random.seed(seed)
    return rng, seed


def ensure_dir(dir):
    if not os.path.exists(dir):
        os.makedirs(dir)


def categorical_error(pred, label):
    """
    Compute categorical error given score vectors and labels as
    numpy.ndarray.
    """
    pred_label = pred.argmax(1)
    return (pred_label != label.flat).mean()


def train():
    bs_train, bs_valid = args.train_batch_size, args.val_batch_size
    extension_module = args.context
    ctx = get_extension_context(
        extension_module, device_id=args.device_id, type_config=args.type_config
    )
    nn.set_default_context(ctx)

    train_data_source = data_source_cifar10(
        train=True, shuffle=True, label_shuffle=args.shuffle_label
    )
    val_data_source = data_source_cifar10(
        train=False, shuffle=False, label_shuffle=False
    )
    n_train_samples = len(train_data_source.labels)
    n_val_samples = len(val_data_source.labels)
    # Data Iterator
    train_loader = data_iterator(
        train_data_source, bs_train, None, False, False)
    val_loader = data_iterator(val_data_source, bs_valid, None, False, False)

    input_save_dir = (
        "./data/input/shuffle" if args.shuffle_label else "./data/input/no_shuffle"
    )
    if not os.path.exists(input_save_dir):
        os.makedirs(input_save_dir)
    np.save(os.path.join(input_save_dir, "x_train.npy"),
            train_data_source.images)
    np.save(os.path.join(input_save_dir, "y_train.npy"),
            train_data_source.raw_label)
    np.save(os.path.join(input_save_dir, "x_val.npy"), val_data_source.images)
    np.save(os.path.join(input_save_dir, "y_val.npy"), val_data_source.labels)
    if args.shuffle_label:
        np.save(
            os.path.join(input_save_dir, "y_shuffle_train.npy"),
            train_data_source.labels,
        )

    model_prediction = vgg16_prediction
    prediction = functools.partial(model_prediction, ncls=10, seed=args.seed)

    # Create training graphs
    test = False
    image_train = nn.Variable((bs_train, 3, 32, 32))
    label_train = nn.Variable((bs_train, 1))
    pred_train, _ = prediction(image_train, test)
    loss_train = loss_function(pred_train, label_train)

    # Create validation graph
    test = True
    image_valid = nn.Variable((bs_valid, 3, 32, 32))
    label_valid = nn.Variable((bs_valid, 1))
    pred_valid, _ = prediction(image_valid, test)
    loss_val = loss_function(pred_valid, label_valid)

    for param in nn.get_parameters().values():
        param.grad.zero()

    cfg = read_yaml("./learning_rate.yaml")
    print(cfg)
    lr_sched = create_learning_rate_scheduler(cfg.learning_rate_config)
    solver = S.Momentum(momentum=0.9, lr=lr_sched.get_lr())
    solver.set_parameters(nn.get_parameters())
    start_point = 0

    if args.checkpoint is not None:
        # load weights and solver state info from specified checkpoint file.
        start_point = load_checkpoint(args.checkpoint, solver)

    # Create monitor
    from nnabla.monitor import Monitor, MonitorSeries, MonitorTimeElapsed

    monitor = Monitor(monitor_path)
    monitor_loss = MonitorSeries("Training loss", monitor, interval=1)
    monitor_err = MonitorSeries("Training error", monitor, interval=1)
    monitor_time = MonitorTimeElapsed("Training time", monitor, interval=1)
    monitor_verr = MonitorSeries("Test error", monitor, interval=1)
    monitor_vloss = MonitorSeries("Test loss", monitor, interval=1)

    train_iter = math.ceil(n_train_samples / bs_train)
    val_iter = math.ceil(n_val_samples / bs_valid)

    # Training-loop
    for i in range(start_point, args.train_epochs):
        lr_sched.set_epoch(i)
        solver.set_learning_rate(lr_sched.get_lr())
        print("Learning Rate: ", lr_sched.get_lr())
        # Validation
        ve = 0.0
        vloss = 0.0
        print("## Validation")
        for j in range(val_iter):
            image, label, _ = val_loader.next()
            image_valid.d = image
            label_valid.d = label
            loss_val.forward()
            vloss += loss_val.data.data.copy() * bs_valid
            ve += categorical_error(pred_valid.d, label)
        ve /= args.val_iter
        vloss /= n_val_samples

        monitor_verr.add(i, ve)
        monitor_vloss.add(i, vloss)

        if int((i + 1) % args.model_save_interval) == 0:
            # save checkpoint file
            save_checkpoint(monitor_path, i + 1, solver)

        # Forward/Zerograd/Backward
        print("## Training")
        e = 0.0
        loss = 0.0
        for k in range(train_iter):

            image, label, shuffle = train_loader.next()

            image_train.d = image
            label_train.d = shuffle
            loss_train.forward()
            solver.zero_grad()
            loss_train.backward()
            solver.update()
            e += categorical_error(pred_train.d, label_train.d)
            loss += loss_train.data.data.copy() * bs_train
        e /= train_iter
        loss /= n_train_samples

        e = categorical_error(pred_train.d, label_train.d)
        monitor_loss.add(i, loss)
        monitor_err.add(i, e)
        monitor_time.add(i)

    # save_nnp_lastepoch
    contents = save_nnp({"x": image_valid}, {"y": pred_valid}, bs_valid)
    save.save(os.path.join(monitor_path, ("vgg16_result.nnp")), contents)


if __name__ == "__main__":
    """
    $ python train.py --context "cudnn" -b 64

    """
    args = get_args()
    save_dir = "shuffle" if args.shuffle_label else "no_shuffle"
    monitor_path = os.path.join("./tmp.monitor", save_dir)
    ensure_dir(monitor_path)

    train()


# responsible_ai/shap/utils.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import numpy as np
import nnabla as nn
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from nnabla.ext_utils import get_extension_context
from nnabla.utils.load import load


def get_context(device_id):
    # for cli app use
    try:
        context = 'cudnn'
        ctx = get_extension_context(context, device_id=device_id)
    except (ModuleNotFoundError, ImportError):
        context = 'cpu'
        ctx = get_extension_context(context, device_id=device_id)
    # for nnc use
    config_filename = 'net.nntxt'
    if os.path.isfile(config_filename):
        config_info = load([config_filename])
        ctx = config_info.global_config.default_context
    return ctx


def red_blue_map():
    colors = []
    for i in np.linspace(1, 0, 100):
        colors.append((30. / 255, 136. / 255, 229. / 255, i))
    for i in np.linspace(0, 1, 100):
        colors.append((255. / 255, 13. / 255, 87. / 255, i))
    return LinearSegmentedColormap.from_list("red_transparent_blue", colors)


def get_model_layers(model_graph):
    model_layers = dict()
    for k, v in model_graph.variables.items():
        if ("/W" in k) or ("/b" in k):
            continue
        else:
            model_layers[k] = v
    return model_layers


def get_middle_layer(model_layers, interim_layer_index):
    return list(model_layers.values())[interim_layer_index]


def gradient(model_layers, input_layer, target_layer, output_layer, inputs, idx):

    for v in model_layers.values():
        v.grad.zero()
        v.need_grad = True
    selected = output_layer[:, idx]
    input_layer.d = inputs
    selected.forward()
    selected.backward()
    grad = target_layer.g.copy()
    return grad


def get_interim_input(input_layer, middle_layer, output_layer, inputs):
    input_layer.d = inputs
    output_layer.forward()
    middle_layer_d = middle_layer.d.copy()
    return middle_layer_d


def shap(model_graph, X, label, interim_layer_index, num_samples,
         dataset, batch_size, num_epochs=1):
    input_layer = list(model_graph.inputs.values())[0]
    model_layers = get_model_layers(model_graph)
    middle_layer = get_middle_layer(model_layers, interim_layer_index)
    output_layer = list(model_graph.outputs.values())[0]
    # get data
    if len(X.shape) == 3:
        batch_size = 1
    else:
        batch_size = len(X)

    x = nn.Variable((batch_size,) + input_layer.shape)
    # set training True since gradient of variable is necessary for SHAP
    if interim_layer_index == 0:
        data = X.reshape((1,) + X.shape)
    else:
        data = get_interim_input(input_layer, middle_layer, output_layer, X)

    samples_input = [np.zeros((num_samples, ) + X.shape)]
    samples_delta = [np.zeros((num_samples, ) + data.shape[1:])]

    rseed = np.random.randint(0, 1e6)
    np.random.seed(rseed)
    phis = [np.zeros((1,) + data.shape[1:])]

    output_phis = []

    for j in range(num_epochs):
        for k in range(num_samples):
            rind = np.random.choice(len(dataset))
            t = np.random.uniform()
            im = dataset[rind]
            x = X.copy()
            samples_input[j][k] = (t * x + (1 - t) * im.copy()).copy()
            if interim_layer_index == 0:
                samples_delta[j][k] = (x - im.copy()).copy()
            else:
                samples_delta[j][k] = get_interim_input(
                    input_layer, middle_layer, output_layer, samples_input[j][k])[0]

        grads = []

        for b in range(0, num_samples, batch_size):
            batch_last = min(b + batch_size, num_samples)
            batch = samples_input[j][b:batch_last].copy()
            grads.append(gradient(model_layers, input_layer,
                         middle_layer, output_layer, batch, label))
        grad = [np.concatenate([g for g in grads], 0)]
        samples = grad[0] * samples_delta[0]
        phis[0][j] = samples.mean(0)

    output_phis.append(phis[0])
    return output_phis


def visualize(X, output_phis, output, ratio_num=10):
    import matplotlib
    matplotlib.use('Agg')
    img = X.copy()
    height = img.shape[1]
    width = img.shape[2]
    ratio = ratio_num / height
    fig_size = np.array([width * ratio, ratio_num])
    fig, ax = plt.subplots(figsize=fig_size, dpi=1 / ratio)
    shap_plot = output_phis[0][0].sum(0)

    if img.max() > 1:
        img = img / 255.
    if img.shape[0] == 3:
        img_gray = (0.2989 * img[0, :, :] +
                    0.5870 * img[1, :, :] + 0.1140 * img[2, :, :])
    else:
        img_gray = img.reshape(img.shape[1:])

    abs_phis = np.abs(output_phis[0].sum(1)).flatten()
    max_border = np.nanpercentile(abs_phis, 99.9)
    min_border = -np.nanpercentile(abs_phis, 99.9)

    ax.imshow(img_gray, cmap=plt.get_cmap('gray'), alpha=0.15,
              extent=(-1, shap_plot.shape[1], shap_plot.shape[0], -1))
    ax.imshow(shap_plot, cmap=red_blue_map(), vmin=min_border, vmax=max_border)
    ax.axis("off")
    fig.tight_layout()
    fig.savefig(output)
    fig.clf()
    plt.close()


def shap_computation(model_graph, X, label, interim_layer_index, num_samples,
                     dataset, batch_size, output):
    ctx = get_context(0)
    nn.set_default_context(ctx)
    output_phis = shap(model_graph, X, label, interim_layer_index, num_samples,
                       dataset, batch_size)
    visualize(X, output_phis, output)


# responsible_ai/tracin/args.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


def get_args(monitor_path='tmp.monitor',
             model_save_path=None,
             description=None):
    import argparse
    import os
    from distutils.util import strtobool

    if model_save_path is None:
        model_save_path = monitor_path
    if description is None:
        parser = argparse.ArgumentParser(
            description='Training TrackIn Base model')

    parser.add_argument('--monitor-path', '-m', type=str, default=monitor_path)

    parser.add_argument('--model_save_path',
                        '-o',
                        type=str,
                        default=model_save_path,
                        help='Directory to save checkpoints and logs.')

    parser.add_argument('--model_save_interval',
                        type=input,
                        default=1)

    parser.add_argument('--shuffle_label',
                        type=strtobool,
                        default=True,
                        help='whether shuffle trainig label or not')

    parser.add_argument('--augmentation', type=strtobool, default=True)

    parser.add_argument('--model', type=str, choices=['resnet23', 'resnet56'])

    parser.add_argument('--input',
                        type=str,
                        default=None,
                        help='input npy file')

    parser.add_argument('--output',
                        type=str,
                        default=None,
                        help='save directory of shuffled numpy array dataset')

    parser.add_argument('--checkpoint',
                        type=str,
                        default=None,
                        help='if resume training, you put your weight')

    parser.add_argument('--resume', type=int, default=0, help='resume epoch')

    parser.add_argument(
        "--device_id",
        "-d",
        type=str,
        default='0',
        help='Device ID the training run on. This is only valid if you specify `-c cudnn`.'
    )

    parser.add_argument("--type_config",
                        "-t",
                        type=str,
                        default='float',
                        help='Type of computation. e.g. "float", "half".')

    parser.add_argument('--context',
                        '-c',
                        type=str,
                        default='cudnn',
                        help="Extension path. ex) cpu, cudnn.")

    parser.add_argument('--train_batch_size',
                        type=int,
                        default=500,
                        help='Batch size for training (per replica).')

    parser.add_argument('--val-iter', type=int, default=100)

    parser.add_argument('--val_batch_size',
                        type=int,
                        default=100,
                        help='Batch size for eval. (per replica)')

    parser.add_argument('--train_epochs',
                        type=int,
                        default=270,
                        help='Number of epochs to train for.')

    parser.add_argument('--seed',
                        '-s',
                        help='random seed number default=0',
                        default=0,
                        type=int)

    args = parser.parse_args()

    if not os.path.isdir(args.model_save_path):
        os.makedirs(args.model_save_path)

    return args


# responsible_ai/tracin/calculate_score.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import argparse
import functools
import os

import numpy as np
import nnabla as nn
import nnabla.functions as F

from nnabla.utils.data_iterator import data_iterator
from nnabla.ext_utils import get_extension_context

from tqdm import tqdm
from distutils.util import strtobool

from model import resnet23_prediction, resnet56_prediction, loss_function
from cifar10_load import Cifar10NumpySource


CHECKPOINTS_PATH_FORMAT = "params_{}.h5"


def load_ckpt_path():
    checkpoints = [os.path.join(args.checkpoint, CHECKPOINTS_PATH_FORMAT.format(
        str(i))) for i in range(29, 299, 30)]
    return checkpoints


def load_data():
    x = np.load(os.path.join(args.input, 'x_train.npy'))[:50]
    y = np.load(os.path.join(args.input, 'y_train.npy'))[:50]
    y_shuffle = np.load(os.path.join(args.input, 'y_shuffle_train.npy'))[:50]

    data_source = Cifar10NumpySource(x, y, y_shuffle)
    data_num = len(data_source.labels)
    loader = data_iterator(data_source, 1, None, False, False)

    return loader, data_num, x, y, y_shuffle


def calculate_ckpt_score(dataloader, data_num, image_val, label_val, pred_val, hidden, loss_val):
    ckpt_scores = []
    labels, shuffle = [], []

    bar = tqdm(total=data_num)
    for i in range(data_num):
        bar.update(1)
        inputs = dataloader.next()
        for name, param in nn.get_parameters().items():
            param.grad.zero()  # grad initialize
            if 'affine' not in name:
                param.need_grad = False

        grads = []
        images, labels, shuffle = inputs
        image_val.d, label_val.d = images, shuffle

        loss_val.forward()
        loss_val.backward()

        for name, param in nn.get_parameters().items():
            if 'affine' in name:
                grads.append(param.grad)
        grad_mul = [F.sum(grad * grad) for grad in grads]
        score = F.add_n(*grad_mul)
        ckpt_scores.append(score.data)
    return ckpt_scores


def get_scores(dataloader, data_num, ckpt_paths):
    if args.model == 'resnet23':
        model_prediction = resnet23_prediction
    elif args.model == 'resnet56':
        model_prediction = resnet56_prediction

    prediction = functools.partial(model_prediction,
                                   ncls=10,
                                   nmaps=64,
                                   act=F.relu,
                                   seed=args.seed)

    test = True
    image_val = nn.Variable((1, 3, 32, 32))
    label_val = nn.Variable((1, 1))
    pred_val, hidden = prediction(image_val, test)
    loss_val = loss_function(pred_val, label_val)

    infl_scores = np.zeros((data_num, len(ckpt_paths)), dtype=np.float32)
    for ckpt_ind, ckpt_path in enumerate(ckpt_paths):
        epoch = os.path.splitext(os.path.basename(ckpt_path))[0].split('_')[-1]
        print(f'Epoch: {epoch}')
        nn.load_parameters(ckpt_path)

        ckpt_influences = calculate_ckpt_score(
            dataloader, data_num, image_val, label_val, pred_val, hidden, loss_val)
        if args.save_every_epoch:
            np.save(os.path.join(args.output, (epoch+'_influence.npy')),
                    np.array([float(score) for score in ckpt_influences]))

        infl_scores[:, ckpt_ind] = ckpt_influences
    sum_ckpt_scores = infl_scores.sum(axis=-1)

    return {
        'scores': sum_ckpt_scores
    }


def main():
    extension_module = args.context
    ctx = get_extension_context(extension_module,
                                device_id=args.device_id,
                                type_config=args.type_config)
    nn.set_default_context(ctx)

    dataloader, data_num, images, labels, labels_shuffle = load_data()
    ckpt_paths = load_ckpt_path()

    results = get_scores(dataloader, data_num, ckpt_paths)

    np.save(os.path.join(args.output, 'influence_all_epoch.npy'),
            results['scores'])


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='evaluation classification')
    parser.add_argument('--input', type=str)
    parser.add_argument("--checkpoint", type=str)
    parser.add_argument('--output', type=str)
    parser.add_argument('--context', '-c', default='cudnn')
    parser.add_argument('--device-id', type=str, default='0')
    parser.add_argument('--model', type=str, choices=['resnet23', 'resnet56'])
    parser.add_argument('--save_every_epoch', type=strtobool, default=False,
                        help='whether save influence score between every epoch or not')
    parser.add_argument("--type_config",
                        "-t",
                        type=str,
                        default='float',
                        help='Type of computation. e.g. "float", "half".')
    parser.add_argument('--seed',
                        '-s',
                        help='random seed number default=0',
                        default=0,
                        type=int)

    args = parser.parse_args()

    if not os.path.exists(args.output):
        os.makedirs(args.output)

    main()


# responsible_ai/tracin/cifar10_load.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
'''
Provide data iterator for CIFAR10 examples.
'''
import numpy as np
import tarfile
import random

from nnabla.logger import logger
from nnabla.utils.data_source import DataSource
from nnabla.utils.data_source_loader import download


class Cifar10NumpySource(DataSource):
    def _get_data(self, position):
        image = self._images[self._indexes[position]]
        label = self._labels[self._indexes[position]]
        if self._shuffle_labels is not False:
            shuffle = self._shuffle_labels[self._indexes[position]]
            return (image, label, shuffle)
        else:
            return (image, label)

    def __init__(self, x, y, y_shuffled=False, shuffle=False, rng=None):
        self._images = x
        self._labels = y
        if y_shuffled is not False:
            self._shuffle_labels = y_shuffled
        else:
            self._shuffle_labels = False
        self._shuffle = shuffle
        self._size = self._labels.size
        if self._shuffle_labels is not False:
            self._variables = ('x', 'y', 'shuffle')
        else:
            self._variables = ('x', 'y')
        if rng is None:
            rng = np.random.RandomState(313)
        self.rng = rng
        self.reset()

    def reset(self):
        if self._shuffle:
            self._indexes = self.rng.permutation(self._size)
        else:
            self._indexes = np.arange(self._size)
        super(Cifar10NumpySource, self).reset()

    @property
    def images(self):
        """Get copy of whole data with a shape of (N, 1, H, W)."""
        return self._images.copy()

    @property
    def labels(self):
        """Get copy of whole label with a shape of (N, 1)."""
        return self._labels.copy()

    @property
    def shuffled_labels(self):
        """Get copy of whole label with a shape of (N, 1)."""
        if self._shuffle_labels:
            return self._shuffle_labels.copy()


# Copyright (c) 2021 Sony Corporation. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


class Cifar10DataSource(DataSource):
    '''
    Get data directly from cifar10 dataset from Internet(yann.lecun.com).
    '''

    def _get_data(self, position):
        image = self._images[self._indexes[position]]
        label = self._labels[self._indexes[position]]
        return (image, label)

    def __init__(self,
                 train=True,
                 shuffle=False,
                 rng=None,
                 label_shuffle=True,
                 label_shuffle_rate=0.1):
        super(Cifar10DataSource, self).__init__(shuffle=shuffle, rng=rng)

        self._train = train

        data_uri = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"
        logger.info('Getting labeled data from {}.'.format(data_uri))
        r = download(data_uri)  # file object returned
        with tarfile.open(fileobj=r, mode="r:gz") as fpin:
            # Training data
            if train:
                images = []
                labels = []
                for member in fpin.getmembers():
                    if "data_batch" not in member.name:
                        continue
                    fp = fpin.extractfile(member)
                    data = np.load(fp, encoding="bytes", allow_pickle=True)
                    images.append(data[b"data"])
                    labels.append(data[b"labels"])
                self._size = 50000
                self._images = np.concatenate(images).reshape(
                    self._size, 3, 32, 32)

                self._labels = np.concatenate(labels).reshape(-1, 1)
                self.raw_label = self._labels.copy()
                if label_shuffle:
                    self.shuffle_rate = label_shuffle_rate
                    self.label_shuffle()
                    print(f"{self.shuffle_rate*100}% of data was shuffled ")
                    print(len(np.where(self._labels != self.raw_label)[0]))
            # Validation data
            else:
                for member in fpin.getmembers():
                    if "test_batch" not in member.name:
                        continue
                    fp = fpin.extractfile(member)
                    data = np.load(fp, encoding="bytes", allow_pickle=True)
                    images = data[b"data"]
                    labels = data[b"labels"]
                self._size = 10000
                self._images = images.reshape(self._size, 3, 32, 32)
                self._labels = np.array(labels).reshape(-1, 1)
        r.close()
        logger.info('Getting labeled data from {}.'.format(data_uri))

        self._size = self._labels.size
        self._variables = ('x', 'y')
        if rng is None:
            rng = np.random.RandomState(313)
        self.rng = rng
        self.reset()

    def label_shuffle(self):
        num_cls = int(np.max(self._labels)) + 1
        raw_label = self._labels.copy()
        extract_num = int(len(self._labels) * self.shuffle_rate // 10)
        for i in range(num_cls):
            extract_ind = np.where(raw_label == i)[0]
            labels = [j for j in range(num_cls)]
            labels.remove(i)  # candidate of shuffle label
            artificial_label = [
                labels[int(i) % (num_cls-1)] for i in range(int(extract_num))
            ]
            artificial_label = np.array(
                random.sample(artificial_label,
                              len(artificial_label))).astype('float32')
            convert_label = np.array([i for _ in range(len(extract_ind))])
            convert_label[-extract_num:] = artificial_label

            self._labels[extract_ind] = convert_label.reshape(-1, 1)

    def reset(self):
        if self._shuffle:
            self._indexes = self.rng.permutation(self._size)
        else:
            self._indexes = np.arange(self._size)
        super(Cifar10DataSource, self).reset()

    def raw_label(self):
        return self.raw_label.copy()

    @property
    def images(self):
        """Get copy of whole data with a shape of (N, 1, H, W)."""
        return self._images.copy()

    @property
    def labels(self):
        """Get copy of whole label with a shape of (N, 1)."""
        return self._labels.copy()


def data_source_cifar10(train=True,
                        rng=None,
                        shuffle=True,
                        label_shuffle=True):
    '''
    Provide DataIterator with :py:class:`Cifar10DataSource`
    with_memory_cache and with_file_cache option's default value is all False,
    because :py:class:`Cifar10DataSource` is able to store all data into memory.
    '''
    return Cifar10DataSource(train=train,
                             shuffle=shuffle,
                             rng=rng,
                             label_shuffle=label_shuffle)


# responsible_ai/tracin/mislabel_analysis.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import argparse
import cv2
import os

import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

from distutils.util import strtobool
from glob import glob


def accumulate_rates(score, raw, shuffle):
    check_rate = [round(i, 2) for i in np.arange(0, 1.1, 0.01)]
    fraction_rate, data_rate = np.array([]), np.array([])

    indices = np.argsort(-score)
    all_num = len(raw)
    fraction_num = len(np.where(raw != shuffle)[0])

    for threshold in check_rate:
        checked_fraction = 0
        check_num = int(all_num * threshold)
        check_data = indices[:check_num]
        for ind in check_data:
            if raw[ind] != shuffle[ind]:
                checked_fraction += 1
        data_rate = np.append(data_rate, threshold)
        fraction_rate = np.append(fraction_rate,
                                  float(checked_fraction / fraction_num))
    return fraction_rate, data_rate


def save_score_curve(score, raw, shuffle):

    plt.figure()
    fraction_rate, data_rate = accumulate_rates(score, raw, shuffle)
    plt.plot(data_rate, fraction_rate)
    plt.xlabel('Fraction of training data checked')
    plt.ylabel('Fraction of mislabeled identified')
    plt.grid(color='gray', linewidth=1, alpha=0.5)
    plt.savefig(os.path.join(args.output, 'score_curve.png'))
    plt.close()


def save_distribution(score, raw, shuffle):
    correct = score[np.where(raw == shuffle)[0]]
    wrong = score[np.where(raw != shuffle)[0]]
    plt.figure()

    sns.distplot(correct, label='correct_label', color='blue')
    sns.distplot(wrong, label='shuffled_label', color='red')
    plt.legend()
    plt.xlabel('Self Influence')
    plt.savefig(os.path.join(args.output, 'self_influence_distribution.png'))
    plt.close()


def data_load(input_dir):
    images = np.load(os.path.join(input_dir, 'x_train.npy'))
    influence = np.load(os.path.join(input_dir, 'influence_all_epoch.npy'))
    label_raw = np.load(os.path.join(input_dir, 'y_train.npy'))
    label_shuffle = np.load(os.path.join(input_dir, 'y_shuffle_train.npy'))

    return images, influence, label_raw, label_shuffle


def rank_ckpt_influence(raw, shuffle, extract_percentage=0.10):
    '''
    analysis of the trend of influence samples between extracted checkpoints
    '''
    raw, shuffle = np.squeeze(raw), np.squeeze(shuffle)
    labels = ['airplane', 'automobile', 'bird', 'cat',
              'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
    ckpt_influences = sorted(glob(os.path.join(args.input, '*_influence.npy')))
    for ckpt_influence in ckpt_influences:
        counter = np.zeros(10)
        influence = np.load(ckpt_influence)
        extract_num = int(len(influence) * extract_percentage)
        indices = np.argsort(-influence)[:extract_num]
        for indice in indices:
            if raw[indice] != shuffle[indice]:
                counter[raw[indice]] += 1
        epoch = os.path.basename(ckpt_influence).split('_')[0]
        plt.figure(figsize=(8, 5))
        plt.barh(np.arange(len(labels)), counter)
        plt.yticks(np.arange(len(labels)), labels)
        plt.title(f'Epoch:{epoch} top_influence_sample')
        plt.xlabel('Number')
        plt.ylabel('Label')
        plt.tight_layout()
        plt.savefig(os.path.join(
            args.output, f'{epoch}_top_influence_sample.png'))
        plt.close()


def save_extracted_images(images, score, raw, shuffle):
    labels_name = ['airplane', 'automobile', 'bird', 'cat',
                   'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
    iteration = int(len(images)/10)
    check_num = [int(i) for i in np.arange(0, len(images), len(images)/10)]
    for num in check_num:
        # prepare save directory
        correct_dir = os.path.join(
            args.output, 'extracted_samples', str(num+iteration), 'correct_label')
        incorrect_dir = os.path.join(
            args.output, 'extracted_samples', str(num+iteration), 'incorrect_label')
        for dir in [correct_dir, incorrect_dir]:
            if not os.path.exists(dir):
                os.makedirs(dir)

        # save correct / incorrect samples
        indices = np.argsort(-score)[num:num+iteration]
        for ind in indices:
            if raw[ind] != shuffle[ind]:
                cv2.imwrite(os.path.join(incorrect_dir, f'{labels_name[int(raw[ind])]}_({labels_name[int(shuffle[ind])]}).jpg'), np.transpose(
                    images[ind], (1, 2, 0)))
            else:
                cv2.imwrite(os.path.join(correct_dir, f'{labels_name[int(raw[ind])]}_({labels_name[int(shuffle[ind])]}) .jpg'), np.transpose(
                    images[ind], (1, 2, 0)))


def main():
    print('save various figures ...')
    images, score, raw, shuffle = data_load(args.input)
    save_distribution(score, raw, shuffle)
    save_score_curve(score, raw, shuffle)
    if args.ckpt_analysis:
        print('analyse checkpoint...')
        rank_ckpt_influence(raw, shuffle)
    if args.save_extracted:
        print('saving extracted samples ...')
        save_extracted_images(images, score, raw, shuffle)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='evaluation classification')
    parser.add_argument('--input', type=str)
    parser.add_argument('--output', type=str)
    parser.add_argument('--ckpt_analysis', type=strtobool, default=False,
                        help='whether verify the difference of most influence sample trend between checkpoint or not')
    parser.add_argument('--save_extracted', type=strtobool, default=False,
                        help='whethre save the extracted images at certain threshold of self influence score')
    args = parser.parse_args()

    if not os.path.exists(args.output):
        os.makedirs(args.output)
    main()


# responsible_ai/tracin/model.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import nnabla as nn
import nnabla.functions as F
import nnabla.parametric_functions as PF


def categorical_error(pred, label):
    # TODO: Use F.top_n_error
    """
    Compute categorical error given score vectors and labels as
    numpy.ndarray.
    """
    pred_label = pred.argmax(1)
    return (pred_label != label.flat).mean()


def resnet23_prediction(image, test=False, ncls=10, nmaps=64, act=F.relu, seed=0):
    """
    Construct ResNet 23
    """
    # Residual Unit
    def res_unit(x, scope_name, dn=False):
        C = x.shape[1]
        with nn.parameter_scope(scope_name):
            # Conv -> BN -> Nonlinear
            with nn.parameter_scope("conv1"):
                h = PF.convolution(x, C // 2, kernel=(1, 1), pad=(0, 0),
                                   with_bias=False)
                h = PF.batch_normalization(h, batch_stat=not test)
                h = act(h)
            # Conv -> BN -> Nonlinear
            with nn.parameter_scope("conv2"):
                h = PF.convolution(h, C // 2, kernel=(3, 3), pad=(1, 1),
                                   with_bias=False)
                h = PF.batch_normalization(h, batch_stat=not test)
                h = act(h)
            # Conv -> BN
            with nn.parameter_scope("conv3"):
                h = PF.convolution(h, C, kernel=(1, 1), pad=(0, 0),
                                   with_bias=False)
                h = PF.batch_normalization(h, batch_stat=not test)
            # Residual -> Nonlinear
            h = act(F.add2(h, x))
            # Maxpooling
            if dn:
                h = F.max_pooling(h, kernel=(2, 2), stride=(2, 2))
            return h
    # Conv -> BN -> Nonlinear
    with nn.parameter_scope("conv1"):
        # Preprocess
        if not test:
            image = F.image_augmentation(image,
                                         min_scale=0.8,
                                         max_scale=1.2,
                                         flip_lr=True,
                                         seed=seed)
            image.need_grad = False
        h = PF.convolution(image, nmaps, kernel=(3, 3),
                           pad=(1, 1), with_bias=False)
        h = PF.batch_normalization(h, batch_stat=not test)
        h = act(h)

    h = res_unit(h, "conv2", False)    # -> 32x32
    h = res_unit(h, "conv3", True)     # -> 16x16
    h = res_unit(h, "conv4", False)    # -> 16x16
    h = res_unit(h, "conv5", True)     # -> 8x8
    h = res_unit(h, "conv6", False)    # -> 8x8
    h = res_unit(h, "conv7", True)     # -> 4x4
    h = res_unit(h, "conv8", False)    # -> 4x4
    h = F.average_pooling(h, kernel=(4, 4))  # -> 1x1
    pred = PF.affine(h, ncls)

    return pred, h


def resnet56_prediction(image, test=False, ncls=10, nmaps=64, act=F.relu, seed=0):
    """
    Construct ResNet 56
    """

    channels = [16, 32, 64]

    # Residual Unit
    def res_unit(x, scope_name, c, i):
        subsampling = i == 0 and c > 16
        strides = (2, 2) if subsampling else (1, 1)

        with nn.parameter_scope(scope_name):
            # Conv -> BN -> Nonlinear
            with nn.parameter_scope("conv1"):
                h = PF.convolution(x,
                                   c,
                                   kernel=(3, 3),
                                   pad=(1, 1),
                                   stride=strides,
                                   with_bias=False)
                h = PF.batch_normalization(h, batch_stat=not test)
                h = act(h)
            # Conv -> BN -> Nonlinear
            with nn.parameter_scope("conv2"):
                h = PF.convolution(h,
                                   c,
                                   kernel=(3, 3),
                                   pad=(1, 1),
                                   with_bias=False)
                h = PF.batch_normalization(h, batch_stat=not test)

            if subsampling:
                # Conv -> BN
                with nn.parameter_scope("conv3"):
                    x = PF.convolution(x,
                                       c,
                                       kernel=(1, 1),
                                       pad=(0, 0),
                                       stride=(2, 2),
                                       with_bias=False)
            # Residual -> Nonlinear
            h = act(F.add2(h, x))

            return h

    # Conv -> BN -> Nonlinear
    with nn.parameter_scope("conv1"):
        # Preprocess
        if not test:
            image = F.image_augmentation(image,
                                         min_scale=0.8,
                                         max_scale=1.2,
                                         flip_lr=True,
                                         seed=seed)

            image.need_grad = False
        h = PF.convolution(image,
                           channels[0],
                           kernel=(3, 3),
                           pad=(1, 1),
                           with_bias=False)
        h = PF.batch_normalization(h, batch_stat=not test)
        h = act(h)

    for c in channels:
        h = res_unit(h, f"{c}_conv2", c, 0)
        h = res_unit(h, f"{c}_conv3", c, 1)
        h = res_unit(h, f"{c}_conv4", c, 2)
        h = res_unit(h, f"{c}_conv5", c, 3)
        h = res_unit(h, f"{c}_conv6", c, 4)
        h = res_unit(h, f"{c}_conv7", c, 5)
        h = res_unit(h, f"{c}_conv8", c, 6)
        h = res_unit(h, f"{c}_conv9", c, 7)
        h = res_unit(h, f"{c}_conv10", c, 8)

    h = F.global_average_pooling(h)  # -> 1x1
    if test:
        h.need_grad = False
    pred = PF.affine(h, ncls)

    return pred, h


def loss_function(pred, label):
    loss = F.mean(F.softmax_cross_entropy(pred, label))
    return loss


# responsible_ai/tracin/train.py

# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from __future__ import absolute_import
import os
import sys
import functools
import math
import nnabla as nn
import nnabla.functions as F
import nnabla.solvers as S
import nnabla.utils.save as save
import numpy as np

from nnabla.ext_utils import get_extension_context
from nnabla.utils.data_iterator import data_iterator
from args import get_args
from cifar10_load import data_source_cifar10, Cifar10NumpySource
from model import resnet56_prediction, loss_function, resnet23_prediction

sys.path.append("../../utils/")
from neu.learning_rate_scheduler import create_learning_rate_scheduler
from neu.yaml_wrapper import read_yaml
from neu.save_nnp import save_nnp
from neu.checkpoint_util import save_checkpoint, load_checkpoint


def set_seed(seed=200):
    np.random.seed(seed)
    rng = np.random.seed(seed)
    return rng, seed


def load_data(bs_train, bs_valid):
    x_train = np.load(os.path.join(args.input, "x_train.npy"))
    y_train = np.load(os.path.join(args.input, "y_shuffle_train.npy"))

    x_val = np.load(os.path.join(args.input, "x_val.npy"))
    y_val = np.load(os.path.join(args.input, "y_val.npy"))

    data_source_train = Cifar10NumpySource(x_train, y_train, shuffle=False)
    data_source_val = Cifar10NumpySource(x_val, y_val)

    train_samples, val_samples = len(data_source_train.labels), len(
        data_source_val.labels
    )
    train_loader = data_iterator(
        data_source_train, bs_train, None, False, False)
    val_loader = data_iterator(data_source_val, bs_valid, None, False, False)

    return train_loader, val_loader, train_samples, val_samples


def categorical_error(pred, label):
    """
    Compute categorical error given score vectors and labels as
    numpy.ndarray.
    """
    pred_label = pred.argmax(1)
    return (pred_label != label.flat).mean()


def train():
    bs_train, bs_valid = args.train_batch_size, args.val_batch_size
    extension_module = args.context
    ctx = get_extension_context(
        extension_module, device_id=args.device_id, type_config=args.type_config
    )
    nn.set_default_context(ctx)

    if args.input:
        train_loader, val_loader, n_train_samples, n_val_samples = load_data(
            bs_train, bs_valid
        )

    else:
        train_data_source = data_source_cifar10(
            train=True, shuffle=True, label_shuffle=True
        )
        val_data_source = data_source_cifar10(train=False, shuffle=False)
        n_train_samples = len(train_data_source.labels)
        n_val_samples = len(val_data_source.labels)
        # Data Iterator
        train_loader = data_iterator(
            train_data_source, bs_train, None, False, False)
        val_loader = data_iterator(
            val_data_source, bs_valid, None, False, False)

        if args.shuffle_label:
            if not os.path.exists(args.output):
                os.makedirs(args.output)
            np.save(os.path.join(args.output, "x_train.npy"),
                    train_data_source.images)
            np.save(
                os.path.join(args.output, "y_shuffle_train.npy"),
                train_data_source.labels,
            )
            np.save(os.path.join(args.output, "y_train.npy"),
                    train_data_source.raw_label)
            np.save(os.path.join(args.output, "x_val.npy"),
                    val_data_source.images)
            np.save(os.path.join(args.output, "y_val.npy"),
                    val_data_source.labels)

    if args.model == "resnet23":
        model_prediction = resnet23_prediction
    elif args.model == "resnet56":
        model_prediction = resnet56_prediction
    prediction = functools.partial(
        model_prediction, ncls=10, nmaps=64, act=F.relu, seed=args.seed)

    # Create training graphs
    test = False
    image_train = nn.Variable((bs_train, 3, 32, 32))
    label_train = nn.Variable((bs_train, 1))
    pred_train, _ = prediction(image_train, test)

    loss_train = loss_function(pred_train, label_train)

    # Create validation graph
    test = True
    image_valid = nn.Variable((bs_valid, 3, 32, 32))
    label_valid = nn.Variable((bs_valid, 1))
    pred_valid, _ = prediction(image_valid, test)
    loss_val = loss_function(pred_valid, label_valid)

    for param in nn.get_parameters().values():
        param.grad.zero()

    cfg = read_yaml("./learning_rate.yaml")
    print(cfg)
    lr_sched = create_learning_rate_scheduler(cfg.learning_rate_config)
    solver = S.Momentum(momentum=0.9, lr=lr_sched.get_lr())
    solver.set_parameters(nn.get_parameters())
    start_point = 0

    if args.checkpoint is not None:
        # load weights and solver state info from specified checkpoint file.
        start_point = load_checkpoint(args.checkpoint, solver)

    # Create monitor
    from nnabla.monitor import Monitor, MonitorSeries, MonitorTimeElapsed

    monitor = Monitor(args.monitor_path)
    monitor_loss = MonitorSeries("Training loss", monitor, interval=1)
    monitor_err = MonitorSeries("Training error", monitor, interval=1)
    monitor_time = MonitorTimeElapsed("Training time", monitor, interval=1)
    monitor_verr = MonitorSeries("Test error", monitor, interval=1)
    monitor_vloss = MonitorSeries("Test loss", monitor, interval=1)

    # save_nnp
    contents = save_nnp({"x": image_valid}, {"y": pred_valid}, bs_valid)
    save.save(
        os.path.join(args.model_save_path,
                     (args.model+"_epoch0_result.nnp")), contents
    )

    train_iter = math.ceil(n_train_samples / bs_train)
    val_iter = math.ceil(n_val_samples / bs_valid)

    # Training-loop
    for i in range(start_point, args.train_epochs):
        lr_sched.set_epoch(i)
        solver.set_learning_rate(lr_sched.get_lr())
        print("Learning Rate: ", lr_sched.get_lr())
        # Validation
        ve = 0.0
        vloss = 0.0
        print("## Validation")
        for j in range(val_iter):
            image, label = val_loader.next()
            image_valid.d = image
            label_valid.d = label
            loss_val.forward()
            vloss += loss_val.data.data.copy() * bs_valid
            ve += categorical_error(pred_valid.d, label)
        ve /= args.val_iter
        vloss /= n_val_samples

        monitor_verr.add(i, ve)
        monitor_vloss.add(i, vloss)

        if int(i % args.model_save_interval) == 0:
            # save checkpoint file
            save_checkpoint(args.model_save_path, i, solver)

        # Forward/Zerograd/Backward
        print("## Training")
        e = 0.0
        loss = 0.0
        for k in range(train_iter):

            image, label = train_loader.next()
            image_train.d = image
            label_train.d = label
            loss_train.forward()
            solver.zero_grad()
            loss_train.backward()
            solver.update()
            e += categorical_error(pred_train.d, label_train.d)
            loss += loss_train.data.data.copy() * bs_train
        e /= train_iter
        loss /= n_train_samples

        e = categorical_error(pred_train.d, label_train.d)
        monitor_loss.add(i, loss)
        monitor_err.add(i, e)
        monitor_time.add(i)

    nn.save_parameters(
        os.path.join(args.model_save_path, "params_%06d.h5" %
                     (args.train_epochs))
    )

    # save_nnp_lastepoch
    contents = save_nnp({"x": image_valid}, {"y": pred_valid}, bs_valid)
    save.save(os.path.join(args.model_save_path,
              (args.model+"_result.nnp")), contents)


if __name__ == "__main__":
    """
    $ python train.py --context "cudnn" -b 64

    """
    args = get_args()
    train()


# setup.py

import setuptools
from setuptools import setup

setup(
    name='neu',
    version='1.23.0.dev',  # TODO: Must be consistent with nnabla version?
    author='Sony Group Corporation',
    author_email='nnabla@googlegroups.com',
    package_dir={"": "utils"},
    packages=setuptools.find_packages(where="utils"),
    scripts=[],
    license='LICENSE',
    description='Utility package for nnabla enhancement',
    long_description=open('./utils/README.md').read(),
    long_description_content_type="text/markdown",
    url='https://github.com/sony/nnabla-examples/',
    project_urls={
        "nnabla": "https://nnabla.org/",
    },
    classifiers=[
        "Programming Language :: Python :: 3",
        "License :: OSI Approved :: Apache Software License",
        "Operating System :: OS Independent",
    ],
    entry_points={
        'console_scripts': [
            'neu_cli = neu.cli:main',
        ],
    },
    python_requires=">=3.6",
    install_requires=[
        "pyyaml",
        "dominate",
    ],
)


# responsible_ai/roc_pr_curve/utils.py
# Copyright 2022 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import sys
import numpy as np
import pandas as pd
import opendatasets as od
from urllib.request import urlopen
from io import BytesIO
from zipfile import ZipFile
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
from sklearn.metrics import precision_recall_curve, roc_curve, auc


def load_german_data():
    """
    Load and preprocess german credit card dataset and returns the feature, target, and sensitive attribute.

    This dataset classifies people described by a set of attributes as good or bad credit risks.
    (For more info : https://archive.ics.uci.edu/ml/datasets/Statlog+%28German+Credit+Data%29)
    Returns:
        features (pandas.DataFrame): features of the german credit card dataset
        target (pandas.Series): target values of the german credit card dataset
        sensitive_attributes(pandas.DataFrame): sensitive attributes values of german credit card dataset
    """
    file_path = r'https://archive.ics.uci.edu/ml/machine-learning-databases/statlog/german/german.data'
    column_names = ['status', 'month', 'credit_history',
                    'purpose', 'credit_amount', 'savings', 'employment',
                    'investment_as_income_percentage', 'personal_status',
                    'other_debtors', 'residence_since', 'property', 'age',
                    'installment_plans', 'housing', 'number_of_credits',
                    'skill_level', 'people_liable_for', 'telephone',
                    'foreign_worker', 'credit']

    dataframe = pd.read_csv(
            file_path, sep=' ', header=None, names=column_names)
    # derive the gender attribute from personal_status (you can refer the german.doc)
    status_map = {'A91': 'male', 'A93': 'male', 'A94': 'male',
                  'A92': 'female', 'A95': 'female'}
    dataframe['gender'] = dataframe['personal_status'].replace(status_map)

    # targets; 1 when someone makes good credit , otherwise 0
    target = (dataframe['credit'] == 1).astype(int)
    # features; note that the 'target' and sensitive attribute columns are dropped
    features = (dataframe
                .drop(columns=['credit', 'age', 'gender', 'personal_status'])
                .fillna('Unknown')
                .pipe(pd.get_dummies, drop_first=True))
    print("features sample", features.shape[0])
    print("features attributes", features.shape[1])
    print("target samples", target.shape[0])
    return features, target


def load_adult_data():
    """
    Load and preprocess Census Income dataset and returns the feature, target, and sensitive attribute.

    This dataset is used to predict whether income exceeds $50K/yr
    (For more info : https://archive.ics.uci.edu/ml/datasets/adult)

    Returns:
        features (pandas.DataFrame): features of the Census income dataset
        target (pandas.Series): target values of the Census income dataset
        sensitive_attributes(pandas.DataFrame): sensitive attributes values of Census income dataset
    """
    file_path = r'https://archive.ics.uci.edu/ml/machine-learning-databases/adult/adult.data'
    column_names = ['age', 'workclass', 'fnlwgt', 'education', 'education_num',
                    'marital_status', 'occupation', 'relationship', 'race', 'sex',
                    'capital_gain', 'capital_loss', 'hours_per_week', 'country', 'target']
    try:
        # Restrict races to White and Black
        input_data = (pd.read_csv(file_path, names=column_names,
                                  na_values="?", sep=r'\s*,\s*', engine='python').loc[lambda df: df['race'].isin(['White', 'Black'])]).dropna()
    except Exception as e:
        print("Error : ", e)
        sys.exit(1)

    # targets; 1 when someone makes over 50k , otherwise 0
    target = (input_data['target'] == '>50K').astype(int)

    # features; note that the 'target' and sensitive attribute columns are dropped
    features = (input_data
                .drop(columns=['target', 'race', 'sex'])
                .pipe(pd.get_dummies, drop_first=True))
    print("features sample", features.shape[0])
    print("features attributes", features.shape[1])
    print("target samples", target.shape[0])
    return features, target


def roc_curve_value(y_label, y_pred, thresholds):
    '''
    calculate TPR and FPR
    Args:
        y_label     (list) : ground truth target values
        y_pred      (list) : estimated targets as returned by a classifier
        thresholds  (array): threshold values

    return:
        fpr (list): False Positive Rate
        tpr (list): True Positive Rate
    '''
    fpr = []
    tpr = []
    positive = sum(y_label)
    negative = len(y_label) - positive
    for thresh in thresholds:
        fp_count = 0
        tp_count = 0
        length = len(y_pred)
        for i in range(length):
            if y_pred[i] >= thresh:
                if y_label[i] == 1:
                    tp_count = tp_count + 1
                if y_label[i] == 0:
                    fp_count = fp_count + 1
        fpr.append(fp_count/float(negative))
        tpr.append(tp_count/float(positive))
    return fpr, tpr


def auc_score(x, y):
    '''
    Compute Area Under Curve
    Args:
        x (list): computed x values
        y (list): computed y values
    Return:
        auc_score (float): auc_score
    '''
    height = 0.5*(np.array(y[1:])+np.array(y[:-1]))
    width = -((np.array(x[1:])-np.array(x[:-1])))
    auc_score = (height*width).sum()
    return auc_score


def calc_precision_recall(y_label, y_pred):
    '''
    Calculate precission and recall
    Args:
        y_label     (list) : ground truth target values
        y_pred      (list) : estimated targets as returned by a classifier

    return:
        precision (float): precision value
        recall    (float): recall value
    '''

    tp_count = 0
    fp_count = 0
    fn_count = 0
    length = len(y_label)
    for i in range(length):
        if y_label[i] == y_pred[i] == 1:
            tp_count += 1
        if y_pred[i] == 1 and y_label[i] != y_pred[i]:
            fp_count += 1
        if y_pred[i] == 0 and y_label[i] != y_pred[i]:
            fn_count += 1
    try:
        precision = tp_count / (tp_count + fp_count)
    except:
        precision = 1
    try:
        recall = tp_count / (tp_count + fn_count)
    except:
        recall = 1
    return precision, recall


def cal_scores(thresholds, y_test_probs, y_test):
    '''
    calculate precision scores, recall scores
    Args:
        thresholds   (float): threshold values
        y_test_probs (list): estimated targets as returned by a classifier
        y_test       (list): target values

    Return:
        precision_scores (list): computed precision values
        recall_scores    (list): computed recall values
    '''
    precision_scores = []
    recall_scores = []
    for thresh in thresholds:
        y_test_preds = []
        y_test_preds = [1 if prob > thresh else 0 for prob in y_test_probs]
        precision, recall = calc_precision_recall(y_test, y_test_preds)
        precision_scores.append(precision)
        recall_scores.append(recall)
    return precision_scores, recall_scores


def get_roc_curve(fpr, tpr, auc):
    """
    plot ROC curve
    """
    plt.rcParams["figure.figsize"] = (10, 8)
    plt.title('Receiver Operating Characteristic curve')
    plt.plot([0, 1], ls="--", color='tan')
    plt.plot([0, 0], [1, 0], c=".7"), plt.plot([1, 1], c=".7")
    plt.plot(fpr, tpr, color='blue', label="ROC curve (area = %0.2f)" % auc)
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.legend(loc=4),
    plt.savefig('ROC_curve.png')
    plt.show()


def get_pr_curve(precision_scores, recall_scores, Y_test, auc):
    """
    plot precision Recall curve
    """
    positive = sum(Y_test)
    baseline = positive / len(Y_test)
    plt.rcParams["figure.figsize"] = (10, 8)
    plt.title('Precision / Recall curve')
    plt.plot([0, 1], [baseline, baseline], linestyle='--',
             color='tan', label='Baseline')
    plt.plot(recall_scores, precision_scores, color='blue',
             label="PR curve (area = %0.2f)" % auc)
    plt.xlabel('Recall')
    plt.ylabel('Precision')
    plt.legend(loc='center left')
    plt.savefig('Precision_Recall_curve.png')
    plt.show()


def get_roc_curve_multiclass(target_label, predicted_label, n_classes, labels):
    """
    plot ROC curve for multiclassification
    """
    fpr = dict()
    tpr = dict()
    roc_auc = dict()
    plt.rcParams["figure.figsize"] = (10, 8)
    for i in range(n_classes):
        fpr[i], tpr[i], _ = roc_curve(
            target_label[:, i], predicted_label[:, i])
        roc_auc[i] = auc(fpr[i], tpr[i])
        plt.plot([0, 1], ls="--", color='tan')
        plt.plot([0, 0], [1, 0], c=".7")
        plt.plot([1, 1], c=".7")
        plt.plot(fpr[i], tpr[i], label='ROC curve (area = %0.2f) for label %i %s' % (
            roc_auc[i], i, labels[i]))
    plt.title('Receiver operating characteristic curve')
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.legend(loc=4)
    plt.savefig('roc_auc_multiclass.png')
    plt.show()


def get_pr_curve_multiclass(target_label, predicted_label, n_classes, labels):
    """
    plot precision Recall curve for multiclassification
    """
    precision = dict()
    recall = dict()
    pr_auc = dict()
    plt.rcParams["figure.figsize"] = (10, 8)
    for i in range(n_classes):
        precision[i], recall[i], _ = precision_recall_curve(
            target_label[:, i], predicted_label[:, i])
        pr_auc[i] = auc(recall[i], precision[i])
        plt.plot(recall[i], precision[i], label='PR curve (area = %0.2f) for label %i %s' % (
            pr_auc[i], i, labels[i]))
    plt.title('Precision / Recall curve ')
    plt.xlabel('Recall')
    plt.ylabel('Precision')
    plt.legend(loc=3)
    plt.savefig('pr_auc_multiclass.png')
    plt.show()
