rop_image(img, l, t, r, b):
    w = r - l + 1
    h = b - t + 1
    M = np.array([[1, 0, l], [0, 1, t]], dtype=np.float32)
    return cv2.warpAffine(img, M, (w, h), flags=cv2.INTER_NEAREST, borderMode=cv2.BORDER_REPLICATE)

    # This was slower.
    # ind_x = np.clip(np.arange(l, r + 1), 0, img.shape[1] - 1)
    # ind_y = np.clip(np.arange(t, b + 1), 0, img.shape[0] - 1)
    # return img[np.ix_(ind_y, ind_x)]


def resize_image(img, shape):
    return cv2.resize(img, shape, interpolation=cv2.INTER_NEAREST)


def crop_resize_image(img, l, t, r, b, shape, flip):
    w, h = shape
    pts1 = np.float32([[l, t], [r, t], [l, b], [r, b]])
    pts2 = np.float32([[0, 0], [w - 1, 0], [0, h - 1], [w - 1, h - 1]])
    # This was wrong
    # pts1 = np.float32([[l, t], [r + 1, t], [l, b + 1], [r + 1, b + 1]])
    # pts2 = np.float32([[0, 0], [w, 0], [0, h], [w, h]])
    if flip:
        pts2[:, 0] = pts2[::-1, 0]
    M = cv2.getPerspectiveTransform(pts1, pts2)
    return cv2.warpPerspective(img, M, shape, flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)


def data_augmentation(img, shape, jitter, hue, saturation, exposure):
    oh, ow = img.shape[:2]

    dw = int(ow*jitter)
    dh = int(oh*jitter)

    pleft = random.randint(-dw, dw)
    pright = random.randint(-dw, dw)
    ptop = random.randint(-dh, dh)
    pbot = random.randint(-dh, dh)

    swidth = ow - pleft - pright
    sheight = oh - ptop - pbot

    sx = float(swidth) / ow
    sy = float(sheight) / oh

    flip = random.randint(1, 10000) % 2

    dx = (float(pleft)/ow)/sx
    dy = (float(ptop) / oh)/sy

    # cropped = crop_image(img, pleft, ptop, pleft +
    #                      swidth - 1, ptop + sheight - 1)
    # sized = resize_image(cropped, shape)

    # if flip:
    #     sized = sized[:, ::-1]

    sized = crop_resize_image(img, pleft, ptop, pleft +
                              swidth - 1, ptop + sheight - 1, shape,
                              flip)

    img = random_distort_image(sized, hue, saturation, exposure)

    return img, flip, dx, dy, sx, sy


def fill_truth_detection(labpath, w, h, flip, dx, dy, sx, sy):
    max_boxes = 50
    label = np.zeros((max_boxes, 5))
    if labpath is not None or (isinstance(labpath, str) and os.path.getsize(labpath)):
        if isinstance(labpath, str):
            bs = np.loadtxt(labpath)
        else:
            if labpath is None:
                return label
            bs = labpath.copy()
        if bs is None:
            return label
        bs = np.reshape(bs, (-1, 5))
        cc = 0
        for i in range(bs.shape[0]):
            x1 = bs[i][1] - bs[i][3]/2
            y1 = bs[i][2] - bs[i][4]/2
            x2 = bs[i][1] + bs[i][3]/2
            y2 = bs[i][2] + bs[i][4]/2

            x1 = min(0.999, max(0, x1 * sx - dx))
            y1 = min(0.999, max(0, y1 * sy - dy))
            x2 = min(0.999, max(0, x2 * sx - dx))
            y2 = min(0.999, max(0, y2 * sy - dy))

            bs[i][1] = (x1 + x2)/2
            bs[i][2] = (y1 + y2)/2
            bs[i][3] = (x2 - x1)
            bs[i][4] = (y2 - y1)

            if flip:
                bs[i][1] = 0.999 - bs[i][1]

            if bs[i][3] < 0.001 or bs[i][4] < 0.001:
                continue
            label[cc] = bs[i]
            cc += 1
            if cc >= 50:
                break

    label = np.reshape(label, (-1))
    return label


def load_data_detection(imgpath, labpath, shape, jitter, hue, saturation, exposure):

    # data augmentation
    if isinstance(imgpath, str):
        img = cv2.imread(imgpath)
    else:
        img = imgpath
    img, flip, dx, dy, sx, sy = data_augmentation(
        img, shape, jitter, hue, saturation, exposure)
    label = fill_truth_detection(
        labpath, img.shape[1], img.shape[0], flip, dx, dy, 1./sx, 1./sy)
    return img[..., ::-1], label


# object-detection/yolov2/region_loss_utils.py

# Copyright 2018,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.


# This file was forked from https://github.com/marvis/pytorch-yolo2 ,
# licensed under the MIT License (see LICENSE.external for more details).


from utils import *
import numpy as np


def build_targets_numpy(pred_boxes, target, anchors, num_anchors,
                        num_classes, nH, nW, coord_scale, noobject_scale,
                        object_scale, cls_scale, sil_thresh, seen):
    nB = target.shape[0]
    nA = num_anchors
    nC = num_classes
    anchor_step = len(anchors)//num_anchors

    # The mask arrays `coord_mask`, `conf_mask` and `cls_mask` indicate not only
    # whether the specific bounding box is evaluated or not, but also are
    # multiplied by loss coefficients `coord_scale`, `noobject_scale`,
    # `object_scale` and `cls_scale`.
    conf_mask = np.ones((nB, nA, nH, nW), dtype=np.float32) * noobject_scale
    coord_mask, cls_mask, tx, ty, tw, th, tconf, tcls = [
        np.zeros((nB, nA, nH, nW), dtype=np.float32) for _ in range(8)]

    nAnchors = nA*nH*nW
    nPixels = nH*nW
    for b in range(nB):
        cur_pred_boxes = np.transpose(pred_boxes[b*nAnchors:(b+1)*nAnchors])
        cur_ious = np.zeros(nAnchors, dtype=np.float32)
        for t in range(50):
            if target[b][t*5+1] == 0:
                break
            gx = target[b][t*5+1]*nW
            gy = target[b][t*5+2]*nH
            gw = target[b][t*5+3]*nW
            gh = target[b][t*5+4]*nH
            cur_gt_boxes = np.array([[gx, gy, gw, gh]], dtype=np.float32)
            cur_gt_boxes = np.repeat(cur_gt_boxes, nAnchors, axis=0)
            cur_gt_boxes = np.transpose(cur_gt_boxes)
            cur_ious = np.maximum(cur_ious, bbox_ious_numpy(
                cur_pred_boxes, cur_gt_boxes, x1y1x2y2=False))
        conf_mask_b = conf_mask[b]
        conf_mask_b_orig_shape = conf_mask_b.shape
        conf_mask_b = np.reshape(conf_mask_b, -1)
        conf_mask_b[cur_ious > sil_thresh] = 0
        conf_mask_b = np.reshape(conf_mask_b, conf_mask_b_orig_shape)
        conf_mask[b] = conf_mask_b
    if seen < 12800:
        if anchor_step == 4:
            tx = anchors.reshape((nA, anchor_step))[
                :, 2].reshape((1, nA, 1, 1))
            for t_axis, t_n in zip([0, 2, 3], [nB, nH, nW]):
                tx = np.repeat(tx, t_n, axis=t_axis)
            ty = anchors.reshape((num_anchors, anchor_step))[
                :, 2].reshape((1, nA, 1, 1))
            for t_axis, t_n in zip([0, 2, 3], [nB, nH, nW]):
                ty = np.repeat(tx, t_n, axis=t_axis)*2
        else:
            tx.fill(0.5)
            ty.fill(0.5)
        tw.fill(0.0)
        th.fill(0.0)
        coord_mask.fill(0.01)

    nGT = 0
    nCorrect = 0
    avg_iou = 0.0
    count = 0
    for b in range(nB):
        for t in range(50):
            if target[b][t*5+1] == 0:
                break
            nGT = nGT + 1
            best_iou = 0.0
            best_n = -1
            min_dist = 10000
            gx = target[b][t*5+1] * nW
            gy = target[b][t*5+2] * nH
            gi = int(gx)
            gj = int(gy)
            truthw = target[b][t*5+3]
            truthh = target[b][t*5+4]
            gw = truthw * nW
            gh = truthh * nH
            gt_box = [0, 0, gw, gh]
            for n in range(nA):
                aw = anchors[anchor_step*n]
                ah = anchors[anchor_step*n+1]
                anchor_box = [0, 0, aw, ah]
                iou = bbox_iou_numpy(anchor_box, gt_box, x1y1x2y2=False)
                if anchor_step == 4:
                    ax = anchors[anchor_step*n+2]
                    ay = anchors[anchor_step*n+3]
                    dist = pow(((gi+ax) - gx), 2) + pow(((gj+ay) - gy), 2)
                if iou > best_iou:
                    best_iou = iou
                    best_n = n
                elif anchor_step == 4 and iou == best_iou and dist < min_dist:
                    best_iou = iou
                    best_n = n
                    min_dist = dist

            gt_box = [gx, gy, gw, gh]
            pred_box = pred_boxes[b*nAnchors+best_n*nPixels+gj*nW+gi]

            coord_mask[b][best_n][gj][gi] = coord_scale * (2 - truthw * truthh)
            cls_mask[b][best_n][gj][gi] = cls_scale
            conf_mask[b][best_n][gj][gi] = object_scale
            tx[b][best_n][gj][gi] = target[b][t*5+1] * nW - gi
            ty[b][best_n][gj][gi] = target[b][t*5+2] * nH - gj
            tw[b][best_n][gj][gi] = np.log(gw/anchors[anchor_step*best_n])
            th[b][best_n][gj][gi] = np.log(gh/anchors[anchor_step*best_n+1])
            iou = bbox_iou_numpy(gt_box, pred_box, x1y1x2y2=False)  # best_iou
            tconf[b][best_n][gj][gi] = iou
            tcls[b][best_n][gj][gi] = target[b][t*5]
            if iou > 0.5:
                nCorrect = nCorrect + 1
            avg_iou += iou
            count += 1
    return nGT, nCorrect, avg_iou / count, coord_mask[:, :, None], conf_mask[:, :, None], cls_mask[:, :, None], np.stack((tx, ty, tw, th), axis=2), tconf[:, :, None], tcls[:, :, None]


# object-detection/yolov2/train.py

# Copyright 2018,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.


# This file was forked from https://github.com/marvis/pytorch-yolo2 ,
# licensed under the MIT License (see LICENSE.external for more details).


from __future__ import division

import dataset
from utils import *
from arg_utils import Yolov2OptionTraining
import os

import nnabla as nn
import nnabla.solvers as S

from train_graph import TrainGraph

import random
random.seed(0)
import numpy as np
np.seterr(**{key: 'raise' for key in ['divide', 'over', 'invalid']})


def get_current_batch(seen, args):
    '''
    Get the current batch number from number of processed examples.
    '''
    return int(seen // (args.batch_size * args.accum_times))


class YoloSolver(object):
    def __init__(self, args):

        # Create solvers (only Convolution kernels require weight decay).
        param_convweights = {
            k: v for k, v in nn.get_parameters().items() if k.endswith("conv/W")}
        param_others = {k: v for k, v in nn.get_parameters().items()
                        if not k.endswith("conv/W")}
        convweights = S.Momentum(args.learning_rate, args.momentum)
        others = S.Momentum(args.learning_rate, args.momentum)
        convweights.set_parameters(param_convweights)
        others.set_parameters(param_others)

        # Init parameter gradients.
        convweights.zero_grad()
        others.zero_grad()

        # Set attributes.
        self.convweights = convweights
        self.others = others
        self.args = args
        self.batch_size = args.batch_size * args.accum_times
        self.rate = args.accum_times
        self.count = 0

    def get_current_lr(self, seen):
        '''
        Get the current learning rate by the current number of processed examples.

        Burn-in and step learning rate scheduling are employed.
        '''
        args = self.args
        batch_num = get_current_batch(seen, args)
        lr = float(args.learning_rate)

        # Burn-in (burn_in=1000, burn_in_power=4 by default)
        if batch_num < args.burn_in:
            return lr * ((float(batch_num) / args.burn_in) ** args.burn_in_power)

        # Steps is only implemented
        for step, scale in zip(args.steps, args.scales):
            if step > batch_num:
                return lr
            lr *= scale
        return lr

    def update_at_rate(self, lr, force=False):
        '''

        Returns: int
            Current learning rate
        '''
        if force and self.count > 0:
            self.update(lr)
            return True
        self.count = (self.count + 1) % self.rate
        if self.count == 0:
            self.update(lr)
            return True
        return False

    def update(self, lr):
        # Get current learning rate and set to solvers
        self.convweights.set_learning_rate(lr / self.batch_size)
        self.others.set_learning_rate(lr / self.batch_size)

        # Update
        self.convweights.weight_decay(self.args.decay * self.batch_size)
        self.convweights.update()
        self.others.update()

        # Reset parameter gradients and accumulation count
        self.convweights.zero_grad()
        self.others.zero_grad()
        self.count = 0


def train(args, sample_iter, nsamples, epoch, max_epochs, train_graph, yolo_solver, on_memory_data, total_batches):

    total_loss = []
    epoch_seen = 0
    tic = time.time()
    for each_batch in range(0, int(total_batches)):
        data_tensor, target_tensor = sample_iter.next()
        lr = yolo_solver.get_current_lr(train_graph.seen)
        nB = data_tensor.shape[0]
        print('size={}, '.format(data_tensor.shape[2]), end='')
        stats = train_graph.forward_backward(data_tensor, target_tensor)
        print('s=%d/%d, b=%d/%d, e=%d/%d, lr %g, mIOU %.3f, nGT %d, recall %d,'
              'proposals %d, loss: %.3f, exec: %.1f [ms]' % (
                    (stats.seen % nsamples), nsamples,
                    get_current_batch(stats.seen, args), args.max_batches,
                    epoch, max_epochs, lr, stats.mIoU, stats.nGT,
                    stats.nCorrect, stats.nProposals, stats.loss, stats.time))
        epoch_seen += nB
        total_loss.append(float(stats.loss))
        # Update parameters for every accum_times iterations.
        updated = yolo_solver.update_at_rate(lr)

    # Update if gradients are computed from previous update.
    yolo_solver.update_at_rate(lr, force=True)
    if (epoch+1) % args.save_interval == 0:
        logging('save weights to %s/%06d.h5' % (args.output, epoch+1))
        nn.save_parameters('%s/%06d.h5' % (args.output, epoch+1))
    return np.sum(total_loss) / epoch_seen


def main():
    # Training settings
    args = Yolov2OptionTraining().parse_args()

    nsamples = file_lines(args.train)

    set_default_context_by_args(args)

    # Training parameters
    max_epochs = args.max_batches * args.batch_size * args.accum_times / nsamples + 1

    # total number of batches
    total_batches = (nsamples+args.batch_size-1)//args.batch_size

    if not os.path.exists(args.output):
        os.mkdir(args.output)

    ###############

    # Load parameters
    print("Load", args.weight, "...")
    if args.fine_tune:
        nn.load_parameters(args.weight)
        nn.parameter.pop_parameter("detection/conv/W")
        nn.parameter.pop_parameter("detection/conv/b")
    else:
        nn.load_parameters(args.weight)

    train_graph = TrainGraph(args, (args.size_aug[-1], args.size_aug[-1]))
    yolo_solver = YoloSolver(args)

    if args.on_memory_data:
        on_memory_data = dataset.load_on_memory_data(args.train)
    else:
        on_memory_data = None

    from nnabla.monitor import Monitor, MonitorSeries, MonitorTimeElapsed
    monitor = Monitor(args.output)
    monitor_loss = MonitorSeries('Training loss', monitor, interval=1)
    monitor_time = MonitorTimeElapsed('Time per epoch', monitor, interval=1)

    sample_iter = dataset.data_iterator_yolo(
        args.train, args,
        shuffle=True,
        train=True,
        image_sizes=args.size_aug,
        image_size_change_freq=args.batch_size * args.accum_times * 10,
        on_memory_data=on_memory_data,
        use_cv2=not args.disable_cv2, batch_size=args.batch_size)

    # Epoch loop
    for epoch in range(0, int(max_epochs)):
        loss = train(args, sample_iter, nsamples, epoch, max_epochs, train_graph,
                     yolo_solver, on_memory_data, total_batches)
        monitor_loss.add(epoch, loss)
        monitor_time.add(epoch)

    # Save the final parameters
    logging('save weights to %s/%06d.h5' % (args.output, epoch+1))
    nn.save_parameters('%s/%06d.h5' % (args.output, epoch+1))


if __name__ == '__main__':
    main()


# object-detection/yolov2/train_graph.py

# Copyright 2018,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 time
import collections

from utils import *
from region_loss_utils import *
import nnabla as nn
import nnabla.functions as F
import nnabla_ext.cuda
import yolov2

from nnabla.function import PythonFunction


class RegionLossTargets(PythonFunction):
    def __init__(self, num_classes, anchors, seen, coord_scale=1.0, noobject_scale=1.0, object_scale=5.0, class_scale=1.0, thresh=0.6):
        '''
        Args:
            thresh (float): boxes with best IoU<=``thresh`` are considered as
                negative samples wrt objectness.

        Returns:

        '''
        self.num_classes = num_classes
        self.anchors = anchors
        self.seen = seen
        self.coord_scale = coord_scale
        self.noobject_scale = noobject_scale
        self.object_scale = object_scale
        self.class_scale = class_scale
        self.thresh = thresh

        self.prev_miou = None
        self.prev_ngt = None
        self.prev_ncorrect = None

    @property
    def name(self):
        return 'PythonRegionLossTargets'

    def min_outputs(self):
        return 6

    def setup_impl(self, inputs, outputs):
        bbox_pred = inputs[0]
        nB, nA, _, nH, nW = bbox_pred.shape
        tcoord, mcoord, tconf, mconf, tcls, mcls = outputs
        mask_shape = (nB, nA, 1, nH, nW)
        tcoord.reset_shape(bbox_pred.shape, True)
        mcoord.reset_shape(mask_shape, True)
        tconf.reset_shape(mask_shape, True)
        mconf.reset_shape(mask_shape, True)
        tcls.reset_shape(mask_shape, True)
        mcls.reset_shape(mask_shape, True)

    def forward_impl(self, inputs, outputs):
        bbox, target = inputs
        nB, nA, _, nH, nW = bbox.shape
        nC = self.num_classes
        anchor_step = len(self.anchors) // nA

        grid_x = np.arange(0, nW)[None, None, None, None, :]
        grid_y = np.arange(0, nH)[None, None, None, :, None]

        pred_boxes = np.zeros((4, nB*nA*nH*nW))
        bb = bbox.data.get_data('r')
        pred_boxes[0, :] = (bb[:, :, 0] + grid_x).flat
        pred_boxes[1, :] = (bb[:, :, 1] + grid_y).flat
        anchors = np.array(self.anchors).reshape(-1, 2)  # nA * 2
        # clip without out buffer is sometimes slow
        tmpbuff = np.empty_like(bb[:, :, 2])
        tmpbuff[...] = bb[:, :, 2]
        tmpbuff[tmpbuff > 5] = 5
        pred_boxes[2, :] = (
            np.exp(tmpbuff) * anchors[None, :, 0, None, None]).flat
        tmpbuff[...] = bb[:, :, 3]
        tmpbuff[tmpbuff > 5] = 5
        pred_boxes[3, :] = (
            np.exp(tmpbuff) * anchors[None, :, 1, None, None]).flat
        pred_boxes = np.transpose(pred_boxes, (1, 0)).reshape(-1, 4)
        # tcoord, mcoord, tconf, mconf, tcls, mcls = outputs
        o = [v.data.get_data('w') for v in outputs]
        nGT, nCorrect, mIoU, o[1][...], o[3][...], o[5][...], o[0][...], o[2][...], o[4][...] = build_targets_numpy(
            pred_boxes, target.d, self.anchors, nA, nC, nH, nW,
            self.coord_scale, self.noobject_scale, self.object_scale,
            self.class_scale, self.thresh, self.seen)
        self.seen += nB
        self.prev_ngt = nGT
        self.prev_ncorrect = nCorrect
        self.prev_miou = mIoU

    def backward_impl(self, inputs, outputs, propagate_down, accum):
        pass


def create_network(batchsize, imheight, imwidth, args, seen):
    import gc
    gc.collect()
    nnabla_ext.cuda.clear_memory_cache()

    anchors = args.num_anchors
    classes = args.num_classes
    yolo_x = nn.Variable((batchsize, 3, imheight, imwidth))
    target = nn.Variable((batchsize, 50 * 5))
    yolo_features = yolov2.yolov2(yolo_x, anchors, classes, test=False)

    nB = yolo_features.shape[0]
    nA = args.num_anchors
    nC = args.num_classes
    nH = yolo_features.shape[2]
    nW = yolo_features.shape[3]

    # Bouding box regression loss
    # pred.shape = [nB, nA, 4, nH, nW]
    output = F.reshape(yolo_features, (nB, nA, (5 + nC), nH, nW))
    xy = F.sigmoid(output[:, :, :2, ...])
    wh = output[:, :, 2:4, ...]
    bbox_pred = F.concatenate(xy, wh, axis=2)
    conf_pred = F.sigmoid(output[:, :, 4:5, ...])
    cls_pred = output[:, :, 5:, ...]

    region_loss_targets = RegionLossTargets(
        nC, args.anchors, seen, args.coord_scale, args.noobject_scale,
        args.object_scale, args.class_scale, args.thresh)

    tcoord, mcoord, tconf, mconf, tcls, mcls = region_loss_targets(
        bbox_pred, target)
    for v in tcoord, mcoord, tconf, mconf, tcls, mcls:
        v.need_grad = False

    # Bounding box regression
    bbox_loss = F.sum(F.squared_error(bbox_pred, tcoord) * mcoord)

    # Conf (IoU) regression loss
    conf_loss = F.sum(F.squared_error(conf_pred, tconf) * mconf)

    # Class probability regression loss
    cls_loss = F.sum(F.softmax_cross_entropy(cls_pred, tcls, axis=2) * mcls)

    # Note:
    # loss is devided by 2.0 due to the fact that the original darknet
    # code doesn't multiply the derivative of square functions by 2.0
    # in region_layer.c.
    loss = (bbox_loss + conf_loss) / 2.0 + cls_loss

    return yolo_x, target, loss, region_loss_targets


Vars = collections.namedtuple('Vars', ['x', 't', 'loss'])
Stats = collections.namedtuple(
    'Stats', ['loss', 'nGT', 'nCorrect', 'nProposals', 'mIoU', 'seen', 'time'])


class TrainGraph(object):

    def __init__(self, args, default_hw=(416, 416)):
        self.args = args
        self.seen = 0
        self.v = None
        self.region_class = None
        self.create_graph([args.batch_size, 3] + list(default_hw))
        self.tic = time.time()

    def create_graph(self, shape):
        self.v = None
        self.region_class = None
        x, t, loss, region_class = create_network(
            shape[0], shape[2], shape[3],
            self.args, self.seen)
        self.v = Vars(x, t, loss)
        self.region_class = region_class

    def forward_backward(self, image, target):
        if self.v.x.shape != image.shape:
            self.create_graph(image.shape)
        self.v.x.d = image
        execution_time = (time.time() - self.tic) * 1000
        self.tic = time.time()
        self.v.t.d = target
        self.v.loss.forward(clear_no_need_grad=True)
        loss = self.v.loss.d.copy()
        self.v.loss.backward(clear_buffer=True)
        nGT = self.region_class.prev_ngt
        nCorrect = self.region_class.prev_ncorrect
        mIoU = self.region_class.prev_miou
        nProposals = -1
        self.seen = self.region_class.seen
        return Stats(loss, nGT, nCorrect, nProposals, mIoU, self.seen, execution_time)


# object-detection/yolov2/utils.py

# Copyright 2018,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.


# This file was forked from https://github.com/marvis/pytorch-yolo2 ,
# licensed under the MIT License (see LICENSE.external for more details).


import time
import math
import numpy as np
import nnabla

import struct  # get_image_size
import imghdr  # get_image_size


def raise_info_thread(f):

    def decorated(*args, **kw):
        try:
            return f(*args, **kw)
        except:
            import sys
            import traceback
            raise sys.exc_info()[0](traceback.format_exc())
    return decorated


def sigmoid(x):
    return 1.0/(math.exp(-x)+1.)


def bbox_iou(box1, box2, x1y1x2y2=True):
    if x1y1x2y2:
        mx = min(box1[0], box2[0])
        Mx = max(box1[2], box2[2])
        my = min(box1[1], box2[1])
        My = max(box1[3], box2[3])
        w1 = box1[2] - box1[0]
        h1 = box1[3] - box1[1]
        w2 = box2[2] - box2[0]
        h2 = box2[3] - box2[1]
    else:
        mx = min(box1[0]-box1[2]/2.0, box2[0]-box2[2]/2.0)
        Mx = max(box1[0]+box1[2]/2.0, box2[0]+box2[2]/2.0)
        my = min(box1[1]-box1[3]/2.0, box2[1]-box2[3]/2.0)
        My = max(box1[1]+box1[3]/2.0, box2[1]+box2[3]/2.0)
        w1 = box1[2]
        h1 = box1[3]
        w2 = box2[2]
        h2 = box2[3]
    uw = Mx - mx
    uh = My - my
    cw = w1 + w2 - uw
    ch = h1 + h2 - uh
    carea = 0
    if cw <= 0 or ch <= 0:
        return 0.0

    area1 = w1 * h1
    area2 = w2 * h2
    carea = cw * ch
    uarea = area1 + area2 - carea
    return carea/uarea


def bbox_iou_numpy(box1, box2, x1y1x2y2=True):
    if x1y1x2y2:
        mx = np.min((box1[0], box2[0]))
        Mx = np.max((box1[2], box2[2]))
        my = np.min((box1[1], box2[1]))
        My = np.max((box1[3], box2[3]))
        w1 = box1[2] - box1[0]
        h1 = box1[3] - box1[1]
        w2 = box2[2] - box2[0]
        h2 = box2[3] - box2[1]
    else:
        mx = np.min((box1[0]-box1[2]/2.0, box2[0]-box2[2]/2.0))
        Mx = np.max((box1[0]+box1[2]/2.0, box2[0]+box2[2]/2.0))
        my = np.min((box1[1]-box1[3]/2.0, box2[1]-box2[3]/2.0))
        My = np.max((box1[1]+box1[3]/2.0, box2[1]+box2[3]/2.0))
        w1 = box1[2]
        h1 = box1[3]
        w2 = box2[2]
        h2 = box2[3]
    uw = Mx - mx
    uh = My - my
    cw = w1 + w2 - uw
    ch = h1 + h2 - uh
    carea = 0
    if cw <= 0 or ch <= 0:
        return 0.0

    area1 = w1 * h1
    area2 = w2 * h2
    carea = cw * ch
    uarea = area1 + area2 - carea
    return carea/uarea


def bbox_ious(boxes1, boxes2, x1y1x2y2=True):
    if x1y1x2y2:
        mx = np.minimum(boxes1[0], boxes2[0])
        Mx = np.maximum(boxes1[2], boxes2[2])
        my = np.minimum(boxes1[1], boxes2[1])
        My = np.maximum(boxes1[3], boxes2[3])
        w1 = boxes1[2] - boxes1[0]
        h1 = boxes1[3] - boxes1[1]
        w2 = boxes2[2] - boxes2[0]
        h2 = boxes2[3] - boxes2[1]
    else:
        mx = np.minimum(boxes1[0]-boxes1[2]/2.0, boxes2[0]-boxes2[2]/2.0)
        Mx = np.maximum(boxes1[0]+boxes1[2]/2.0, boxes2[0]+boxes2[2]/2.0)
        my = np.minimum(boxes1[1]-boxes1[3]/2.0, boxes2[1]-boxes2[3]/2.0)
        My = np.maximum(boxes1[1]+boxes1[3]/2.0, boxes2[1]+boxes2[3]/2.0)
        w1 = boxes1[2]
        h1 = boxes1[3]
        w2 = boxes2[2]
        h2 = boxes2[3]
    uw = Mx - mx
    uh = My - my
    cw = w1 + w2 - uw
    ch = h1 + h2 - uh
    mask = ((cw <= 0) + (ch <= 0) > 0)
    area1 = w1 * h1
    area2 = w2 * h2
    carea = cw * ch
    carea[mask] = 0
    uarea = area1 + area2 - carea
    return carea/uarea


def bbox_ious_numpy(boxes1, boxes2, x1y1x2y2=True):
    if x1y1x2y2:
        mx = np.minimum(boxes1[0], boxes2[0])
        Mx = np.maximum(boxes1[2], boxes2[2])
        my = np.minimum(boxes1[1], boxes2[1])
        My = np.maximum(boxes1[3], boxes2[3])
        w1 = boxes1[2] - boxes1[0]
        h1 = boxes1[3] - boxes1[1]
        w2 = boxes2[2] - boxes2[0]
        h2 = boxes2[3] - boxes2[1]
    else:
        mx = np.minimum(boxes1[0]-boxes1[2]/2.0, boxes2[0]-boxes2[2]/2.0)
        Mx = np.maximum(boxes1[0]+boxes1[2]/2.0, boxes2[0]+boxes2[2]/2.0)
        my = np.minimum(boxes1[1]-boxes1[3]/2.0, boxes2[1]-boxes2[3]/2.0)
        My = np.maximum(boxes1[1]+boxes1[3]/2.0, boxes2[1]+boxes2[3]/2.0)
        w1 = boxes1[2]
        h1 = boxes1[3]
        w2 = boxes2[2]
        h2 = boxes2[3]
    uw = Mx - mx
    uh = My - my
    cw = w1 + w2 - uw
    ch = h1 + h2 - uh
    mask = ((cw <= 0) + (ch <= 0) > 0)
    area1 = w1 * h1
    area2 = w2 * h2
    carea = cw * ch
    carea[mask] = 0
    uarea = area1 + area2 - carea
    return carea/uarea


def nms(boxes, nms_thresh):
    if len(boxes) == 0:
        return boxes

    det_confs = np.zeros(len(boxes))
    for i in range(len(boxes)):
        det_confs[i] = 1-boxes[i][4]

    sortIds, _ = zip(*sorted(enumerate(det_confs), key=lambda x: x[1]))
    out_boxes = []
    for i in range(len(boxes)):
        box_i = boxes[sortIds[i]]
        if box_i[4] > 0:
            out_boxes.append(box_i)
            for j in range(i+1, len(boxes)):
                box_j = boxes[sortIds[j]]
                if bbox_iou(box_i, box_j, x1y1x2y2=False) > nms_thresh:
                    box_j[4] = 0
    return out_boxes


def get_region_boxes(output, conf_thresh, num_classes, anchors, num_anchors, only_objectness=1, validation=False):
    anchor_step = len(anchors)//num_anchors
    if output.ndim == 3:
        output = output.reshape(-1)
    batch = output.shape[0]
    assert(output.shape[1] == (5+num_classes)*num_anchors)
    h = output.shape[2]
    w = output.shape[3]

    t0 = time.time()
    all_boxes = []
    output = output.reshape((batch*num_anchors, 5+num_classes, h*w)).transpose(
        (1, 0, 2)).reshape((5+num_classes, batch*num_anchors*h*w))

    grid_x = np.tile(np.tile(np.linspace(0, w-1, w), (h, 1)),
                     (batch*num_anchors, 1, 1)).reshape(batch*num_anchors*h*w)
    grid_y = np.tile(np.tile(np.linspace(0, h-1, h), (w, 1)).transpose(),
                     (batch*num_anchors, 1, 1)).reshape(batch*num_anchors*h*w)
    anchor_w = np.array(anchors).reshape((num_anchors, anchor_step))[:, 0:1]
    anchor_h = np.array(anchors).reshape((num_anchors, anchor_step))[:, 1:2]
    anchor_w = np.tile(np.tile(anchor_w, (batch, 1)),
                       (1, 1, h*w)).reshape(batch*num_anchors*h*w)
    anchor_h = np.tile(np.tile(anchor_h, (batch, 1)),
                       (1, 1, h*w)).reshape(batch*num_anchors*h*w)

    def npv2nnablav(x):
        v = nnabla.Variable(x.shape)
        v.d = x
        return v
    grid_x, grid_y, anchor_h, anchor_w = map(
        npv2nnablav, [grid_x, grid_y, anchor_h, anchor_w])

    output_v = nnabla.Variable(output.shape)
    output_v.d = output
    outputs = nnabla.functions.split(output_v)
    xs = nnabla.functions.sigmoid(outputs[0]) + grid_x
    ys = nnabla.functions.sigmoid(outputs[1]) + grid_y
    ws = nnabla.functions.exp(outputs[2]) * anchor_w
    hs = nnabla.functions.exp(outputs[3]) * anchor_h
    det_confs = nnabla.functions.sigmoid(outputs[4])

    cls = nnabla.functions.stack(*outputs[5:])
    cls = cls.reshape((batch*num_anchors, num_classes, h*w))
    cls = nnabla.functions.transpose(cls, [0, 2, 1]).reshape(
        (batch*num_anchors*h*w, num_classes))

    o = output[5:5+num_classes].transpose()
    v = nnabla.Variable(o.shape)
    v.d = o
    cls_confs = nnabla.functions.softmax(v)
    cls_confs.forward()

    def nnablav2np(v):
        v.forward()
        return v.d
    xs, ys, ws, hs, det_confs, cls_confs = map(
        nnablav2np, [xs, ys, ws, hs, det_confs, cls_confs])
    cls_max_confs = np.max(cls_confs, axis=1)
    cls_max_ids = np.argmax(cls_confs, axis=1)
    t1 = time.time()

    sz_hw = h*w
    sz_hwa = sz_hw*num_anchors

    if validation:
        cls_confs = cls_confs.reshape(-1, num_classes)
    t2 = time.time()
    for b in range(batch):
        boxes = []
        for cy in range(h):
            for cx in range(w):
                for i in range(num_anchors):
                    ind = b*sz_hwa + i*sz_hw + cy*w + cx
                    det_conf = det_confs[ind]
                    if only_objectness:
                        conf = det_confs[ind]
                    else:
                        conf = det_confs[ind] * cls_max_confs[ind]

                    if conf > conf_thresh:
                        bcx = xs[ind]
                        bcy = ys[ind]
                        bw = ws[ind]
                        bh = hs[ind]
                        cls_max_conf = cls_max_confs[ind]
                        cls_max_id = cls_max_ids[ind]
                        box = [bcx/w, bcy/h, bw/w, bh/h,
                               det_conf, cls_max_conf, cls_max_id]
                        if (not only_objectness) and validation:
                            for c in range(num_classes):
                                tmp_conf = cls_confs[ind][c]
                                if c != cls_max_id and det_confs[ind]*tmp_conf > conf_thresh:
                                    box.append(tmp_conf)
                                    box.append(c)
                        boxes.append(box)
        all_boxes.append(boxes)
    t3 = time.time()
    return all_boxes


def load_class_names(namesfile):
    class_names = []
    with open(namesfile, 'r') as fp:
        lines = fp.readlines()
    for line in lines:
        line = line.rstrip()
        class_names.append(line)
    return class_names


def scale_bboxes(bboxes, width, height):
    import copy
    dets = copy.deepcopy(bboxes)
    for i in range(len(dets)):
        dets[i][0] = dets[i][0] * width
        dets[i][1] = dets[i][1] * height
        dets[i][2] = dets[i][2] * width
        dets[i][3] = dets[i][3] * height
    return dets


def file_lines(thefilepath):
    count = 0
    thefile = open(thefilepath, 'rb')
    while True:
        buffer = thefile.read(8192*1024)
        if not buffer:
            break
        count += buffer.count(b'\n')
    thefile.close()
    return count


def get_image_size(fname):
    '''Determine the image type of fhandle and return its size.
    from draco'''
    with open(fname, 'rb') as fhandle:
        head = fhandle.read(24)
        if len(head) != 24:
            return
        if imghdr.what(fname) == 'png':
            check = struct.unpack('>i', head[4:8])[0]
            if check != 0x0d0a1a0a:
                return
            width, height = struct.unpack('>ii', head[16:24])
        elif imghdr.what(fname) == 'gif':
            width, height = struct.unpack('<HH', head[6:10])
        elif imghdr.what(fname) == 'jpeg' or imghdr.what(fname) == 'jpg':
            try:
                fhandle.seek(0)  # Read 0xff next
                size = 2
                ftype = 0
                while not 0xc0 <= ftype <= 0xcf:
                    fhandle.seek(size, 1)
                    byte = fhandle.read(1)
                    while ord(byte) == 0xff:
                        byte = fhandle.read(1)
                    ftype = ord(byte)
                    size = struct.unpack('>H', fhandle.read(2))[0] - 2
                # We are at a SOFn block
                fhandle.seek(1, 1)  # Skip `precision' byte.
                height, width = struct.unpack('>HH', fhandle.read(4))
            except Exception:  # IGNORE:W0703
                return
        else:
            return
        return width, height


def logging(message):
    print(('%s %s' % (time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), message)))


def set_default_context_by_args(args):
    ngpus = len(args.gpus.split(','))
    assert ngpus == 1, 'muti-gpu training is not supported so far. Given {} (size={})'.format(
        args.gpus, ngpus)

    from nnabla.ext_utils import get_extension_context
    ext = 'cpu'
    if args.use_cuda:
        ext = 'cudnn'
    ctx = get_extension_context(ext, device_id=args.gpus)
    nnabla.set_default_context(ctx)


# object-detection/yolov2/valid.py

# Copyright 2018,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.


# This file was forked from https://github.com/marvis/pytorch-yolo2 ,
# licensed under the MIT License (see LICENSE.external for more details).


import dataset
import utils
import numpy as np
import os
import itertools
from multiprocessing.pool import ThreadPool
import nnabla
import nnabla_ext.cuda
import yolov2

from arg_utils import Yolov2OptionValid
args = Yolov2OptionValid().parse_args()


def valid(weightfile, outfile, outdir):
    pool = ThreadPool(1)
    valid_images = args.valid
    name_list = args.names
    prefix = outdir
    names = utils.load_class_names(name_list)

    utils.set_default_context_by_args(args)

    with open(valid_images) as fp:
        tmp_files = fp.readlines()
        valid_files = [item.rstrip() for item in tmp_files]

    # Build the YOLO v2 network
    def create_losses(batchsize, imheight, imwidth, test=True):
        import gc
        gc.collect()
        nnabla_ext.cuda.clear_memory_cache()

        anchors = args.num_anchors
        classes = args.num_classes
        yolo_x = nnabla.Variable((batchsize, 3, imheight, imwidth))
        yolo_features = yolov2.yolov2(yolo_x, anchors, classes, test=test)
        return yolo_x, yolo_features

    yolo_x_nnabla, yolo_features_nnabla = create_losses(
        args.valid_batchsize, args.height, args.width, test=True)
    nnabla.load_parameters(weightfile)

    valid_dataset = dataset.data_iterator_yolo(valid_images, args,
                                               train=False,
                                               shape=(args.width, args.height), shuffle=False, batch_size=args.valid_batchsize)
    assert(args.valid_batchsize > 1)

    fps = [0]*args.num_classes
    if not os.path.exists(outdir):
        os.mkdir(outdir)
    for i in range(args.num_classes):
        buf = '%s/%s%s.txt' % (prefix, outfile, names[i])
        fps[i] = open(buf, 'w')

    lineId = 0
    total_samples = len(valid_files)
    total_batches = (total_samples+args.valid_batchsize -
                     1)//args.valid_batchsize

    for each_batch in range(0, total_batches):
        ret = valid_dataset.next()
        data, target = ret
        yolo_x_nnabla.d = data
        yolo_features_nnabla.forward(clear_buffer=True)
        batch_boxes = utils.get_region_boxes(
            yolo_features_nnabla.d, args.conf_thresh, args.num_classes, args.anchors, args.num_anchors, 0, 1)
        for i in range(yolo_features_nnabla.d.shape[0]):
            if lineId >= total_samples:
                print("Reached End of total_samples")
                break
            fileId = os.path.basename(valid_files[lineId]).split('.')[0]
            width, height = utils.get_image_size(valid_files[lineId])
            print(valid_files[lineId])
            lineId += 1
            boxes = batch_boxes[i]
            boxes = utils.nms(boxes, args.nms_thresh)
            for box in boxes:
                x1 = (box[0] - box[2]/2.0) * width
                y1 = (box[1] - box[3]/2.0) * height
                x2 = (box[0] + box[2]/2.0) * width
                y2 = (box[1] + box[3]/2.0) * height

                det_conf = box[4]
                for j in range((len(box)-5)//2):
                    cls_conf = box[5+2*j]
                    cls_id = box[6+2*j]
                    prob = det_conf * cls_conf
                    fps[cls_id].write('%s %f %f %f %f %f\n' %
                                      (fileId, prob, x1, y1, x2, y2))

    for i in range(args.num_classes):
        fps[i].close()


if __name__ == '__main__':
    weightfile = args.weight
    outdir = args.output
    outfile = 'comp4_det_test_'
    valid(weightfile, outfile, outdir)


# object-detection/yolov2/yolov2.py

# Copyright 2018,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 darknet19 import conv_bn_pool, darknet19_feature


def reorg_ref(x, stride):
    import numpy as np
    batch, out_c, h, w = x.shape
    c = int(out_c * stride * stride)
    h = int(h / stride)
    w = int(w / stride)
    y = np.zeros_like(x).flatten()
    x = x.flatten()
    out_c = int(c / (stride * stride))
    for i in range(x.size):
        in_index = i
        in_w = i % w
        i = int(i / w)
        in_h = i % h
        i = int(i / h)
        in_c = i % c
        i = int(i / c)
        b = i % batch
        c2 = in_c % out_c
        offset = int(in_c / out_c)
        w2 = in_w * stride + offset % stride
        h2 = in_h * stride + int(offset / stride)
        out_index = w2 + w * stride * (h2 + h * stride * (c2 + out_c * b))
        y[in_index] = x[out_index]
    return y.reshape(batch, c, h, w)


def reorg_ref_darknet(x, stride):
    import numpy as np
    batch, c, h, w = x.shape
    y = np.zeros_like(x).flatten()
    x = x.flatten()
    out_c = int(c / (stride * stride))
    for i in range(x.size):
        in_index = i
        in_w = i % w
        i = int(i / w)
        in_h = i % h
        i = int(i / h)
        in_c = i % c
        i = int(i / c)
        b = i % batch
        c2 = in_c % out_c
        offset = int(in_c / out_c)
        w2 = in_w * stride + offset % stride
        h2 = in_h * stride + int(offset / stride)
        out_index = w2 + w * stride * (h2 + h * stride * (c2 + out_c * b))
        y[in_index] = x[out_index]
    return y.reshape(batch, c * stride * stride, int(h / stride), int(w / stride))


def reorg(x, stride):
    # Input shape
    b, c, h, w = x.shape
    # Output shape
    assert h % stride == 0
    assert w % stride == 0
    C, H, W = stride * stride * c, int(h / stride), int(w / stride)
    # Reorg operation in Darknet can be done by transpose and reshape.
    r = F.reshape(x, (b, c, H, stride, W, stride))
    r = F.transpose(r, (0, 3, 5, 1, 2, 4))
    r = F.reshape(r, (b, C, H, W))
    # Note that an actual computation is only fired at the transpose
    # because `reshape` doesn't involve any computation.
    return r


def reorg_darknet_bug(x, stride):
    '''Simulate DarkNet's reorg layer including an indexing bug.
    '''
    b, c, h, w = x.shape
    assert h % stride == 0
    assert w % stride == 0
    c_bug, h_bug, h_bug = int(c / stride / stride), h * stride, w * stride
    C_bug, H_bug, W_bug = c, h, w
    C, H, W = stride * stride * c, int(h / stride), int(w / stride)
    r = F.reshape(x, (b, c_bug, h, stride, w, stride))
    r = F.transpose(r, (0, 3, 5, 1, 2, 4))
    r = F.reshape(r, (b, C, H, W))
    return r


def yolov2_feature(c13, c18, test=False, feature_dict=None):
    '''
    '''
    if feature_dict is None:
        feature_dict = {}
    # Extra feature extraction for c18
    h = conv_bn_pool(c18, 1024, 3, pool=False, test=test, name='c18_19')
    feature_dict['c18_19'] = h
    h = conv_bn_pool(h, 1024, 3, pool=False, test=test, name='c18_20')
    feature_dict['c18_20'] = h

    # Extra feature extraction for c13
    c13_h = conv_bn_pool(c13, 64, 1, pool=False, test=test, name='c13_14')
    feature_dict['c13_14'] = c13_h
    c13_h = reorg_darknet_bug(c13_h, 2)
    feature_dict['reorg'] = c13_h

    # Concatenate c13 and c18 features together
    h = F.concatenate(c13_h, h, axis=1)
    feature_dict['route'] = h

    # Extra feature extraction of the multi-scale features
    h = conv_bn_pool(h, 1024, 3, pool=False, test=test, name='c21')
    feature_dict['c21'] = h
    return h


def yolov2_detection_layer(h, num_anchors, num_classes):
    '''
    '''
    return PF.convolution(h, (4 + 1 + num_classes) *
                          num_anchors, (1, 1), name='detection')


def yolov2(x, num_anchors, num_classes, test=False, feature_dict=None):
    feature_dict_d = {}
    c18 = darknet19_feature(x, test=test, feature_dict=feature_dict_d)
    c13 = feature_dict_d['c13']
    c21 = yolov2_feature(c13, c18, test=test, feature_dict=feature_dict)
    det = yolov2_detection_layer(c21, num_anchors, num_classes)
    if feature_dict is not None:
        feature_dict.update(feature_dict_d)
        feature_dict['c21'] = c21
        feature_dict['det'] = det
    return det


def yolov2_activate(x, anchors, biases):
    shape = x.shape
    y = F.reshape(x, (shape[0], anchors, -1,) + shape[2:])
    stop = list(y.shape)
    stop[2] = 2
    t_xy = F.slice(y, (0, 0, 0, 0, 0), stop)
    stop[2] = 4
    t_wh = F.slice(y, (0, 0, 2, 0, 0), stop)
    stop[2] = 5
    t_o = F.slice(y, (0, 0, 4, 0, 0), stop)
    stop[2] = y.shape[2]
    t_p = F.slice(y, (0, 0, 5, 0, 0), stop)
    t_xy = F.sigmoid(t_xy)
    t_wh = F.exp(t_wh)
    t_o = F.sigmoid(t_o)
    t_p = F.softmax(t_p, axis=2)
    t_x, t_y, t_wh = yolov2_image_coordinate(t_xy, t_wh, biases)
    y = F.concatenate(t_x, t_y, t_wh, t_o, t_p, axis=2)
    y = F.transpose(y, (0, 1, 3, 4, 2)).reshape((
        shape[0], -1, shape[1] / anchors))
    return y


def yolov2_image_coordinate(t_xy, t_wh, biases):
    import numpy as np
    from nnabla.parameter import pop_parameter, set_parameter
    h, w = t_xy.shape[-2:]
    xs = pop_parameter('xs')
    ys = pop_parameter('ys')
    if xs is None or (h != xs.shape[-1]):
        xs = F.arange(0, w).reshape((1, 1, 1, w))
        xs.need_grad = False
        set_parameter('xs', xs)
    if ys is None or (h != ys.shape[-2]):
        ys = F.arange(0, h).reshape((1, 1, h, 1))
        ys.need_grad = False
        set_parameter('ys', ys)
    t_x, t_y = F.split(t_xy, axis=2)
    oshape = list(t_x.shape)
    oshape.insert(2, 1)
    t_x = F.reshape((t_x + xs) / w, oshape)
    t_y = F.reshape((t_y + ys) / h, oshape)
    pop_parameter('biases')
    biases = biases.reshape(
        1, biases.shape[0], biases.shape[1], 1, 1) / np.array([w, h]).reshape(1, 1, 2, 1, 1)
    b = nn.Variable.from_numpy_array(biases)
    b.need_grad = False
    set_parameter('biases', b)
    t_wh = t_wh * b
    return t_x, t_y, t_wh


# object-detection/yolov2/yolov2_detection.py

# Copyright 2018,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 yolov2
from draw_utils import DrawBoundingBoxes

import nnabla as nn
import nnabla.functions as F

import time
import numpy as np
from nnabla.utils.image_utils import imread, imresize, imsave


from arg_utils import get_anchors_by_name_or_parse


def get_args():
    import argparse
    from os.path import dirname, basename, join
    p = argparse.ArgumentParser()
    p.add_argument('--width', type=int, default=608)
    p.add_argument('--weights', type=str, default='yolov2.h5')
    p.add_argument('--context', '-c', type=str, default='cudnn')
    p.add_argument('--device-id', '-d', type=str, default='0')
    p.add_argument('--type-config', '-t', type=str, default='float')
    p.add_argument('--output', type=str, default=None)
    p.add_argument('--classes', type=int, default=80)
    p.add_argument('--class-names', type=str, default='coco.names')
    p.add_argument('--thresh', type=float, default=.5)
    p.add_argument('--nms', type=float, default=.45)
    p.add_argument('--nms-per-class', type=bool, default=True)
    p.add_argument(
        '--anchors', type=str,
        default='coco')
    p.add_argument('input', type=str, default='dog.jpg')
    args = p.parse_args()
    assert args.width % 32 == 0
    args.anchors = get_anchors_by_name_or_parse(args.anchors)
    args.num_anchors = int(len(args.anchors) // 2)
    args.anchors = np.array(args.anchors).reshape(-1, 2)
    if args.output is None:
        args.output = join('.', 'detect.' + basename(args.input))
    return args


def draw_bounding_boxes(img, bboxes, im_w, im_h, names, colors, sub_w, sub_h, thresh):
    draw = DrawBoundingBoxes(img, colors)
    for bb in bboxes:
        if bb[4] <= 0:
            continue
        # x, y, w, h = bb[:4] * np.array([im_w, im_h, im_w, im_h])
        x, y, w, h = bb[:4]
        x = (x - (1 - sub_w) / 2.) / sub_w * im_w
        y = (y - (1 - sub_h) / 2.) / sub_h * im_h
        w = w * im_w / sub_w
        h = h * im_h / sub_h
        dw = w / 2.
        dh = h / 2.
        x0 = int(np.clip(x - dw, 0, im_w))
        y0 = int(np.clip(y - dh, 0, im_h))
        x1 = int(np.clip(x + dw, 0, im_w))
        y1 = int(np.clip(y + dh, 0, im_h))
        det_ind = np.where(bb[5:] > thresh)[0]
        if len(det_ind) == 0:
            continue
        prob = bb[5 + det_ind]
        # Object detection with deep learning and OpenCV
        # https://goo.gl/q4RdcZ
        label = ', '.join("{}: {:.2f}%".format(
            names[det_ind[j]], prob[j] * 100) for j in range(len(det_ind)))
        print("[INFO] {}".format(label))
        draw.draw((x0, y0, x1, y1), det_ind[0], label)
    return draw.get()


def main():
    args = get_args()
    names = np.genfromtxt(args.class_names, dtype=str, delimiter='?')
    rng = np.random.RandomState(1223)
    colors = rng.randint(0, 256, (args.classes, 3)).astype(np.uint8)
    colors = [tuple(c.tolist()) for c in colors]

    # Set context
    from nnabla.ext_utils import get_extension_context
    ctx = get_extension_context(
        args.context, device_id=args.device_id, type_config=args.type_config)
    nn.set_default_context(ctx)

    # Load parameter
    _ = nn.load_parameters(args.weights)

    # Build a YOLO v2 network
    feature_dict = {}
    x = nn.Variable((1, 3, args.width, args.width))
    y = yolov2.yolov2(x, args.num_anchors, args.classes,
                      test=True, feature_dict=feature_dict)
    y = yolov2.yolov2_activate(y, args.num_anchors, args.anchors)
    y = F.nms_detection2d(y, args.thresh, args.nms, args.nms_per_class)

    # Read image
    img_orig = imread(args.input, num_channels=3)
    im_h, im_w, _ = img_orig.shape
    # letterbox
    w = args.width
    h = args.width

    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 = imresize(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

    # Execute YOLO v2
    print("forward")
    in_img = img.transpose(2, 0, 1).reshape(1, 3, args.width, args.width)
    x.d = in_img
    y.forward(clear_buffer=True)
    print("done")

    bboxes = y.d[0]
    img_draw = draw_bounding_boxes(
        img_orig, bboxes, im_w, im_h, names, colors, new_w * 1.0 / w, new_h * 1.0 / h, args.thresh)
    imsave(args.output, img_draw)

    # Timing
    s = time.time()
    n_time = 10
    for i in range(n_time):
        x.d = in_img
        y.forward(clear_buffer=True)
        # Invoking device-to-host copy if CUDA
        # so that time contains data transfer.
        _ = y.d
    print("Processing time: {:.1f} [ms/image]".format(
        (time.time() - s) / n_time * 1000))


if __name__ == '__main__':
    main()


# object-detection/yolov2/yolov2_nnp.py

# Copyright 2018,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 yolov2

import nnabla as nn
import nnabla.functions as F

import time
import numpy as np
from nnabla.utils.image_utils import imread, imresize, imsave


from arg_utils import get_anchors_by_name_or_parse


def get_args():
    import argparse
    from os.path import dirname, basename, join
    p = argparse.ArgumentParser()
    p.add_argument('--width', type=int, default=608)
    p.add_argument('--height', type=int, default=608)
    p.add_argument('--weights', type=str, default='yolov2.h5')
    p.add_argument('--classes', type=int, default=80)
    p.add_argument('--thresh', type=float, default=.5)
    p.add_argument('--nms', type=float, default=.45)
    p.add_argument('--nms-per-class', type=bool, default=True)
    p.add_argument(
        '--anchors', type=str,
        default='coco')
    p.add_argument('--nnp', type=str, default='yolov2.nnp')
    args = p.parse_args()
    assert args.width % 32 == 0
    assert args.height % 32 == 0
    args.anchors = get_anchors_by_name_or_parse(args.anchors)
    args.num_anchors = int(len(args.anchors) // 2)
    args.anchors = np.array(args.anchors).reshape(-1, 2)
    return args


def main():
    args = get_args()

    # Load parameter
    _ = nn.load_parameters(args.weights)

    # Build a YOLO v2 network
    x = nn.Variable((1, 3, args.height, args.width))
    y = yolov2.yolov2(x / 255.0, args.num_anchors, args.classes, test=True)
    y = yolov2.yolov2_activate(y, args.num_anchors, args.anchors)
    y = F.nms_detection2d(y, args.thresh, args.nms, args.nms_per_class)

    # Save NNP file (used in C++ inference later.).
    runtime_contents = {
        'networks': [
            {'name': 'runtime',
             'batch_size': 1,
             'outputs': {'y': y},
             'names': {'x': x}}],
        'executors': [
            {'name': 'runtime',
             'network': 'runtime',
             'data': ['x'],
             'output': ['y']}]}
    import nnabla.utils.save
    nnabla.utils.save.save(args.nnp, runtime_contents,
                           variable_batch_size=False)


if __name__ == '__main__':
    main()


# parametric_function_classes/mnist-collection/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 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 mnist_data import data_iterator_mnist
from nnabla.monitor import Monitor, MonitorSeries

import nnabla.experimental.parametric_function_classes as PFC


class Model0(PFC.Module):
    def __init__(self):
        self.conv0 = PFC.Conv2d(1, 16, (3, 3), pad=(1, 1), with_bias=False)
        self.conv1 = PFC.Conv2d(16, 32, (3, 3), pad=(1, 1), with_bias=False)
        self.bn0 = PFC.BatchNorm2d(16)
        self.bn1 = PFC.BatchNorm2d(32)
        self.act = F.relu
        self.pool = F.max_pooling

    def __call__(self, x, test=False):
        h = x
        h = self.conv0(h)
        h = self.bn0(h, test)
        h = self.act(h)
        h = self.pool(h, (2, 2))
        h = self.conv1(h)
        h = self.bn1(h, test)
        h = self.act(h)
        h = self.pool(h, (2, 2))
        return h


class Model1(PFC.Module):
    def __init__(self):
        self.conv0 = PFC.Conv2d(32, 64, (3, 3), pad=(1, 1), with_bias=False)
        self.bn0 = PFC.BatchNorm2d(64)
        self.affine = PFC.Affine(64, 10)
        self.act = F.relu
        self.pool = F.average_pooling

    def __call__(self, x, test=False):
        h = x
        h = self.conv0(h)
        h = self.bn0(h, test)
        h = self.act(h)
        h = self.pool(h, h.shape[2:])
        h = self.affine(h)
        return h


class Model(PFC.Module):
    def __init__(self):
        self.model0 = Model0()
        self.model1 = Model1()

    def __call__(self, x, test=False):
        h = x if test else F.image_augmentation(x, flip_lr=True, angle=0.26)
        h = self.model0(h, test)
        h = self.model1(h, test)
        return h


def main():
    # Context
    ctx = get_extension_context("cudnn", device_id="0")
    nn.set_default_context(ctx)
    nn.auto_forward(False)
    # Inputs
    b, c, h, w = 64, 1, 28, 28
    x = nn.Variable([b, c, h, w])
    t = nn.Variable([b, 1])
    vx = nn.Variable([b, c, h, w])
    vt = nn.Variable([b, 1])
    # Model
    model = Model()
    pred = model(x)
    loss = F.softmax_cross_entropy(pred, t)
    vpred = model(vx, test=True)
    verror = F.top_n_error(vpred, vt)
    # Solver
    solver = S.Adam()
    solver.set_parameters(model.get_parameters(grad_only=True))
    # Data Iterator
    tdi = data_iterator_mnist(b, train=True)
    vdi = data_iterator_mnist(b, train=False)
    # Monitor
    monitor = Monitor("tmp.monitor")
    monitor_loss = MonitorSeries("Training loss", monitor, interval=10)
    monitor_verr = MonitorSeries("Test error", monitor, interval=1)

    # Training loop
    for e in range(1):
        for j in range(tdi.size // b):
            i = e * tdi.size // b + j
            x.d, t.d = tdi.next()
            solver.zero_grad()
            loss.forward(clear_no_need_grad=True)
            loss.backward(clear_buffer=True)
            solver.update()
            monitor_loss.add(i, loss.d)
        error = 0.0
        for _ in range(vdi.size // b):
            vx.d, vt.d = vdi.next()
            verror.forward(clear_buffer=True)
            error += verror.d
        error /= vdi.size // b
        monitor_verr.add(i, error)


if __name__ == '__main__':
    main()


# parametric_function_classes/mnist-collection/mnist_data.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.

'''
Provide data iterator for MNIST examples.
'''
import numpy
import struct
import zlib

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


def load_mnist(train=True):
    '''
    Load MNIST dataset images and labels from the original page by Yan LeCun or the cache file.

    Args:
        train (bool): The testing dataset will be returned if False. Training data has 60000 images, while testing has 10000 images.

    Returns:
        numpy.ndarray: A shape of (#images, 1, 28, 28). Values in [0.0, 1.0].
        numpy.ndarray: A shape of (#images, 1). Values in {0, 1, ..., 9}.

    '''
    if 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))
    #       labels = numpy.frombuffer(f.read(), numpy.uint8).reshape(-1, 1)
    #
    r = download(label_uri)
    data = zlib.decompress(r.read(), zlib.MAX_WBITS | 32)
    _, size = struct.unpack('>II', data[0:8])
    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))
    r = download(image_uri)
    data = zlib.decompress(r.read(), zlib.MAX_WBITS | 32)
    _, size, height, width = struct.unpack('>IIII', data[0:16])
    images = numpy.frombuffer(data[16:], numpy.uint8).reshape(
        size, 1, height, width)
    r.close()
    logger.info('Getting image data done.')

    return images, labels


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):
        super(MnistDataSource, self).__init__(shuffle=shuffle)
        self._train = train

        self._images, self._labels = load_mnist(train)

        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()


def data_iterator_mnist(batch_size,
                        train=True,
                        rng=None,
                        shuffle=True,
                        with_memory_cache=False,
                        with_file_cache=False):
    '''
    Provide DataIterator with :py:class:`MnistDataSource`
    with_memory_cache 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.

    '''
    return data_iterator(MnistDataSource(train=train, shuffle=shuffle, rng=rng),
                         batch_size,
                         rng,
                         with_memory_cache,
                         with_file_cache)


# reinforcement_learning/dqn/atari_utils.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 gym
import numpy as np


class Squeeze(gym.ObservationWrapper):
    '''Assume wrap_deepmind with scale=True'''

    def __init__(self, env):
        from gym import spaces
        gym.ObservationWrapper.__init__(self, env)
        self.observation_space = spaces.Box(
            low=0, high=1.0,
            shape=(84, 84), dtype=np.float32)

    def observation(self, observation):
        return np.squeeze(observation)


def make_atari_deepmind(rom_name, valid=False):
    from external.atari_wrappers import make_atari, wrap_deepmind
    env = make_atari(rom_name)
    # framestack is handled by sampler.py
    env = wrap_deepmind(env, episode_life=not valid,
                        frame_stack=False, scale=True)
    env = Squeeze(env)
    return env


# reinforcement_learning/dqn/eval_atari.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 sampler import ObsSampler
from replay_memory import ReplayMemory
from validator import Validator
from explorer import GreedyExplorer
from output_path import OutputPath

from nnabla.ext_utils import get_extension_context
import nnabla as nn

import os


def get_args():

    import argparse

    p = argparse.ArgumentParser()
    p.add_argument('--gym-env', '-g', default='BreakoutNoFrameskip-v4')
    p.add_argument('--nnp', '-m', default=None)
    p.add_argument('--nnp_dir', '-p', default=None)
    p.add_argument('--num_frames', '-f', type=int, default=4)
    p.add_argument('--no-render', '-n', action='store_true', default=True)
    p.add_argument('--extension', '-e', default='cpu')
    p.add_argument('--device-id', '-d', default='0')
    p.add_argument('--log_path', '-l', default='./tmp.output')

    return p.parse_args()


def main():

    args = get_args()

    nn.set_default_context(get_extension_context(
        args.extension, device_id=args.device_id))

    from atari_utils import make_atari_deepmind
    env = make_atari_deepmind(args.gym_env, valid=True)
    print('Observation:', env.observation_space)
    print('Action:', env.action_space)
    obs_sampler = ObsSampler(args.num_frames)
    val_replay_memory = ReplayMemory(env.observation_space.shape,
                                     env.action_space.shape, max_memory=args.num_frames)

    # for one file
    explorer = GreedyExplorer(
        env.action_space.n, use_nnp=True, nnp_file=args.nnp, name='qnet')
    validator = Validator(env, val_replay_memory, explorer, obs_sampler,
                          num_episodes=30, clip_episode_step=True,
                          render=not args.no_render)

    mean_reward = validator.step()
    with open(os.path.join(args.log_path, 'mean_reward.txt'), 'a') as f:
        print("{} {}".format(args.gym_env, str(mean_reward)), file=f)


if __name__ == '__main__':
    main()


# reinforcement_learning/dqn/explorer.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
from six.moves import range

from collections import namedtuple

import numpy as np

import nnabla as nn

from abc import ABC, abstractmethod


class Explorer(ABC):
    def __init__(self, num_actions, name='q', network=None):
        self.num_actions = num_actions
        self.name = name
        self.network = network

    @abstractmethod
    def select_action(self, obs, val=False):
        raise NotImplementedError()

    def build_network(self, obs):
        s = nn.Variable(obs.shape)
        with nn.parameter_scope(self.name):
            q = self.q_builder(s, self.num_actions, test=True)
        Variables = namedtuple(
            'Variables', ['s', 'q'])
        self.network = Variables(s, q)

    def build_network_from_nnp(self, nnp_file):
        from nnabla.utils import nnp_graph
        nnp = nnp_graph.NnpLoader(nnp_file)
        net = nnp.get_network(self.name, batch_size=1)
        s = net.inputs['s']
        q = net.outputs['q']
        Variables = namedtuple(
            'Variables', ['s', 'q'])
        assert q.shape[1] == self.num_actions
        self.network = Variables(s, q)


class GreedyExplorer(Explorer):
    def __init__(self, num_actions, use_nnp=False,
                 q_builder=None, nnp_file=None, name='q'):
        self.num_actions = num_actions
        self.q_builder = q_builder
        self.network = None
        self.name = name
        if use_nnp:
            if nnp_file is None:
                return
            super().build_network_from_nnp(nnp_file)

    def select_action(self, obs, val=False):
        if not self.network:
            super().build_network(obs)
        self.network.s.d = obs
        self.network.q.forward(clear_buffer=True)
        return np.argmax(self.network.q.d, axis=1)


class EGreedyExplorer(Explorer):
    def __init__(self, num_actions, epsilon=0.0, use_nnp=False,
                 q_builder=None, nnp_file=None, name='q', rng=np.random):
        self.num_actions = num_actions
        self.epsilon = epsilon
        self.q_builder = q_builder
        self.network = None
        self.name = name
        self.rng = rng
        if use_nnp:
            if nnp_file is None:
                return
            super().build_network_from_nnp(nnp_file)

    def select_action(self, obs, val=False):
        if not self.network:
            super().build_network(obs)
        self.network.s.d = obs
        self.network.q.forward(clear_buffer=True)
        if self.rng.rand() >= self.epsilon:
            return np.argmax(self.network.q.d, axis=1)
        return self.rng.randint(self.num_actions, size=(obs.shape[0],))


class LinearDecayEGreedyExplorer(EGreedyExplorer):
    def __init__(self, num_actions, eps_val=0.05, eps_start=.9, eps_end=.1, eps_steps=1e6,
                 use_nnp=False, q_builder=None, nnp_file=None,
                 name='q', rng=np.random):
        super().__init__(num_actions, epsilon=eps_start, use_nnp=use_nnp,
                         q_builder=q_builder, nnp_file=nnp_file,
                         name=name, rng=rng)
        self.eps_start = eps_start
        self.eps_end = eps_end
        self.eps_steps = eps_steps
        self.eps_val = eps_val
        self.time = 0

    def update(self):
        self.time += 1

    def linear_decay_epsilon(self):
        self.epsilon = max(
           self.eps_end,
           self.eps_start +
           (self.eps_end - self.eps_start) * self.time / self.eps_steps)

    def select_action(self, obs, val=False):
        if val:
            self.epsilon = self.eps_val
        else:
            self.linear_decay_epsilon()
        return super().select_action(obs)


# reinforcement_learning/dqn/learner.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
from six.moves import range

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.utils.save as save

from collections import namedtuple
import os


def update_variable_dict(d, h, name):
    if d is None:
        return
    d[name] = h


def _q_cnn(inp, convs, hiddens, num_actions, test=False, variable_dict=None):
    h = inp
    # Build convs
    for i, (o, k, s) in enumerate(convs):
        name = 'conv{}'.format(i + 1)
        h = F.relu(PF.convolution(h, o, (k, k), stride=(
            s, s), fix_parameters=test, name=name))
        update_variable_dict(variable_dict, h, name)
    # Build affines
    for i, o in enumerate(hiddens):
        name = 'fc{}'.format(i + 1)
        h = F.relu(PF.affine(h, o, fix_parameters=test, name=name))
        update_variable_dict(variable_dict, h, name)

    return PF.affine(h, num_actions, fix_parameters=test, name='fc_fin')


def q_cnn(inp, num_actions, test=False):
    return _q_cnn(
        inp,
        [(32, 8, 4), (64, 4, 2), (64, 3, 1)],
        [256],
        num_actions, test, None)


def q_mlp(s, num_actions, test=False):
    h = s
    for i, n in enumerate([64]):
        with nn.parameter_scope('fc%d' % (i + 1)):
            h = PF.affine(h, n, fix_parameters=test)
            h = F.relu(h)
    return PF.affine(h, num_actions, name='fc_fin', fix_parameters=test)


class QLearner(object):
    def __init__(self, q_builder, num_actions, clip_reward=1.0, sync_freq=500,
                 save_freq=10000, save_path=None,
                 gamma=0.99, learning_rate=5e-4, weight_decay=0, name_q='q', name_qnext='qnext'):
        self.built = False
        self.q_builder = q_builder
        self.num_actions = num_actions
        self.clip_reward = clip_reward
        self.sync_freq = sync_freq
        self.save_freq = save_freq
        if save_path is None:
            import output_path
            save_path = output_path.default_output_path()
        self.save_path = save_path
        self.gamma = gamma
        self.learning_rate = learning_rate
        self.weight_decay = weight_decay
        self.name_q = name_q
        self.name_qnext = name_qnext

        # count of neural network update (steps / train_freq)
        self.update_count = 0

    def build_train_graph(self, batch):
        self.solver = S.Adam(self.learning_rate)

        obs, action, reward, terminal, newobs = batch
        # Create input variables
        s = nn.Variable(obs.shape)
        a = nn.Variable(action.shape)
        r = nn.Variable(reward.shape)
        t = nn.Variable(terminal.shape)
        snext = nn.Variable(newobs.shape)
        with nn.parameter_scope(self.name_q):
            q = self.q_builder(s, self.num_actions, test=False)
            self.solver.set_parameters(nn.get_parameters())
        with nn.parameter_scope(self.name_qnext):
            qnext = self.q_builder(snext, self.num_actions, test=True)
        qnext.need_grad = False
        clipped_r = F.minimum_scalar(F.maximum_scalar(
            r, -self.clip_reward), self.clip_reward)
        q_a = F.sum(
            q * F.one_hot(F.reshape(a, (-1, 1), inplace=False), (q.shape[1],)), axis=1)
        target = clipped_r + self.gamma * (1 - t) * F.max(qnext, axis=1)
        loss = F.mean(F.huber_loss(q_a, target))
        Variables = namedtuple(
            'Variables', ['s', 'a', 'r', 't', 'snext', 'q', 'loss'])
        self.v = Variables(s, a, r, t, snext, q, loss)
        self.sync_models()
        self.built = True

    def sync_models(self):
        with nn.parameter_scope(self.name_q):
            q_params = nn.get_parameters(grad_only=False)
        with nn.parameter_scope(self.name_qnext):
            qnext_params = nn.get_parameters(grad_only=False)
        for k, v in q_params.items():
            qnext_params[k].data.copy_from(v.data)

    def save_model(self):
        from nnabla.utils.save import save
        with nn.parameter_scope(self.name_q):
            save(self.save_path.get_filepath('qnet_{:08d}.nnp'.format(self.update_count)),
                 {'networks': [
                     {'name': 'qnet',
                      'batch_size': self.v.s.shape[0],
                      'outputs': {'q': self.v.q},
                      'names': {'s': self.v.s}}
                 ]
            })

    def update(self, batch):
        if not self.built:
            self.build_train_graph(batch)
        self.v.s.d = batch[0]
        self.v.a.d = batch[1]
        self.v.r.d = batch[2]
        self.v.t.d = batch[3]
        self.v.snext.d = batch[4]
        self.solver.zero_grad()
        self.v.loss.forward(clear_no_need_grad=True)
        self.v.loss.backward(clear_buffer=True)
        if self.weight_decay:
            self.solver.weight_decay(self.weight_decay)
        self.solver.update()
        self.update_count += 1
        if self.update_count % self.sync_freq == 0:
            self.sync_models()
        if self.update_count % self.save_freq == 0:
            self.save_model()
        return self.v.loss.d.copy()


# reinforcement_learning/dqn/output_path.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 os


class OutputPath(object):
    def __init__(self, path='./tmp.output/'):
        if not os.path.isdir(path):
            os.makedirs(path)
        self.path = path

    def get_filepath(self, name):
        return os.path.join(self.path, name)


_default_output_path = None


def default_output_path():
    global _default_output_path
    if _default_output_path is None:
        _default_output_path = OutputPath()
    return _default_output_path


# reinforcement_learning/dqn/play_atari.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 sampler import ObsSampler
from replay_memory import ReplayMemory
from validator import Validator
from explorer import GreedyExplorer

from nnabla.ext_utils import get_extension_context
from nnabla.utils.data_source_loader import download
from nnabla.logger import logger

import nnabla as nn

import os


def get_args():

    import argparse

    p = argparse.ArgumentParser()
    p.add_argument('--gym-env', '-g', default='BreakoutNoFrameskip-v4')
    p.add_argument('--nnp', '-m', default=None)
    p.add_argument('--num_frames', '-f', type=int, default=4)
    p.add_argument('--no-render', '-n', action='store_true', default=False)
    p.add_argument('--extension', '-e', default='cpu')
    p.add_argument('--device-id', '-d', default='0')

    return p.parse_args()


def find_local_nnp(env_name):
    nnp_fpath = os.path.join("asset", env_name, "qnet.nnp")
    return os.path.exists(nnp_fpath)


def main():

    args = get_args()

    nn.set_default_context(get_extension_context(
        args.extension, device_id=args.device_id))

    if args.nnp is None:
        local_nnp_dir = os.path.join("asset", args.gym_env)
        local_nnp_file = os.path.join(local_nnp_dir, "qnet.nnp")

        if not find_local_nnp(args.gym_env):
            logger.info("Downloading nnp data since you didn't specify...")
            nnp_uri = os.path.join("https://nnabla.org/pretrained-models/nnp_models/examples/dqn",
                                   args.gym_env,
                                   "qnet.nnp")
            if not os.path.exists(local_nnp_dir):
                os.mkdir(local_nnp_dir)
            download(nnp_uri, output_file=local_nnp_file, open_file=False)
            logger.info("Download done!")

        args.nnp = local_nnp_file

    from atari_utils import make_atari_deepmind
    env = make_atari_deepmind(args.gym_env, valid=False)
    print('Observation:', env.observation_space)
    print('Action:', env.action_space)
    obs_sampler = ObsSampler(args.num_frames)
    val_replay_memory = ReplayMemory(env.observation_space.shape,
                                     env.action_space.shape,
                                     max_memory=args.num_frames)
    # just play greedily
    explorer = GreedyExplorer(
        env.action_space.n, use_nnp=True, nnp_file=args.nnp, name='qnet')
    validator = Validator(env, val_replay_memory, explorer, obs_sampler,
                          num_episodes=1, render=not args.no_render)
    while True:
        validator.step()


if __name__ == '__main__':
    main()


# reinforcement_learning/dqn/replay_memory.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 numpy as np


class ReplayMemory(object):

    def __init__(self, obs_dims, action_dims, max_memory=1000000,
                 obs_dtype=np.float32, action_dtype=np.float32,
                 reward_dtype=np.float32):
        self.obs = np.empty((max_memory,) + obs_dims, dtype=obs_dtype)
        self.action = np.empty((max_memory,) + action_dims, dtype=action_dtype)
        self.reward = np.empty((max_memory,), dtype=reward_dtype)
        self.terminal = np.ones((max_memory,), dtype=np.uint8)
        self.add_st_called = False
        self.cursor = -1
        self.seen = -1

    def add_st(self, obs, terminal):
        assert not self.add_st_called, "calling add_s consecutively is prohibited."
        self.cursor += 1
        self.seen += 1
        if self.cursor >= self.reward.shape[0]:
            self.cursor = 0
        # print("cursor={}, seen={}".format(self.cursor, self.seen))
        self.obs[self.cursor] = obs
        self.terminal[self.cursor] = terminal
        if terminal:
            # Do not toggle add_st_called because add_st is called
            # at the beginning of the next episode.
            return
        self.add_st_called = True

    def add_ar(self, action, reward):
        assert self.add_st_called, "add_art must be called after add_s is called."
        self.action[self.cursor] = action
        self.reward[self.cursor] = reward
        self.add_st_called = False

    def sample(self, num_samples, sampler, rng=np.random):
        max_ind = min(max(self.seen, self.cursor), self.obs.shape[0])
        obs = []
        action = []
        reward = []
        terminal = []
        newobs = []
        i = 0
        while i < num_samples:
            j = rng.randint(max_ind)
            sample = sampler(self, j)
            if sample is None:
                continue
            obs += [sample[0]]
            action += [sample[1]]
            reward += [sample[2]]
            terminal += [sample[3]]
            newobs += [sample[4]]
            i += 1
        obs = np.stack(obs)
        action = np.stack(action)
        reward = np.stack(reward)
        terminal = np.stack(terminal)
        newobs = np.stack(newobs)
        return obs, action, reward, terminal, newobs


# reinforcement_learning/dqn/sampler.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 numpy as np


class Sampler(object):
    def __init__(self, num_frames=4):
        self.num_frames = num_frames

    def __call__(self, replay_memory, i):
        if replay_memory.terminal[i] or replay_memory.cursor == i:
            return None
        inds = np.arange(i - self.num_frames + 1, i + 2)
        t_flags = replay_memory.terminal.take(inds, axis=0, mode='wrap')
        obs = replay_memory.obs.take(inds, axis=0, mode='wrap')
        action = replay_memory.action[i]
        reward = replay_memory.reward[i]
        terminal = t_flags[-1]
        if np.all(t_flags[:-1] == 0):
            return obs[:-1], action, reward, terminal, obs[1:]
        first = np.cumsum(t_flags[:-1]).argmax() + 1
        obs[:first] = obs[first]
        return obs[:-1], action, reward, terminal, obs[1:]


class ObsSampler(object):
    def __init__(self, num_frames=4):
        self.num_frames = num_frames

    def __call__(self, replay_memory):
        i = replay_memory.cursor
        if self.num_frames == 1:
            return replay_memory.obs[i][np.newaxis]
        inds = np.arange(i - self.num_frames + 1, i + 1)
        obs = replay_memory.obs.take(inds, axis=0, mode='wrap')
        t_flags = replay_memory.terminal.take(inds, axis=0, mode='wrap')
        if np.all(t_flags[:-1] == 0):
            return obs
        first = np.cumsum(t_flags[:-1]).argmax() + 1
        obs[:first] = obs[first]
        return obs


# reinforcement_learning/dqn/train_atari.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
from six.moves import range

import os
import numpy as np

from replay_memory import ReplayMemory
from sampler import Sampler, ObsSampler
from learner import QLearner, q_cnn
from explorer import LinearDecayEGreedyExplorer
from trainer import Trainer
from validator import Validator
from output_path import OutputPath

from nnabla.ext_utils import get_extension_context
import nnabla as nn
from nnabla.monitor import Monitor
from tensorboardX import SummaryWriter


def get_args():

    import argparse

    p = argparse.ArgumentParser()
    p.add_argument('--gym-env', '-g', default='BreakoutNoFrameskip-v4')
    p.add_argument('--num_epochs', '-E', type=int, default=1000000)
    p.add_argument('--num_episodes', '-T', type=int, default=10)
    p.add_argument('--num_val_episodes', '-V', type=int, default=1)
    p.add_argument('--num_eval_steps', '-S', type=int, default=125000*4)
    p.add_argument('--inter_eval_steps', '-i', type=int, default=250000*4)
    p.add_argument('--num_frames', '-f', type=int, default=4)
    p.add_argument('--render-train', '-r', action='store_true')
    p.add_argument('--render-val', '-v', action='store_true')
    p.add_argument('--extension', '-e', default='cpu')
    p.add_argument('--device-id', '-d', default='0')
    p.add_argument('--log_path', '-l', default='./tmp.output')

    return p.parse_args()


def main():

    args = get_args()

    nn.set_default_context(get_extension_context(
        args.extension, device_id=args.device_id))

    if args.log_path:
        output_path = OutputPath(args.log_path)
    else:
        output_path = OutputPath()
    monitor = Monitor(output_path.path)

    tbw = SummaryWriter(output_path.path)

    # Create an atari env.
    from atari_utils import make_atari_deepmind
    env = make_atari_deepmind(args.gym_env, valid=False)
    env_val = make_atari_deepmind(args.gym_env, valid=True)
    print('Observation:', env.observation_space)
    print('Action:', env.action_space)

    # 10000 * 4 frames
    val_replay_memory = ReplayMemory(
        env.observation_space.shape, env.action_space.shape, max_memory=args.num_frames)
    replay_memory = ReplayMemory(
        env.observation_space.shape, env.action_space.shape, max_memory=40000)

    learner = QLearner(q_cnn, env.action_space.n, sync_freq=1000, save_freq=250000,
                       gamma=0.99, learning_rate=1e-4, name_q='q', save_path=output_path)

    explorer = LinearDecayEGreedyExplorer(
        env.action_space.n, eps_start=1.0, eps_end=0.01, eps_steps=1e6,
        q_builder=q_cnn, name='q')

    sampler = Sampler(args.num_frames)
    obs_sampler = ObsSampler(args.num_frames)

    validator = Validator(env_val, val_replay_memory, explorer, obs_sampler,
                          num_episodes=args.num_val_episodes, num_eval_steps=args.num_eval_steps,
                          render=args.render_val, monitor=monitor, tbw=tbw)

    trainer_with_validator = Trainer(env, replay_memory, learner, sampler, explorer, obs_sampler, inter_eval_steps=args.inter_eval_steps,
                                     num_episodes=args.num_episodes, train_start=10000, batch_size=32,
                                     render=args.render_train, validator=validator, monitor=monitor, tbw=tbw)

    for e in range(args.num_epochs):
        trainer_with_validator.step()


if __name__ == '__main__':
    main()


# reinforcement_learning/dqn/trainer.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
from six.moves import range

import numpy as np
import time

from nnabla.monitor import MonitorSeries


class Trainer(object):

    def __init__(self, env, replay_memory, learner, sampler, explorer,
                 obs_sampler, inter_eval_steps=-1, num_episodes=10,
                 train_start=1000, train_freq=4, batch_size=32, render=False,
                 validator=None, monitor=None, tbw=None):
        self.env = env
        self.replay_memory = replay_memory
        self.learner = learner
        self.sampler = sampler
        self.explorer = explorer
        self.obs_sampler = obs_sampler
        self.inter_eval_steps = inter_eval_steps
        self.num_episodes = num_episodes
        self.train_start = train_start
        self.train_started = False
        self.train_freq = train_freq
        self.batch_size = batch_size
        self.render = render
        self.validator = validator
        self.steps = 0
        self.losses = []
        self.monitor = None
        if monitor:
            monitor_loss = MonitorSeries("QNetwork Loss", monitor)
            monitor_epsilon = MonitorSeries("Epsilon", monitor)
            monitor_train_reward = MonitorSeries("Train Reward", monitor)
            self.monitor = {
                'loss': monitor_loss,
                'epsilon': monitor_epsilon,
                'reward': monitor_train_reward
            }
        self.tbw = tbw

    def trained_steps(self):
        return self.steps-self.train_start

    def step_train(self):
        batch = self.replay_memory.sample(self.batch_size, self.sampler)
        loss = self.learner.update(batch)
        self.losses += [loss]
        self.train_started = True

    def step_episode(self):
        # A. initialize env
        obs = self.env.reset()
        self.replay_memory.add_st(obs, False)

        # B. Time steps(=frames) loop
        total_reward = 0.0
        while True:
            if self.render:
                self.env.render()

            # B-1. determine action
            obs_net = self.obs_sampler(self.replay_memory)[np.newaxis]
            action = self.explorer.select_action(obs_net)
            if self.train_started:
                self.explorer.update()  # for epsilon linear decay

            # B-2. one step forward simulator
            obs, reward, done, _ = self.env.step(action[0])
            self.replay_memory.add_ar(action, reward)
            self.replay_memory.add_st(obs, done)
            total_reward += reward

            # B-3. Train step
            if (self.steps+1) > self.train_start and \
               (self.steps+1) % self.train_freq == 0:
                self.step_train()

            # B-. Increment
            self.steps += 1

            # B-. Force evaluate in a certain number of intermediate step
            if self.trained_steps() != 0 and \
               self.trained_steps() % self.inter_eval_steps == 0:
                self.validator.evaluate(self.trained_steps())

            # B-. Break if it's terminal
            if done:
                break

        return total_reward

    def step(self):
        total_rewards = []
        for _ in range(self.num_episodes):
            total_rewards += [self.step_episode()]

        # log output
        if self.trained_steps() >= 0:
            if self.monitor:
                self.monitor['loss'].add(self.trained_steps(),
                                         np.mean(self.losses))
                self.monitor['reward'].add(self.trained_steps(),
                                           np.mean(total_rewards))
                self.monitor['epsilon'].add(self.trained_steps(),
                                            self.explorer.epsilon)
            if self.tbw:
                self.tbw.add_scalar('training/loss',
                                    np.mean(self.losses), self.trained_steps())
                self.tbw.add_scalar('training/reward',
                                    np.mean(total_rewards), self.trained_steps())
                self.tbw.add_scalar('training/epsilon',
                                    self.explorer.epsilon, self.trained_steps())

        # validate if validator is available
        if self.validator:
            self.validator.step(self.trained_steps())

        self.losses = []


# reinforcement_learning/dqn/validator.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
from six.moves import range

import time
import os
import numpy as np

from nnabla.monitor import MonitorSeries


class Validator(object):

    def __init__(self, env, replay_memory, explorer, obs_sampler, num_episodes=10, clip_episode_step=False, num_ces=4500*4, num_eval_steps=125000*4, render=False, monitor=None, tbw=None):
        self.env = env
        self.replay_memory = replay_memory
        self.obs_sampler = obs_sampler
        self.explorer = explorer
        self.num_episodes = num_episodes
        self.clip_episode_step = clip_episode_step
        self.num_ces = num_ces
        self.num_eval_steps = num_eval_steps
        self.render = render
        self.steps = 0
        self.losses = []
        self.start_time = time.time()
        self.monitor = None
        if monitor:
            monitor_val_score = MonitorSeries("Validation Score", monitor)
            monitor_eval_score = MonitorSeries("Eval Score", monitor)
            self.monitor = {
                'val_score': monitor_val_score,
                'eval_score': monitor_eval_score
            }
        self.tbw = tbw

    def step_episode(self):
        # A. initialize env
        obs = self.env.reset()
        self.replay_memory.add_st(obs, False)

        if self.clip_episode_step:
            self.steps = 0

        # B. Time steps loop
        total_reward = 0.0
        while True:
            if self.render:
                self.env.render()
            # B-1. determine action
            obs_net = self.obs_sampler(self.replay_memory)[np.newaxis]
            action = self.explorer.select_action(obs_net, val=True)
            # B-2. one step forward simulator
            obs, reward, done, _ = self.env.step(action[0])
            self.replay_memory.add_ar(action, reward)
            total_reward += reward

            # B-. Increment
            self.steps += 1

            if self.clip_episode_step and self.steps >= self.num_ces:
                self.replay_memory.add_st(obs, True)
                break
            else:
                self.replay_memory.add_st(obs, done)

            # B-. Break if it's terminal
            if done:
                break

        return total_reward

    def step(self, cur_train_steps=-1):
        total_rewards = []
        for episode in range(self.num_episodes):
            total_rewards += [self.step_episode()]
            #print('episode{} ended in {} steps'.format(str(episode), self.steps))
        mean_reward = np.mean(total_rewards)
        if cur_train_steps >= 0:
            if self.monitor:
                self.monitor['val_score'].add(cur_train_steps, mean_reward)
            if self.tbw:
                self.tbw.add_scalar('validation/score',
                                    mean_reward, cur_train_steps)
        return mean_reward

    # This evaluation is mainly for DQN atari evaluation
    def evaluate(self, cur_train_steps):
        self.steps = 0
        total_rewards = []
        while True:
            reward = self.step_episode()
            if self.steps > self.num_eval_steps:
                mean_reward = np.mean(total_rewards)
                if self.monitor:
                    self.monitor['eval_score'].add(
                        cur_train_steps, mean_reward)
                if self.tbw:
                    self.tbw.add_scalar('evaluation/score',
                                        mean_reward, cur_train_steps)
                break
            else:
                total_rewards += [reward]


# semantic-segmentation/deeplabv3plus/_checkpoint_nnp_util.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 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


# semantic-segmentation/deeplabv3plus/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=10000, model_save_path=None, learning_rate=1e-3, batch_size=2, weight_decay=1e-4, description=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
    if description is None:
        description = "Examples on data iterator examples. The following help shared among examples in this folder. Some arguments are valid or invalid in some examples."
    parser = argparse.ArgumentParser(description)
    parser.add_argument('--fine-tune', action='store_true',
                        default=False, help="Whether to fine tune model or not; False by default")
    parser.add_argument('--distributed', action='store_true',
                        default=False, help="Whether to use distributed/single gpu training; False by default")
    parser.add_argument('--compute-acc', action='store_true',
                        default=False, help="Whether to compute the accuracy mean IOU value during training/validation; False by default")
    parser.add_argument("--input-ckpt-file", type=str)
    parser.add_argument("--output-nnabla-file", type=str,
                        default='deeplab_nnabla.h5')
    parser.add_argument("--batch-size", "-b", type=int, default=batch_size)
    parser.add_argument("--label-path", type=str)
    parser.add_argument("--data-dir", type=str, help='Path to VOC datset.')
    parser.add_argument("--train-file", type=str,
                        help='VOC train split text file')
    parser.add_argument("--val-file",
                        type=str, help='VOC val split text file')
    parser.add_argument("--train-dir", "-t",
                        type=str, default=model_save_path,
                        help='Path to training data.')
    parser.add_argument("--val-dir", "-v",
                        type=str, default=model_save_path,
                        help='Path to validation data.')
    parser.add_argument("--train-label-dir",
                        type=str, default=model_save_path,
                        help='Path to training data-labels.')
    parser.add_argument("--val-label-dir",
                        type=str, default=model_save_path,
                        help='Path to validation data-labels.')
    parser.add_argument("--learning-rate", "-l",
                        type=float, default=learning_rate)
    parser.add_argument("--output-stride",
                        type=int, default=16)
    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", 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("--accum-grad",
                        type=int, default=1,
                        help='Weight decay factor of SGD update.')
    parser.add_argument("--weight-decay", "-w",
                        type=float, default=weight_decay,
                        help='Weight decay factor of SGD update.')
    parser.add_argument("--warmup-epoch", type=int, default=5)
    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 cuda.cudnn`.')
    parser.add_argument("--type-config", type=str, default='float',
                        help='Type of computation. e.g. "float", "half".')
    parser.add_argument("--model-save-interval", "-s", type=int, default=1000,
                        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("--pretrained-model-path",
                        type=str, default=model_save_path,
                        help='Path the pretrained model parameters saved.')
    parser.add_argument('--context', '-c', type=str,
                        default='cpu', help="Extension modules. ex) 'cpu', 'cudnn'.")
    parser.add_argument('--channel', default=1, type=int)
    parser.add_argument('--image-width', default=513, type=int)
    parser.add_argument('--image-height', default=513, type=int)
    parser.add_argument('--dataset-path', type=str)
    parser.add_argument("--model-load-path", "-T",
                        type=str, default=model_save_path,
                        help='Path the model parameters loaded.')
    parser.add_argument('--label-file-path', type=str)
    parser.add_argument('--test-image-file', type=str)
    parser.add_argument('--num-class', default=10, type=int)
    parser.add_argument('--train-samples', default=10, type=int)
    parser.add_argument('--val-samples', default=10, type=int)
    parser.add_argument("--sync-weight-every-itr",
                        type=int, default=100,
                        help="Sync weights every specified iteration. NCCL uses\
 the ring all reduce, so gradients in each device are not exactly same. When it\
 is accumulated in the weights, the weight values in each device diverge.")
    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


# semantic-segmentation/deeplabv3plus/convert_tf_nnabla.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 tensorflow.python import pywrap_tensorflow

import nnabla as nn
import nnabla.parametric_functions as pf

import numpy

from args import get_args


def tf_to_nn_param_map():
    '''map from tensor name to Nnabla default parameter names 
    '''
    return {
        '/xception_module': '',
        'BatchNorm': 'bn',
        'moving_mean': 'mean',
        'moving_variance': 'var',
        'depthwise_weights': 'depthwise_conv/W',
        'weights': 'conv/W',
        '_pointwise': '/pointwise',
        '_depthwise': '/depthwise',
        'conv1_1': '1',
        'conv1_2': '2',
        'biases': 'conv/b',
        'logits/semantic': 'decoder/logits/affine'
    }


def rename_params(param_name):
    '''
       Rename the ckpt tensor names to corresp. Nnabla param names
    '''

    tf_to_nn_dict = tf_to_nn_param_map()
    d_conv = '_depthwise'
    p_conv = '_pointwise'

    for key in tf_to_nn_dict:
        if key in param_name:
            param_name = param_name.replace(key, tf_to_nn_dict[key])

    return param_name


def convert(input_file, output_file):
    ''' Convert the input checkpoint file to output hdf5 file
    '''
    parse_tf_ckpt(input_file, output_file)


def parse_tf_ckpt(ckpt_file, h5_file):
    ''' Parse the TF checkpoint file and save as nnabla parameters
    '''
    # Get tensorflow checkpoint reader
    reader = pywrap_tensorflow.NewCheckpointReader(ckpt_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 sorted(var_to_shape_map):

        # Read tensor values for each tensor name
        weight = reader.get_tensor(key)
        if 'depthwise' in key and weight.ndim == 4:
            weight = numpy.squeeze(weight, axis=3)
            weight = numpy.transpose(weight, (2, 0, 1))

        else:
            if(weight.ndim == 4):
                weight = numpy.transpose(weight, (3, 2, 0, 1))

        if 'BatchNorm' in key:
            weight = weight.reshape((1, -1, 1, 1))

        if 'Momentum' in key or 'ExponentialMovingAverage' in key:
            continue

        key = rename_params(key)

        # 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():
    ''' 
    Main

    Usage: python convert_tf_nnabla.py --input-ckpt-file=/path to ckpt file --output-nnabla-file=/output .h5 file

    '''

    # Parse the arguments
    args = get_args()

    # convert the input file(.ckpt) to the output file(.h5)
    convert(args.input_ckpt_file, args.output_nnabla_file)


if __name__ == '__main__':
    main()


# semantic-segmentation/deeplabv3plus/dataset_utils.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 numpy as np
import os
from imageio import imread
from args import get_args
import matplotlib.pyplot as plt


def get_color():
    # RGB format
    return np.array([[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0], [0, 0, 128], [120, 0, 128], [0, 128, 128], [128, 128, 128], [64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0], [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128], [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0], [0, 64, 128], [224, 224, 192]])


def encode_label(label):
    '''
    Converting pixel values to corresponding class numbers. Assuming that the input label in 3-dim(h,w,c) and in BGR fromat read from cv2
    '''

    h, w, c = label.shape
    new_label = np.zeros((h, w, 1), dtype=np.int32)

    cls_to_clr_map = get_color()

    for i in range(cls_to_clr_map.shape[0]):
        #new_label[(label == cls_to_clr_map[i])[:,:,0]] = i
        #new_label[np.argwhere((label.astype(np.int32) == cls_to_clr_map[i]).all(axis=2))]=i
        print(np.where((label.astype(np.int32) == [120, 0, 128]).all(axis=2)))
        if i == 21:
            new_label[np.where(
                (label.astype(np.int32) == cls_to_clr_map[i]).all(axis=2))] = 255
        else:
            new_label[np.where(
                (label.astype(np.int32) == cls_to_clr_map[i]).all(axis=2))] = i

    return new_label


# this method should generate train-image.txt and train-label.txt
def generate_path_files(data_dir, train_file, val_file):

    ti = open('train_image.txt', 'w')
    tl = open('train_label.txt', 'w')
    vi = open('val_image.txt', 'w')
    vl = open('val_label.txt', 'w')

    rootdir = data_dir

    train_text_file = open(train_file, "r")
    lines = [line[:-1] for line in train_text_file]
    for line in lines:
        if os.path.exists(data_dir+'JPEGImages/'+line+'.jpg'):
            ti.write(data_dir+'JPEGImages/'+line+'.jpg' + '\n')
            assert (os.path.isfile(data_dir+'SegmentationClass/encoded/'+line +
                                   '.npy')), "No matching label file for image : " + line + '.jpg'
            tl.write(data_dir+'SegmentationClass/encoded/'+line + '.npy' + '\n')

    val_text_file = open(val_file, "r")
    lines = [line[:-1] for line in val_text_file]
    for line in lines:
        if os.path.exists(data_dir+'JPEGImages/'+line+'.jpg'):
            vi.write(data_dir+'JPEGImages/'+line+'.jpg' + '\n')
            assert (os.path.isfile(data_dir+'SegmentationClass/encoded/'+line +
                                   '.npy')), "No matching label file for image : " + line + '.jpg'
            vl.write(data_dir+'SegmentationClass/encoded/'+line + '.npy' + '\n')

    ti.close()
    tl.close()
    vi.close()
    vl.close()


def main():
    '''
    Arguments:
    train-file = txt file containing randomly selected image filenames to be taken as training set.
    val-file = txt file containing randomly selected image filenames to be taken as validation set.
    data-dir = dataset directory
    Usage: python dataset_utils.py --train-file="" --val-file="" --data_dir=""
    '''

    args = get_args()
    data_dir = args.data_dir

    if not os.path.exists(data_dir+'SegmentationClass/' + 'encoded/'):
        os.makedirs(data_dir+'SegmentationClass/' + 'encoded/')
    for filename in os.listdir(data_dir+'SegmentationClass/'):
        if os.path.isdir(data_dir+'SegmentationClass/' + filename):
            continue
        label = imread(data_dir+'SegmentationClass/' +
                       filename).astype('float32')
        label = encode_label(label)
        np.save(data_dir+'SegmentationClass/' + 'encoded/' +
                filename.split('.')[0] + '.npy', label)

    generate_path_files(args.data_dir, args.train_file, args.val_file)


if __name__ == '__main__':
    main()


# semantic-segmentation/deeplabv3plus/download_pretrained_tf_deeplabv3plus_coco_voc_trainaug.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 tarfile
import nnabla


def download(url):
    from nnabla.utils.data_source_loader import download as dl
    dl(url, url.split('/')[-1], False)


def main():
    model_url = 'http://download.tensorflow.org/models/deeplabv3_pascal_train_aug_2018_01_04.tar.gz'

    print('Downloading Xception COCO VOC trainaug pretrained model ...')
    download(model_url)

    # untar the downloaded compressed file
    compressed_file = tarfile.open(model_url.split('/')[-1], "r:gz")
    compressed_file.extractall()
    compressed_file.close()


if __name__ == '__main__':
    main()


# semantic-segmentation/deeplabv3plus/eval.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 train
import model_inference

import nnabla as nn
import nnabla.logger as logger
import nnabla.monitor as M
import numpy as np

from args import get_args
from segmentation_data import data_iterator_segmentation


def validate(args):

   # load trained param file
    _ = nn.load_parameters(args.model_load_path)

    # get data iterator
    vdata = data_iterator_segmentation(args.val_samples, args.batch_size, args.val_dir,
                                       args.val_label_dir, target_width=args.image_width, target_height=args.image_height)

    # get deeplabv3plus model
    v_model = train.get_model(args, test=True)
    v_model.pred.persistent = True  # Not clearing buffer of pred in forward
    v_pred2 = v_model.pred.unlinked()

    # Create monitor
    monitor = M.Monitor(args.monitor_path)
    monitor_miou = M.MonitorSeries("mean IOU", monitor, interval=1)

    l = 0.0
    e = 0.0
    vmiou = 0.
    # Evaluation loop
    for j in range(args.val_samples // args.batch_size):
        images, labels, masks = vdata.next()
        v_model.image.d = images
        v_model.label.d = labels
        v_model.mask.d = masks
        v_model.pred.forward(clear_buffer=True)
        miou = train.compute_miou(
            args.num_class, labels, np.argmax(v_model.pred.d, axis=1), masks)
        vmiou += miou
        print(j, miou)

    monitor_miou.add(0, vmiou / (args.val_samples / args.batch_size))
    return vmiou / args.val_samples


def main():
    args = get_args()
    rng = np.random.RandomState(1223)

    # Get context
    from nnabla.ext_utils import get_extension_context
    logger.info("Running in %s" % args.context)
    ctx = get_extension_context(
        args.context, device_id=args.device_id, type_config=args.type_config)
    nn.set_default_context(ctx)

    miou = validate(args)


if __name__ == '__main__':
    '''
    Usage : python eval.py --model-load-path=/path to trained .h5 file --val-samples=no. of validation examples --batch-size=1 --val-dir=file containing paths to val images --val-label-dir=file containing paths to val labels --num-class=no. of categories
    '''

    main()


# semantic-segmentation/deeplabv3plus/image_preprocess.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 cv2
import random
import numpy as np


def convert_to_rgb(image, label=None):

    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    if label.ndim == 2:
        label = cv2.cvtColor(label, cv2.COLOR_BGR2RGB)
    return image, label


def resize(image, desired_size):

    old_size = image.shape[:2]  # old_size is in (height, width) format
    ratio = min(np.divide(desired_size, old_size))
    new_size = (int(old_size[0]*ratio), int(old_size[1]*ratio))

    # new_size should be in (width, height) format
    if image.shape[2] == 1:
        image = cv2.resize(
            image, (new_size[1], new_size[0]), interpolation=cv2.INTER_NEAREST)
        return image
    image = cv2.resize(image, (new_size[1], new_size[0]))

    return image


def random_flip(image, label=None):
    # random no from 0-1(exclude 1)
    rand = random.uniform(0, 1)
    prob = 0.5
    if rand <= prob:
        image = cv2.flip(image, 1)
        label = cv2.flip(label, 1)
        if label is not None:
            label = np.expand_dims(label, axis=2)

    return image, label


def random_crop(image, label, crop_h, crop_w):
    h, w, c = image.shape
    if h >= crop_h and w >= crop_w:
        max_offset_h = h - crop_h + 1
        max_offset_w = w - crop_w + 1

        # rand no between 0 to max_offset --> int
        offset_h = int(random.uniform(0, max_offset_h))
        offset_w = int(random.uniform(0, max_offset_w))

        # crop image to reach crop_h begin from offset_h, crop_w begin from offset_w
        crop_img = image[offset_h:crop_h+offset_h, offset_w:crop_w+offset_w]
        crop_lab = label[offset_h:crop_h+offset_h, offset_w:crop_w+offset_w]

    return crop_img, crop_lab


def random_scale(image, label):

    min_scale_factor = 0.5
    max_scale_factor = 2.
    scale_factor_step_size = 0.25

    num_steps = int((max_scale_factor - min_scale_factor) /
                    scale_factor_step_size + 1)
    scale_factors = np.linspace(min_scale_factor, max_scale_factor, num_steps)
    random.shuffle(scale_factors)

    scale = scale_factors[0].astype(np.float32)
    scaled_image = cv2.resize(image, (int(
        image.shape[1]*scale), int(image.shape[0]*scale)), interpolation=cv2.INTER_LINEAR)

    scaled_label = cv2.resize(label, (int(
        label.shape[1]*scale), int(label.shape[0]*scale)), interpolation=cv2.INTER_NEAREST)
    if scaled_label.ndim == 2:
        scaled_label = np.expand_dims(scaled_label, axis=2)

    return scaled_image, scaled_label


def zero_mean_unit_range(image):
    return ((2.0 / 255.0) * image - 1.0)


def cast(image, label=None):
    image = image.astype(np.float32)
    if label is not None:
        label = label.astype(np.int32)
    return image, label


def pad(image, label=None, desired_size=(513, 513)):

    # mean pixel value
    color = [127.5, 127.5, 127.5]

    old_size = image.shape[:2]  # old_size is in (height, width) format

    new_size = tuple([x + max(ds - x, 0)
                      for x, ds in zip(old_size, desired_size)])

    # Pad image with mean pixel value
    new_im = cv2.copyMakeBorder(
        image, 0, new_size[0]-old_size[0], 0, new_size[1]-old_size[1], cv2.BORDER_CONSTANT, value=color)
    if label is not None:
        label = cv2.copyMakeBorder(
            label, 0, new_size[0]-old_size[0], 0, new_size[1]-old_size[1], cv2.BORDER_CONSTANT, value=[255, 255, 255])
        if label.ndim == 2:
            label = np.expand_dims(label, axis=2)

    return new_im, label


def create_mask(label):
    # cretae masks of 1, 0 for ignore labels
    mask = (label != 255)
    mask = mask.astype(np.int32)
    label[label == 255] = 0
    label = label.astype(np.int32)
    return label, mask


def preprocess_image_and_label(image, label=None, target_width=513, target_height=513, train=True):
    '''
    image - 4d array (batch, ch, h, w)
    label - 4d array (batch, ch, h, w)

    '''

    crop_w = target_width
    crop_h = target_height

    if train == True:
        image, label = cast(image, label)

        image, label = random_scale(image, label)

        image, label = pad(image, label, desired_size=(
            target_height, target_width))

        image, label = random_crop(image, label, target_height, target_width)

        image, label = random_flip(image, label)

        image = zero_mean_unit_range(image)
        if label is not None:
            label, mask = create_mask(label)

        return image, label, mask

    else:
        image, label = cast(image)
        image = resize(image, (target_height, target_width))
        image, label = pad(image, desired_size=(crop_h, crop_w))
        image = zero_mean_unit_range(image)
        return image


# semantic-segmentation/deeplabv3plus/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

import numpy as np
import math

import xception_65 as xception


def atrous_spatial_pyramid_pooling(x, output_stride, test=False, fix_params=False):
    if output_stride not in [8, 16]:
        raise ValueError('output_stride neither 8 nor 16.')

    depth = 256
    atrous_rates = [6, 12, 18]
    if output_stride == 8:
        atrous_rates = [2*rate for rate in atrous_rates]

    with nn.parameter_scope("aspp0"):
        atrous_conv0 = PF.convolution(
            x, depth, (1, 1), with_bias=False, fix_parameters=fix_params)
        atrous_conv0 = F.relu(PF.batch_normalization(
            atrous_conv0, batch_stat=not test, fix_parameters=fix_params))

    atrous_conv = []
    for i in range(3):
        with nn.parameter_scope("aspp"+str(i+1)):
            atrous_conv.append(xception.separable_conv_with_bn(x, depth, stride=False, aspp=True,
                                                               atrous_rate=atrous_rates[i], last_block=True, eps=1e-05, test=test, fix_params=fix_params))

    # Image-level features
    with nn.parameter_scope("image_pooling"):
        x_shape = x.shape[2]
        h = F.average_pooling(
            x, (x.shape[2], x.shape[3]), stride=(x.shape[2], x.shape[2]))

        h = PF.convolution(h, depth, (1, 1), with_bias=False,
                           fix_parameters=fix_params)
        h = F.relu(PF.batch_normalization(
            h, batch_stat=not test, fix_parameters=fix_params))

        h = F.interpolate(h, output_size=(
            x.shape[2], x.shape[3]), mode='linear')

    with nn.parameter_scope("concat_projection"):
        h5 = F.concatenate(
            h, atrous_conv0, atrous_conv[0], atrous_conv[1], atrous_conv[2], axis=1)

    return h5


def decoder(x, upsampled, num_classes, test=False, fix_params=False):

    # Project low-level features
    with nn.parameter_scope("feature_projection0"):
        h = PF.convolution(x, 48, (1, 1), with_bias=False,
                           fix_parameters=fix_params)
        h = F.relu(PF.batch_normalization(
            h, batch_stat=not test, fix_parameters=fix_params))

    h = F.concatenate(upsampled, h, axis=1)

    for i in range(2):
        with nn.parameter_scope("decoder_conv"+str(i)):
            h = xception.separable_conv_with_bn(
                h, 256, last_block=True, eps=1e-05, out=True, test=test, fix_params=fix_params)

    with nn.parameter_scope("logits/affine"):
        h = PF.convolution(h, num_classes, (1, 1), with_bias=True,
                           fix_parameters=fix_params)  # no activation

    with nn.parameter_scope("upsample2"):
        h = F.interpolate(h, output_size=(
            h.shape[2]*4 - 3, h.shape[3]*4 - 3), mode='linear')

    return h


def deeplabv3plus_model(x, output_stride, num_classes, test=False, fix_params=False):
    '''Encoder 
    '''
    # Get decoder endpoints from backbone
    endpoints = xception.xception_65(x, test=test, fix_params=fix_params)
    low_level_features = endpoints['Decoder End Point 1']

    encoder_output = atrous_spatial_pyramid_pooling(
        endpoints['Decoder End Point 2'], output_stride, test=test, fix_params=fix_params)

    with nn.parameter_scope("concat_projection"):
        encoder_output = PF.convolution(
            encoder_output, 256, (1, 1), with_bias=False, fix_parameters=fix_params)
        encoder_output = F.relu(PF.batch_normalization(
            encoder_output, batch_stat=not test, fix_parameters=fix_params))

    '''Decoder 
    '''
    with nn.parameter_scope("decoder"):
        with nn.parameter_scope("upsample1"):
            upsampled = F.interpolate(encoder_output, output_size=(
                low_level_features.shape[2], low_level_features.shape[3]), mode='linear')

        h = decoder(low_level_features, upsampled, num_classes,
                    test=test, fix_params=fix_params)

    return h


# semantic-segmentation/deeplabv3plus/model_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 cv2
import numpy as np
import nnabla as nn
from nnabla.logger import logger
import model
import model as net
from args import get_args
import imageio
from PIL import Image
import image_preprocess
import dataset_utils
import matplotlib.pyplot as plt
import time


def visualize(label):
    h, w = label.shape
    vis = np.zeros((h, w, 3), dtype=np.int32)
    clr_map = dataset_utils.get_color()
    for y in range(h):
        for x in range(w):
            if label[y][x] == 22:
                vis[y][x] = 0
            else:
                vis[y][x] = clr_map[label[y][x]]

    #plt.imshow(vis, interpolation='none')
    # plt.show()
    imageio.imwrite('output.png', vis)


def post_process(output, old_size, target_size):
    ratio = min(np.divide(target_size, old_size))
    new_size = (int(old_size[0]*ratio), int(old_size[1]*ratio))

    post_processed = output[0:new_size[0], 0:new_size[1]]

    post_processed = cv2.resize(
        post_processed, (old_size[1], old_size[0]), interpolation=cv2.INTER_NEAREST)

    return post_processed


def main():
    args = get_args()
    rng = np.random.RandomState(1223)

    # Get context
    from nnabla.ext_utils import get_extension_context, import_extension_module
    logger.info("Running in %s" % args.context)
    ctx = get_extension_context(
        args.context, device_id=args.device_id, type_config=args.type_config)
    nn.set_default_context(ctx)
    ext = import_extension_module(args.context)

    # read label file
    f = open(args.label_file_path, "r")
    labels_dict = f.readlines()

    # Load parameters
    _ = nn.load_parameters(args.model_load_path)

    # Build a Deeplab v3+ network
    x = nn.Variable(
        (1, 3, args.image_height, args.image_width), need_grad=False)
    y = net.deeplabv3plus_model(
        x, args.output_stride, args.num_class, test=True)

    # preprocess image
    image = imageio.imread(args.test_image_file, as_gray=False, pilmode="RGB")
    #image = imread(args.test_image_file).astype('float32')
    orig_h, orig_w, orig_c = image.shape
    old_size = (orig_h, orig_w)

    input_array = image_preprocess.preprocess_image_and_label(
        image, label=None, target_width=args.image_width, target_height=args.image_height, train=False)
    print('Input', input_array.shape)
    input_array = np.transpose(input_array, (2, 0, 1))
    input_array = np.reshape(
        input_array, (1, input_array.shape[0], input_array.shape[1], input_array.shape[2]))

    # Compute inference and inference time
    t = time.time()

    x.d = input_array
    y.forward(clear_buffer=True)
    print("done")
    available_devices = ext.get_devices()
    ext.device_synchronize(available_devices[0])
    ext.clear_memory_cache()

    elapsed = time.time() - t
    print('Inference time : %s seconds' % (elapsed))

    output = np.argmax(y.d, axis=1)  # (batch,h,w)

    # Apply post processing
    post_processed = post_process(
        output[0], old_size, (args.image_height, args.image_width))

    # Get the classes predicted
    predicted_classes = np.unique(post_processed)
    for i in range(predicted_classes.shape[0]):
        print('Classes Segmented: ', labels_dict[predicted_classes[i]])

    # Visualize inference result
    visualize(post_processed)


if __name__ == '__main__':
    '''
    Usage : python model_inference.py --model-load-path=/path to trained .h5 file --image-width=target width for input image --test-image-file=image file for inference --num-class=no. of categories --label-file-path=txt file having categories --output-stride=8 or 16
    '''

    main()


# semantic-segmentation/deeplabv3plus/plot_accuracy.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 eval
import nnabla.logger as logger
import nnabla.monitor as M
import numpy as np
import nnabla as nn
from args import get_args
import os
import matplotlib.pyplot as plt


def main():
    args = get_args()
    rng = np.random.RandomState(1223)

    # Get context
    from nnabla.ext_utils import get_extension_context
    logger.info("Running in %s" % args.context)
    ctx = get_extension_context(
        args.context, device_id=args.device_id, type_config=args.type_config)
    nn.set_default_context(ctx)
    iterations = []
    mean_iou = []
    model_dir = args.model_load_path
    for filename in os.listdir(model_dir):
        args.model_load_path = model_dir+filename
        miou = eval.validate(args)
        iterations.append(filename.split('.')[0])
        mean_iou.append(miou)

    for i in range(len(iterations)):
        iterations[i] = iterations[i].replace('param_', '')

    itr = list(map(int, iterations))

    # Plot Iterations Vs mIOU
    plt.axes([0, max(itr), 0.0, 1.0])
    plt.xlabel('Iterations')
    plt.ylabel('Accuracy - mIOU')
    plt.scatter(itr, mean_iou)
    plt.show()

    print(iterations)
    print(mean_iou)
    with open('iterations.txt', 'w') as f:
        for item in iterations:
            f.write('%s\n' % item)
    with open('miou.txt', 'w') as f2:
        for item in mean_iou:
            f2.write('%s\n' % item)

    #plt.plot(iterations, mean_iou)
if __name__ == '__main__':
    main()


# semantic-segmentation/deeplabv3plus/prepare_lfw_data.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 numpy as np
import os
from scipy.misc import imread
from args import get_args
import nnabla.utils.image_utils as utils


def get_prefix(word):
    '''
    Get the prefix to be used to search for image/label
    '''
    if len(str(int(word))) == 1:
        prefix = '_000'
    elif len(str(int(word))) == 2:
        prefix = '_00'
    else:
        prefix = '_0'
    return prefix


def encode_label(label):
    '''
    Converting pixel values to corresponding class numbers. Assuming that the input label in 2-dim(h,w)
    '''

    h, w, _ = label.shape
    new_label = np.zeros((h, w, 1), dtype=np.int32)

    for c in range(3):
        new_label[label[..., c] != 0] = c

    return new_label


def encode_and_write_to_path_files(filename, data_dir, ti, tl):
    '''
    Calling encode_label for each label and writing image and label paths to path files
    '''
    train_f = open(filename, 'r')
    label_path = data_dir+'parts_lfw_funneled_gt_images/'
    image_path = data_dir+'lfw_funneled/'
    for line in train_f:
        words = line.split(' ')
        prefix = get_prefix(words[1])

        if os.path.isdir(image_path + words[0] + '/'):
            ti.write(image_path + words[0] + '/' + words[0] +
                     prefix+str(int(words[1]))+'.jpg' + '\n')
            assert (os.path.isfile(label_path + words[0]+prefix+str(int(words[1]))+'.ppm')
                    ), "No matching label file for image : " + words[0]+prefix+str(int(words[1])) + '.jpg'
            label = utils.imread(
                label_path + words[0]+prefix+str(int(words[1]))+'.ppm')
            label = encode_label(label)
            np.save(label_path + 'encoded/' +
                    words[0]+prefix+str(int(words[1])) + '.npy', label)
            tl.write(label_path + 'encoded/' +
                     words[0]+prefix+str(int(words[1])) + '.npy' + '\n')


# this method should generate lfw_train_image.txt, lfw_train_label.txt, lfw_val_image.txt, lfw_val_label.txt
def generate_path_files(data_dir, train_file, val_file):

    ti = open('lfw_train_image.txt', 'w')
    tl = open('lfw_train_label.txt', 'w')
    vi = open('lfw_val_image.txt', 'w')
    vl = open('lfw_val_label.txt', 'w')

    if not os.path.exists(data_dir+'parts_lfw_funneled_gt_images/' + 'encoded/'):
        os.makedirs(data_dir+'parts_lfw_funneled_gt_images/' + 'encoded/')

    encode_and_write_to_path_files(train_file, data_dir, ti, tl)

    encode_and_write_to_path_files(val_file, data_dir, vi, vl)

    ti.close()
    tl.close()
    vi.close()
    vl.close()


def main():
    '''
    Arguments:
    train-file = txt file containing randomly selected image filenames to be taken as training set.
    val-file = txt file containing randomly selected image filenames to be taken as validation set.
    data-dir = dataset directory
    Usage: python dataset_utils.py --train-file="" --val-file="" --data_dir=""
    '''

    args = get_args()
    data_dir = args.data_dir

    generate_path_files(data_dir, args.train_file, args.val_file)


if __name__ == '__main__':
    main()


# semantic-segmentation/deeplabv3plus/segmentation_data.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 nnabla.utils.data_iterator import data_iterator_simple
import cv2
import image_preprocess
import numpy as np
import time


def data_iterator_segmentation(num_examples, batch_size, image_path_file, label_path_file, rng=None, target_width=513, target_height=513, train=True):

    image_paths = load_paths(image_path_file)
    label_paths = load_paths(label_path_file)

    def image_label_load_func(i):
        '''
        Returns:
            image: c x h x w array
            label: c x h x w array
            mask: c x h x w array
        '''

        img = cv2.imread(image_paths[i]).astype('float32')
        b, g, r = cv2.split(img)
        img = cv2.merge([r, g, b])
        if 'png' in label_paths[i]:
            lab = imageio.imread(
                label_paths[i], as_gray=False, pilmode="RGB").astype('int32')
        else:
            lab = np.load(label_paths[i], allow_pickle=True).astype('int32')
        if lab.ndim == 2:
            lab = lab[..., None]
        # Compute image preprocessing time
        #t = time.time()
        img, lab, mask = image_preprocess.preprocess_image_and_label(
            img, lab, target_width, target_height, train=train)
        #elapsed = time.time() - t

        return np.rollaxis(img, 2), np.rollaxis(lab, 2), np.rollaxis(mask, 2)

    return data_iterator_simple(image_label_load_func, num_examples, batch_size, shuffle=True, rng=rng, with_file_cache=False, with_memory_cache=False)


def load_paths(path_file):
    text_file = open(path_file, "r")
    lines = [line[:-1] for line in text_file]

    return lines


# semantic-segmentation/deeplabv3plus/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.

import numpy as np

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

import nnabla.utils.save as save
from args import get_args
from segmentation_data import data_iterator_segmentation
import model

import os
from collections import namedtuple
from _checkpoint_nnp_util import save_checkpoint, load_checkpoint, save_nnp


def get_model(args, test=False):
    """
    Create computation graph and variables.

    """

    image = nn.Variable(
        [args.batch_size, 3, args.image_height, args.image_width])
    label = nn.Variable(
        [args.batch_size, 1, args.image_height, args.image_width])
    mask = nn.Variable(
        [args.batch_size, 1, args.image_height, args.image_width])

    pred = model.deeplabv3plus_model(
        image, args.output_stride, args.num_class, test=test, fix_params=False)

    if pred.shape != label.shape:
        pred = F.interpolate(pred, output_size=(
            label.shape[2], label.shape[3]), mode='linear')

    loss = F.sum(F.softmax_cross_entropy(
        pred, label, axis=1) * mask) / F.sum(mask)
    Model = namedtuple('Model', ['image', 'label', 'mask', 'pred', 'loss'])
    return Model(image, label, mask, pred, loss)


def one_hot_encode(label, num_class):
    # label is 2-d of shape (batch, pixels(513x513))
    one_hot_vector = np.zeros(
        (label.shape[0], label.shape[1], num_class+1), dtype=np.int32)

    batch_id = np.arange(label.shape[0]).reshape(label.shape[0], 1)
    pixel_id = np.tile(np.arange(label.shape[1]), (label.shape[0], 1))

    one_hot_vector[batch_id, pixel_id, label.astype(np.int32)] = 1

    return one_hot_vector


def compute_miou(num_classes, gt, pred, mask):

    gt = np.squeeze(np.transpose(gt, (1, 0, 2, 3)), axis=0)
    gt = gt.reshape(gt.shape[0], gt.shape[1]*gt.shape[2])
    gt = one_hot_encode(gt, num_classes)

    pred = pred.reshape(pred.shape[0], pred.shape[1]*pred.shape[2])
    pred = one_hot_encode(pred, num_classes)

    # Now, gt: [B, P, C] ; pred: [B, P, C]
    mask = np.squeeze(np.transpose(mask, (1, 0, 2, 3)), axis=0)
    mask = mask.reshape(mask.shape[0], mask.shape[1]
                        * mask.shape[2])  # mask: [B, P]
    mask = mask[..., None]

    # gt: [B, P, C], mask: [B, P, 1]
    numer = np.sum(np.logical_and(gt, pred) * mask, axis=1)
    denom = np.sum(np.logical_or(gt, pred) * mask, axis=1)
    iou = numer / np.maximum(denom, 1)

    cat_mask = np.max(gt, axis=1)
    iou_per_image = np.sum(iou * cat_mask, axis=1) / np.sum(cat_mask, axis=1)
    miou = np.mean(iou_per_image)

    return miou.mean()


def train():
    """
    Main script.
    """

    args = get_args()

    _ = nn.load_parameters(args.pretrained_model_path)
    if args.fine_tune:
        nnabla.parameter.pop_parameter('decoder/logits/affine/conv/W')
        nnabla.parameter.pop_parameter('decoder/logits/affine/conv/b')

    n_train_samples = args.train_samples
    n_val_samples = args.val_samples
    distributed = args.distributed
    compute_acc = args.compute_acc

    if distributed:
        # Communicator and Context
        from nnabla.ext_utils import get_extension_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
        device_id = mpi_rank
        ctx.device_id = str(device_id)
        nn.set_default_context(ctx)
    else:
        # Get context.
        from nnabla.ext_utils import get_extension_context
        extension_module = args.context
        if args.context is None:
            extension_module = 'cpu'
        logger.info("Running in %s" % extension_module)
        ctx = get_extension_context(
            extension_module, device_id=args.device_id, type_config=args.type_config)
        nn.set_default_context(ctx)
        n_devices = 1
        device_id = 0

    # training data
    data = data_iterator_segmentation(
            args.train_samples, args.batch_size, args.train_dir, args.train_label_dir, target_width=args.image_width, target_height=args.image_height)
    # validation data
    vdata = data_iterator_segmentation(args.val_samples, args.batch_size, args.val_dir,
                                       args.val_label_dir, target_width=args.image_width, target_height=args.image_height)

    if distributed:
        data = data.slice(
            rng=None, num_of_slices=n_devices, slice_pos=device_id)
        vdata = vdata.slice(
            rng=None, num_of_slices=n_devices, slice_pos=device_id)
    num_classes = args.num_class

    # Workaround to start with the same initialized weights for all workers.
    np.random.seed(313)
    t_model = get_model(
        args, test=False)
    t_model.pred.persistent = True  # Not clearing buffer of pred in backward
    t_pred2 = t_model.pred.unlinked()
    t_e = F.sum(F.top_n_error(t_pred2, t_model.label, axis=1)
                * t_model.mask) / F.sum(t_model.mask)

    v_model = get_model(
        args, test=True)
    v_model.pred.persistent = True  # Not clearing buffer of pred in forward
    v_pred2 = v_model.pred.unlinked()
    v_e = F.sum(F.top_n_error(v_pred2, v_model.label, axis=1)
                * v_model.mask) / F.sum(t_model.mask)

    # Create Solver
    solver = S.Momentum(args.learning_rate, 0.9)
    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)

    # Setting warmup.
    base_lr = args.learning_rate / n_devices
    warmup_iter = int(1. * n_train_samples /
                      args.batch_size / args.accum_grad / n_devices) * args.warmup_epoch
    warmup_slope = base_lr * (n_devices - 1) / warmup_iter
    solver.set_learning_rate(base_lr)

    # Create monitor
    import nnabla.monitor as M
    monitor = M.Monitor(args.monitor_path)
    monitor_loss = M.MonitorSeries("Training loss", monitor, interval=10)
    monitor_err = M.MonitorSeries("Training error", monitor, interval=10)
    monitor_vloss = M.MonitorSeries("Validation loss", monitor, interval=1)
    monitor_verr = M.MonitorSeries("Validation error", monitor, interval=1)
    monitor_time = M.MonitorTimeElapsed("Training time", monitor, interval=10)
    monitor_miou = M.MonitorSeries("mean IOU", monitor, interval=10)
    monitor_vtime = M.MonitorTimeElapsed(
        "Validation time", monitor, interval=1)

    # save_nnp
    contents = save_nnp({'x': v_model.image}, {
                        'y': v_model.pred}, args.batch_size)
    save.save(os.path.join(args.model_save_path,
                           'Deeplabv3plus_result_epoch0.nnp'), contents, variable_batch_size=False)

    # Training loop
    for i in range(start_point, int(args.max_iter / n_devices)):
        # Save parameters
        if i % (args.model_save_interval // n_devices) == 0 and device_id == 0:
            save_checkpoint(args.model_save_path, i, solver)
        # Validation
        if i % (args.val_interval // n_devices) == 0 and i != 0:
            vmiou_local = 0.
            val_iter_local = n_val_samples // args.batch_size
            vl_local = nn.NdArray()
            vl_local.zero()
            ve_local = nn.NdArray()
            ve_local.zero()
            for j in range(val_iter_local):
                images, labels, masks = vdata.next()
                v_model.image.d = images
                v_model.label.d = labels
                v_model.mask.d = masks
                v_model.image.data.cast(np.float32, ctx)
                v_model.label.data.cast(np.int32, ctx)
                v_model.loss.forward(clear_buffer=True)
                v_e.forward(clear_buffer=True)
                vl_local += v_model.loss.data
                ve_local += v_e.data
                # Mean IOU computation
                if compute_acc:
                    vmiou_local += compute_miou(num_classes, labels,
                                                np.argmax(v_model.pred.d, axis=1), masks)

            vl_local /= val_iter_local
            ve_local /= val_iter_local
            if compute_acc:
                vmiou_local /= val_iter_local
                vmiou_ndarray = nn.NdArray.from_numpy_array(
                    np.array(vmiou_local))
            if distributed:
                comm.all_reduce(vl_local, division=True, inplace=True)
                comm.all_reduce(ve_local, division=True, inplace=True)
                if compute_acc:
                    comm.all_reduce(vmiou_ndarray, division=True, inplace=True)

            if device_id == 0:
                monitor_vloss.add(i * n_devices, vl_local.data.copy())
                monitor_verr.add(i * n_devices, ve_local.data.copy())
                if compute_acc:
                    monitor_miou.add(i * n_devices, vmiou_local)
                monitor_vtime.add(i * n_devices)

        # Training
        l = 0.0
        e = 0.0
        solver.zero_grad()

        e_acc = nn.NdArray(t_e.shape)
        e_acc.zero()
        l_acc = nn.NdArray(t_model.loss.shape)
        l_acc.zero()
        # Gradient accumulation loop
        for j in range(args.accum_grad):
            images, labels, masks = data.next()
            t_model.image.d = images
            t_model.label.d = labels
            t_model.mask.d = masks
            t_model.image.data.cast(np.float32, ctx)
            t_model.label.data.cast(np.int32, ctx)
            t_model.loss.forward(clear_no_need_grad=True)
            t_model.loss.backward(clear_buffer=True)  # Accumulating gradients
            t_e.forward(clear_buffer=True)
            e_acc += t_e.data
            l_acc += t_model.loss.data

        # AllReduce
        if distributed:
            params = [x.grad for x in nn.get_parameters().values()]
            comm.all_reduce(params, division=False, inplace=False)
            comm.all_reduce(l_acc, division=True, inplace=True)
            comm.all_reduce(e_acc, division=True, inplace=True)
        solver.scale_grad(1./args.accum_grad)
        solver.weight_decay(args.weight_decay)
        solver.update()

        # Linear Warmup
        if i <= warmup_iter:
            lr = base_lr + warmup_slope * i
            solver.set_learning_rate(lr)

        if distributed:
            # Synchronize by averaging the weights over devices using allreduce
            if (i+1) % args.sync_weight_every_itr == 0:
                weights = [x.data for x in nn.get_parameters().values()]
                comm.all_reduce(weights, division=True, inplace=True)

        if device_id == 0:
            monitor_loss.add(
                i * n_devices, (l_acc / args.accum_grad).data.copy())
            monitor_err.add(
                i * n_devices, (e_acc / args.accum_grad).data.copy())
            monitor_time.add(i * n_devices)

        # Learning rate decay at scheduled iter --> changed to poly learning rate decay policy
        # if i in args.learning_rate_decay_at:
        solver.set_learning_rate(base_lr * ((1 - i / args.max_iter)**0.1))

    if device_id == 0:
        nn.save_parameters(os.path.join(args.model_save_path,
                                        'param_%06d.h5' % args.max_iter))

    contents = save_nnp({'x': v_model.image}, {
                        'y': v_model.pred}, args.batch_size)
    save.save(os.path.join(args.model_save_path,
                           'Deeplabv3plus_result.nnp'), contents, variable_batch_size=False)


if __name__ == '__main__':
    train()


# semantic-segmentation/deeplabv3plus/xception_65.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.parametric_functions as PF
import nnabla.functions as F
from nnabla.logger import logger
import numpy as np

endpoints = {}


def block(x, sh, num_units, stride, end_point, depth_list, sh_cut, shortcut_channels, act_fn=True, atrous_conv=False, atrous_rate=1, res=False, last_block=False, test=False, fix_params=False):

    for i in range(0, num_units):

        with nn.parameter_scope("unit_"+str(i+1)):

            x = unit(x, depth_list, stride, end_point, act_fn=act_fn, atrous_conv=atrous_conv,
                     atrous_rate=atrous_rate, last_block=last_block, test=test, fix_params=fix_params)

            if sh_cut == True:
                sh = shortcut(sh, shortcut_channels, not atrous_conv,
                              test=test, fix_params=fix_params)
                x = x+sh
            if res == True:
                x = x+sh
                sh = x

    return x


def unit(x, depth_list, stride, end_point, act_fn=True, atrous_conv=False, atrous_rate=1, last_block=False, test=False, fix_params=False):

    if last_block == False:
        x = F.relu(x)
    for i in range(0, len(depth_list)):
        if i == 2:
            act_fn = False
        with nn.parameter_scope("separable_conv"+str(i+1)):

            if end_point == True and i == 1:
                x = separable_conv_with_bn(x, depth_list[i], stride=False, aspp=False, atrous_rate=atrous_rate,
                                           act_fn=act_fn, last_block=last_block, end_point=end_point, test=test, fix_params=fix_params)

            else:
                if stride == True and i == 2:
                    x = separable_conv_with_bn(x, depth_list[i], stride=True, aspp=atrous_conv, atrous_rate=atrous_rate,
                                               act_fn=act_fn, last_block=last_block, test=test, fix_params=fix_params)
                else:
                    x = separable_conv_with_bn(x, depth_list[i], stride=False, aspp=atrous_conv, atrous_rate=atrous_rate,
                                               act_fn=act_fn, last_block=last_block, test=test, fix_params=fix_params)

    return x


def shortcut(x, f, stride=True, test=False, fix_params=False):

    with nn.parameter_scope("shortcut"):
        if(stride == False):
            h = PF.convolution(x, f, (1, 1), with_bias=False,
                               fix_parameters=fix_params)
        else:
            h = PF.convolution(x, f, (1, 1), stride=(
                2, 2), with_bias=False, fix_parameters=fix_params)
        h = PF.batch_normalization(
            h, batch_stat=not test, eps=1e-03, fix_parameters=fix_params)
    return h


def separable_conv_with_bn(x, f, stride=False, aspp=False, atrous_rate=1, act_fn=True, last_block=False, end_point=False, eps=1e-03, out=False, test=False, fix_params=False):

    with nn.parameter_scope("depthwise"):
        if(stride == True):
            h = PF.depthwise_convolution(x, (3, 3), stride=(2, 2), pad=(
                1, 1), with_bias=False, fix_parameters=fix_params)
        elif(aspp == True):
            h = PF.depthwise_convolution(x, (3, 3), pad=(atrous_rate, atrous_rate), stride=(
                1, 1), dilation=(atrous_rate, atrous_rate), with_bias=False, fix_parameters=fix_params)

        else:
            h = PF.depthwise_convolution(x, (3, 3), pad=(
                1, 1), with_bias=False, fix_parameters=fix_params)

        h = PF.batch_normalization(
            h, batch_stat=not test, eps=eps, fix_parameters=fix_params)
        if last_block == True:
            h = F.relu(h)

    with nn.parameter_scope("pointwise"):
        h = PF.convolution(h, f, (1, 1), stride=(
            1, 1), with_bias=False, fix_parameters=fix_params)
        h = PF.batch_normalization(
            h, batch_stat=not test, eps=eps, fix_parameters=fix_params)
        if end_point == True:
            global endpoints
            endpoints['Decoder End Point 1'] = h

        if act_fn == True:
            h = F.relu(h)

    return h


def entry_flow(x, num_blocks, depth_list, test=False, fix_params=False):

    shortcut_channels = [128, 256, 728]
    global endpoints

    with nn.parameter_scope("1"):
        h = PF.convolution(x, 32, (3, 3), pad=(1, 1), stride=(
            2, 2), with_bias=False, fix_parameters=fix_params)
        h = F.relu(PF.batch_normalization(h, batch_stat=not test,
                                          eps=1e-03, fix_parameters=fix_params))

    with nn.parameter_scope("2"):
        h = PF.convolution(h, 64, (3, 3), pad=(
            1, 1), with_bias=False, fix_parameters=fix_params)
        h = F.relu(PF.batch_normalization(h, batch_stat=not test,
                                          eps=1e-03, fix_parameters=fix_params))

    x = h
    sh = x
    for i in range(0, num_blocks):
        with nn.parameter_scope("block"+str(i+1)):

            if i == 1:
                x = block(x, sh, 1, True, True,
                          depth_list[i], True, shortcut_channels[i], test=test)
            else:
                x = block(x, sh, 1, True, False,
                          depth_list[i], True, shortcut_channels[i], test=test)
            sh = x
        if i == 2:
            endpoints['conv1'] = x

    return x


def middle_flow(x, num_blocks, depth_list, test=False, fix_params=False):

    shortcut_channels = [0]
    sh = x
    for i in range(0, num_blocks):
        with nn.parameter_scope("block"+str(i+1)):
            x = block(x, sh, 16, False, False,
                      depth_list[i], False, shortcut_channels, res=True, test=test)
            sh = x

    return x


def exit_flow(x, num_blocks, depth_list, test=False, fix_params=False):

    shortcut_channels = [1024, 0]
    sh = x
    for i in range(0, num_blocks):
        with nn.parameter_scope("block"+str(i+1)):
            if i == 0:
                x = block(x, sh, 1, False, False,
                          depth_list[i], True, shortcut_channels[i], atrous_conv=True, test=test)

            else:
                x = block(x, sh, 1, False, False, depth_list[i], False, shortcut_channels[i],
                          atrous_conv=True, atrous_rate=2, last_block=True, test=test)

            sh = x

    return x


def xception_65(x, test=False, fix_params=False):

    entry_flow_depth_list = [[128, 128, 128], [256, 256, 256], [728, 728, 728]]
    middle_flow_depth_list = [[728, 728, 728]]
    exit_flow_depth_list = [[728, 1024, 1024], [1536, 1536, 2048]]

    with nn.parameter_scope("xception_65"):
        with nn.parameter_scope("entry_flow"):
            x = entry_flow(x, 3, entry_flow_depth_list,
                           test=test, fix_params=fix_params)

        with nn.parameter_scope("middle_flow"):
            x = middle_flow(x, 1, middle_flow_depth_list,
                            test=test, fix_params=fix_params)

        with nn.parameter_scope("exit_flow"):
            x = exit_flow(x, 2, exit_flow_depth_list,
                          test=test, fix_params=fix_params)
            x = F.relu(x)
        global endpoints
        endpoints['Decoder End Point 2'] = x

    return endpoints


# shape-reconstruction/implicit-differentiable-renderer/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.

import nnabla as nn
import numpy as np
from nnabla.utils.data_iterator import data_iterator
from nnabla.utils.data_source import DataSource
from nnabla.utils import image_utils

import os
import glob
import imageio

from helper import load_K_Rt_from_P, generate_raydir_camloc


class DTUMVSDataSource(DataSource):
    '''
    Load DTUMVS dataset from the zip file created by the author of "Multiview Neural Surface Reconstruction by Disentangling Geometry and Appearance".
    '''

    def _get_data(self, position):
        img_idx = self._img_indices[position]
        image = self._images[img_idx]
        mask = self._masks[img_idx]
        intrinsic = self._intrinsics[img_idx]
        pose = self._poses[img_idx]

        H, W, _ = image.shape
        R = H * W

        color = image.reshape((R, 3))[self.pixel_idx]
        mask = mask.reshape((R, 1))[self.pixel_idx]
        xy = self._xy[self.pixel_idx]
        return color, mask, intrinsic, pose, xy

    def _load_dtumvs(self, path):
        # Images
        image_files = sorted(glob.glob(os.path.join(path, "image", "*")))
        images = np.asarray([image_utils.imread(f) for f in image_files])
        images = images * (1.0 / 127.5) - 1.0

        # Masks
        mask_files = sorted(glob.glob(os.path.join(path, "mask", "*")))
        masks = np.asarray([imageio.imread(f, as_gray=True)[:, :, np.newaxis] > 127.5
                            for f in mask_files]) * 1.0

        # Camera projection matrix and scale matrix for special correctness
        cameras = np.load(os.path.join(path, "cameras.npz"))
        world_mats = [cameras['world_mat_%d' % idx].astype(
            np.float32) for idx in range(len(images))]
        scale_mats = [cameras['scale_mat_%d' % idx].astype(
            np.float32) for idx in range(len(images))]

        intrinsics, poses = [], []
        for W, S in zip(world_mats, scale_mats):
            P = W @ S
            P = P[:3, :4]
            intrinsic, pose = load_K_Rt_from_P(P)
            intrinsics.append(intrinsic[:3, :3])
            poses.append(pose)

        # return images[0:1, ...], masks[0:1, ...], np.asarray(intrinsics)[0:1, ...], np.asarray(poses)[0:1, ...]
        return images, masks, np.asarray(intrinsics), np.asarray(poses)

    def __init__(self, path, n_rays, train=True, shuffle=False, rng=None):
        super(DTUMVSDataSource, self).__init__(shuffle=shuffle)
        self._n_rays = n_rays
        self._train = train

        self._images, self._masks, self._intrinsics, self._poses = self._load_dtumvs(
            path)

        # assume all images have same resolution
        H, W, _ = self._images[0].shape
        x = np.arange(W)
        y = np.arange(H)
        xx, yy = np.meshgrid(x, y)
        self._xy = np.asarray([xx.flatten(), yy.flatten()]).T

        self._size = len(self._images)
        self._pixels = H * W
        self._variables = ('image', 'mask', 'intrinsic', 'pose', 'xy')
        if rng is None:
            rng = np.random.RandomState(313)
        self.rng = rng
        self.reset()

        self.pixel_idx = self._sampling_idx()

        dname = os.path.split(path.rstrip("/"))[-1]
        nn.logger.info(f"--- Finish loading DTU MVS dataset ({dname}). ---")
        nn.logger.info(f"Num. of images = {self._size}")
        nn.logger.info(f"Num. of pixels (H x W) = {self._pixels} ({H} x {W})")
        nn.logger.info(f"Num. of random rays = {self._n_rays}")

    def reset(self):
        if self._shuffle:
            self._img_indices = self.rng.permutation(self._size)
        else:
            self._img_indices = np.arange(self._size)
        super(DTUMVSDataSource, self).reset()

    @property
    def images(self):
        """Get copy of whole data with a shape of (B, H, W, 3)."""
        return self._images.copy()

    @property
    def poses(self):
        return self._poses.copy()

    @property
    def intrinsics(self):
        return self._intrinsics.copy()

    @property
    def masks(self):
        return self._masks.copy()

    def change_sampling_idx(self):
        self.pixel_idx = self._sampling_idx()

    def _sampling_idx(self):
        return self.rng.randint(0, self._pixels, self._n_rays)


def data_iterator_dtumvs(data_source,
                         batch_size,
                         rng=None,
                         with_memory_cache=False,
                         with_file_cache=False):
    '''
    Provide DataIterator with :py:class:`DTUMVSDataSource`
    with_memory_cache and with_file_cache option's default value is all False,
    because :py:class:`DTUMVSDataSource` is able to store all data into memory.
    '''
    return data_iterator(data_source,
                         batch_size,
                         rng,
                         with_memory_cache,
                         with_file_cache)


def main(args):
    # Data Iterator
    ds = DTUMVSDataSource(args.path, args.n_rays, shuffle=True)
    di = data_iterator_dtumvs(ds, 1)
    for i in range(args.iters):
        pcolor, mask, intrinsic, pose, xy = di.next()
        print(f"pcolor.shape = {pcolor.shape}")
        print(f"mask.shape = {mask.shape}")
        print(f"intrinsic.shape = {intrinsic.shape}")
        print(f"pose.shape = {pose.shape}")
        print(f"xy.shape = {xy.shape}")
        print(f"Pcolor (min, max) = ({pcolor.min()}, {pcolor.max()})")
        print(f"Mask (min, max) = ({mask.min()}, {mask.max()})")

    # Generate rays
    raydir, camloc = generate_raydir_camloc(pose, intrinsic, xy)
    print(f"raydir.shape = {raydir.shape}")
    np.testing.assert_allclose(
        np.sum(raydir ** 2, axis=-1) ** 0.5, 1.0, atol=1e-6)


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(description="DTU MVS Dataset.")
    parser.add_argument('--path', type=str, required=True,
                        help="Path to scale directory, Ex, DTU/scan24.")
    parser.add_argument('--batch-size', '-b', type=int, default=4)
    parser.add_argument('--n-rays', '-n', type=int, default=800)
    parser.add_argument('--iters', '-i', type=int, default=100)

    args = parser.parse_args()
    main(args)


# shape-reconstruction/implicit-differentiable-renderer/evaluate.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 numpy as np
from nnabla.ext_utils import get_extension_context
from nnabla.monitor import Monitor, MonitorSeries

from dataset import DTUMVSDataSource
from synthesize import render


def psnr(image_list, image_gt_list, mask_obj_list, monitor_psnrs):
    squared_error_sum = 0
    mask_sum = 0
    for i, elms in enumerate(zip(image_list, image_gt_list, mask_obj_list)):
        image, image_gt, mask_obj = elms
        image_gt = image_gt.transpose((2, 0, 1))
        mask_obj = mask_obj.transpose((2, 0, 1))

        image = (image + 1) * 127.5
        image_gt = (image_gt + 1) * 127.5
        squared_error_sum_i = np.sum(((image - image_gt) * mask_obj) ** 2)
        mask_sum_i = np.sum(mask_obj)

        squared_error_sum += squared_error_sum_i
        mask_sum += mask_sum_i

        mse_i = squared_error_sum_i / mask_sum_i
        psnr_i = 20 * np.log10(255) - 10 * np.log10(mse_i)
        monitor_psnrs.add(i, psnr_i)

    mse = squared_error_sum / mask_sum
    psnr = 20 * np.log10(255) - 10 * np.log10(mse)
    return psnr


def main(args):
    # Setting
    device_id = args.device_id
    conf = args.conf
    path = conf.data_path
    B = conf.batch_size
    R = conf.n_rays
    L = conf.layers
    D = conf.depth
    feature_size = conf.feature_size

    ctx = get_extension_context('cudnn', device_id=device_id)
    nn.set_default_context(ctx)

    # Dataset
    ds = DTUMVSDataSource(path, R, shuffle=True)

    # Monitor
    monitor_path = "/".join(args.model_load_path.split("/")[0:-1])
    monitor = Monitor(monitor_path)
    monitor_psnrs = MonitorSeries(f"PSNRs", monitor, interval=1)
    monitor_psnr = MonitorSeries(f"PSNR", monitor, interval=1)

    # Load model
    nn.load_parameters(args.model_load_path)

    # Evaluate
    image_list = []
    for pose, intrinsic, mask_obj in zip(ds.poses, ds.intrinsics, ds.masks):
        image = render(pose[np.newaxis, ...],
                       intrinsic[np.newaxis, ...],
                       mask_obj[np.newaxis, ...],
                       conf)
        image_list.append(image)

    metric = psnr(image_list, ds.images, ds.masks, monitor_psnrs)
    monitor_psnr.add(0, metric)


if __name__ == '__main__':
    import argparse
    from ruamel.yaml import YAML
    from collections import namedtuple

    parser = argparse.ArgumentParser(
        description="Implicit Differentiable Renderer Training.")
    parser.add_argument('--device-id', '-d', type=int, default="0")
    parser.add_argument('--model-load-path', type=str, required=True)
    parser.add_argument('--config', type=str,
                        default="conf/default.yaml", required=True)

    args = parser.parse_args()
    with open(args.config, "r") as f:
        conf = YAML(typ='safe').load(f)
        conf = namedtuple("Conf", conf)(**conf)
    args.conf = conf

    main(args)


# shape-reconstruction/implicit-differentiable-renderer/extract.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 numpy as np
from nnabla.ext_utils import get_extension_context

import os
from functools import partial
from tqdm import tqdm

import trimesh
from skimage import measure

from network import sdf_net


def create_mesh_from_volume(volume, spacing, gradient_direction="ascent"):
    spacing = [np.max(spacing)] * 3
    verts, faces, normals, values = measure.marching_cubes_lewiner(volume,
                                                                   0.0,
                                                                   spacing=spacing,
                                                                   gradient_direction=gradient_direction)
    mesh = trimesh.Trimesh(verts, faces, normals)
    return mesh


def compute_pts_vol(model, mins, maxs, step, sub_batch_size, bias=None, V=None):

    x = np.arange(mins[0], maxs[0], step).astype(np.float32)
    y = np.arange(mins[1], maxs[1], step).astype(np.float32)
    z = np.arange(mins[2], maxs[2], step).astype(np.float32)
    X, Y, Z = np.meshgrid(x, y, z)
    grid_shape = X.shape
    print(f"Grid shape for (x, y, z) = {grid_shape}")

    X = X.reshape(-1)
    Y = Y.reshape(-1)
    Z = Z.reshape(-1)
    pts = np.stack((X, Y, Z), axis=1)

    vol = []
    for b in tqdm(range(0, pts.shape[0], sub_batch_size),
                  desc="compute-volume"):
        p = pts[b:b+sub_batch_size, :]
        p = (p - bias) @ V.T if V is not None else p
        #p = p @ V.T if V is not None else p
        v = model(nn.NdArray.from_numpy_array(p))
        v = v.data.copy().reshape(-1)
        vol.append(v)
    pts = pts.reshape((-1, 3))

    s0, s1, s2 = grid_shape
    vol = np.concatenate(vol).reshape((s0, s1, s2)).transpose((0, 1, 2))
    return pts, vol


def create_largest_mesh(mesh):
    meshes = mesh.split(only_watertight=False)
    areas = np.array([m.area for m in meshes], dtype=np.float)
    mesh = meshes[areas.argmax()]
    return mesh


def main(args):
    conf = args.conf

    # Context
    ctx = get_extension_context("cudnn", device_id=args.device_id)
    nn.set_default_context(ctx)
    nn.set_auto_forward(True)

    # Network
    nn.load_parameters(args.model_load_path)

    # Extract rough mesh
    mins = np.asarray([-conf.bounding_box_size] * 3)
    maxs = np.asarray([conf.bounding_box_size] * 3)
    step = 2 * conf.bounding_box_size / 100
    pts, vol = compute_pts_vol(partial(sdf_net, conf=conf),
                               mins, maxs, step,
                               conf.sub_batch_size)
    mesh_low = create_mesh_from_volume(
        vol, [step] * 3, conf.gradient_direction)
    mesh_large = create_largest_mesh(mesh_low)

    # Samples
    points_all = trimesh.sample.sample_surface(mesh_low, 10000)[0]
    points_large = trimesh.sample.sample_surface(mesh_large, 10000)[0]

    # Centering and rotate
    mean = np.mean(points_large, axis=0, keepdims=True)
    cov = (points_large - mean).transpose((1, 0)) @ (points_large - mean)
    _, V = np.linalg.eig(cov)
    points_all = (points_all - mean) @ V
    points_large = (points_large - mean) @ V

    # Create scale and bias
    mins_all = np.asarray([min(m0, m1) for m0, m1 in zip(
        points_all.min(axis=0), points_large.min(axis=0))])
    maxs_all = np.asarray([max(m0, m1) for m0, m1 in zip(
        points_all.max(axis=0), points_large.max(axis=0))])
    idx_max_length = np.argmax(maxs_all - mins_all)
    m = mins_all[idx_max_length]
    M = maxs_all[idx_max_length]
    l0 = -conf.bounding_box_size
    l1 = conf.bounding_box_size
    scale = (l1 - l0) / (M - m)
    bias = l0 - scale * m

    # Extract fine mesh
    mins_large = points_large.min(axis=0)
    maxs_large = points_large.max(axis=0)
    min_length = np.min(maxs_large - mins_large)
    grid_step = min_length * scale / conf.grid_size
    print(f"grid_step = {grid_step}")
    mins = scale * mins_large + bias
    maxs = scale * maxs_large + bias
    pts, vol = compute_pts_vol(partial(sdf_net, conf=conf),
                               mins, maxs, grid_step,
                               conf.sub_batch_size,
                               bias, V)
    mesh = create_mesh_from_volume(
        vol, [grid_step] * 3, conf.gradient_direction)
    mesh = create_largest_mesh(mesh)

    # Save
    dirname, pname = args.model_load_path.split("/")
    fname = os.path.splitext(pname)[0]
    fpath = f"{dirname}/{fname}_mesh.ply"
    nn.logger.info(f"Saving the mesh file to {fpath}")
    mesh.export(fpath, "ply")


if __name__ == '__main__':
    import argparse
    from ruamel.yaml import YAML
    from collections import namedtuple

    parser = argparse.ArgumentParser(
        description="Implicit Differentiable Renderer Training.")
    parser.add_argument('--device-id', '-d', type=int, default="0")
    parser.add_argument('--model-load-path', type=str, required=True)
    parser.add_argument('--config', type=str,
                        default="conf/default.yaml", required=True)

    args = parser.parse_args()
    with open(args.config, "r") as f:
        conf = YAML(typ='safe').load(f)
        conf = namedtuple("Conf", conf)(**conf)
    args.conf = conf

    main(args)


# shape-reconstruction/implicit-differentiable-renderer/helper.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
from enum import Enum
from dataclasses import dataclass
from typing import Any

import nnabla as nn
import nnabla.functions as F

import cv2

inch_to_mm = 25.4


class FitResolutionGate(Enum):
    FILL = 0
    OVER_SCREEN = 1


@dataclass
class Camera:
    """
    Camera is assumed to be oriended to negative z-axis as default.
    """

    # Intrinsic Parameter
    focal_length: float = 20              # in mm
    film_aperture_width: float = 0.980    # 35mm Full Aperture in inches
    film_aperture_height: float = 0.735   # 35mm Full Aperture in inches
    z_near: float = 1
    z_far: float = 1000
    image_width: int = 640
    image_height: int = 480
    k_fit_resolution_gate: FitResolutionGate = FitResolutionGate.OVER_SCREEN
    # Extrinsic Parameter
    camera_to_world: Any = None
    # Other
    z_orientation: int = -1  # {+1 or -1}

    def field_of_view(self, film_aperture):
        return 2 * 180 / np.pi * np.arctan((film_aperture * inch_to_mm / 2) / self.focal_length)

    def compute_screen(self):
        # fovs, top, right
        fov_h = self.field_of_view(self.film_aperture_width)
        fov_v = self.field_of_view(self.film_aperture_height)
        top = ((self.film_aperture_height * inch_to_mm / 2) /
               self.focal_length) * self.z_near
        right = ((self.film_aperture_width * inch_to_mm / 2) /
                 self.focal_length) * self.z_near

        # Scale if necessary
        x_scale = 1.0
        y_scale = 1.0
        film_aspect_ratio = self.compute_film_aspect_ratio()
        image_aspect_ratio = self.compute_image_aspect_ratio()
        if self.k_fit_resolution_gate == FitResolutionGate.FILL:
            if film_aspect_ratio > image_aspect_ratio:
                x_scale *= image_aspect_ratio / film_aspect_ratio
            else:
                y_scale *= film_aspect_ratio / image_aspect_ratio
        elif self.k_fit_resolution_gate == FitResolutionGate.OVER_SCREEN:
            if film_aspect_ratio > image_aspect_ratio:
                y_scale *= film_aspect_ratio / image_aspect_ratio
            else:
                x_scale *= image_aspect_ratio / film_aspect_ratio
        else:
            raise ValueError(
                f"k_fit_resolution_gate(={k_fit_resolution_gate}) must be in {list(FitResolutionGate)}.")
        top *= y_scale
        right *= x_scale
        bottom = -top
        left = -right

        # Screen (Image Plane)
        screen = Screen(left=left, right=right, top=top, bottom=bottom)
        return screen

    def compute_film_aspect_ratio(self):
        return 1.0 * self.film_aperture_width / self.film_aperture_height

    def compute_image_aspect_ratio(self):
        return 1.0 * self.image_width / self.image_height

    def compute_projection_matrix(self, screen):
        """Projection matrix (OpneGC style).
        """
        l = screen.left
        r = screen.right
        t = screen.top
        b = screen.bottom
        n = self.z_near
        f = self.z_far

        row0 = [2 * n / (r - l), 0, (r + l) / (r - l), 0]
        row1 = [0, 2 * n / (t - b), (t + b) / (t - b), 0]
        row2 = [0, 0, - (f + n) / (f - n), - 2 * f * n / (f - n)]
        row3 = [0, 0, self.z_orientation, 0]
        P = np.array([row0,
                      row1,
                      row2,
                      row3])
        return P

    def compute_intrinsic_inv(self, fov=None):
        """Inverse camera intrintic matrix.

        Transofmration formula is as follows: 

        Px = (2 * ((x + 0.5) / imageWidth) - 1) * tan(fov / 2 * M_PI / 180) * imageAspectRatio; 
        Py = (1 - 2 * ((y + 0.5) / imageHeight) * tan(fov / 2 * M_PI / 180)

        x: pixel x
        y: pixel y
        """
        W = self.image_width
        H = self.image_height
        fov = self.field_of_view(
            self.film_aperture_height) if fov is None else fov * np.pi / 180
        tan_fov_2 = np.tan(fov / 2)
        A = self.compute_image_aspect_ratio()

        row0 = [2 / W * tan_fov_2 * A, 0, (1 - W) / W * tan_fov_2 * A]
        row1 = [0, -2 / H * tan_fov_2, (H - 1) / H * tan_fov_2]
        row2 = [0, 0, self.z_orientation]
        K = np.array([row0,
                      row1,
                      row2])
        return K

    def compute_intrinsic(self, fov=None):
        return np.linalg.inv(self.compute_intrinsic_inv(fov))


@dataclass
class Screen:
    left: float = 0
    right: float = 0
    top: float = 0
    bottom: float = 0


def normalize(vertex):
    assert len(vertex) == 3
    vertex = vertex / np.sqrt(np.sum(vertex ** 2))
    return np.array([vertex[0], vertex[1], vertex[2]])


def look_at(F, T, tmp=np.array([0.0, 1.0, 0.0]), z_orientation=-1):
    """Return a camera-to-world transformation matrix.

    Camera is assumed to be oriended to negative z-axis.
    """
    if len(F) == 4 or len(T) == 4:
        F = F[:3]
        T = T[:3]
    forward = normalize(F - T) if z_orientation == -1 else normalize(T - F)
    right = np.cross(normalize(tmp), forward)
    up = np.cross(forward, right)
    M = np.stack([right, up, forward, F])
    H = np.array([0.0, 0.0, 0.0, 1.0])[:, np.newaxis]
    M = np.concatenate([M, H], axis=1)
    return M.T


class Light:
    """Base Light
    """

    def __init__(self,
                 color=nn.Variable.from_numpy_array([1., 1., 1.])):
        """
        Args:
          color (np.ndarray, nn.NdArray, nn.Variable): Light color.
        """
        if isinstance(color, np.ndarray):
            self.color = nn.Variable.from_numpy_array(color)
        elif isinstance(color, nn.NdArray):
            self.color = nn.Variable(color.shape)
            self.color.data = color
        elif isinstance(color, nn.Variable):
            self.color = color


class DistantLight(Light):
    """Distant Light
    """

    def __init__(self,
                 color=nn.Variable.from_numpy_array([1., 1., 1.]),
                 direction=nn.Variable.from_numpy_array([1., 1., 1.]),
                 ):
        """
        Args:
          color (np.ndarray, nn.NdArray, nn.Variable): Light color.
          direction (np.ndarray, nn.NdArray, nn.Variable): Direction of the light.
        """
        super(DistantLight, self).__init__(color)

        import nnabla.functions as F
        if isinstance(direction, np.ndarray):
            direction = nn.Variable.from_numpy_array(direction)
            self.direction = F.norm_normalization(direction, eps=1e-6)
        elif isinstance(direction, nn.NdArray):
            direction = nn.Variable(direction.shape)
            direction.data = direction
            self.direction = F.norm_normalization(direction, eps=1e-6)
        elif isinstance(direction, nn.Variable):
            self.direction = F.norm_normalization(direction, eps=1e-6)


class PointLight(Light):
    """Point Light
    """

    def __init__(self,
                 color=nn.Variable.from_numpy_array([1., 1., 1.]),
                 point=nn.Variable.from_numpy_array([0., 0., 0.]),
                 ):
        """
        Args:
          color (np.ndarray, nn.NdArray, nn.Variable): Light color.
          point (np.ndarray, nn.NdArray, nn.Variable): Position of the light.
        """
        super(PointLight, self).__init__(color)

        if isinstance(point, np.ndarray):
            self.point = nn.Variable.from_numpy_array(point)
        elif isinstance(point, nn.NdArray):
            self.point = nn.Variable(point.shape)
            self.point.data = point
        elif isinstance(point, nn.Variable):
            self.point = point


def lambert(normal, light_dir):
    """Lambertian (diffuse) lighting.

    Args:
      normal (nn.Variable): Normal of (B, R, 3) shape.
      light_dir (nn.Variable): Light direction of (3, ) shape.

    """
    base_axis = len(normal.shape) - 1
    cos = F.relu(F.affine(normal, light_dir.reshape(
        [3, 1]), base_axis=base_axis))
    return cos


def facing_ratio(normal, raydir):
    """Facing ratio lighting.

    Args:
      normal (nn.Variable): Normal of (B, R, 3) shape.
      raydir (nn.Variable): Ray direction of (B, R, 3) shape.

    """
    factor = F.relu(F.sum(normal * (-raydir), axis=1, keepdims=True))
    return factor


def load_K_Rt_from_P(P):
    out = cv2.decomposeProjectionMatrix(P)
    K = out[0]  # intrinsic matrix
    R = out[1]  # world-to-camera matrix
    c = out[2]  # camera location

    K = K / K[2, 2]
    intrinsic = np.eye(4)
    intrinsic[:3, :3] = K

    pose = np.eye(4, dtype=np.float32)  # [camera-to-world | camera location]
    pose[:3, :3] = R.transpose()
    pose[:3, 3] = (c[:3] / c[3])[:, 0]

    return intrinsic, pose


def generate_raydir_camloc(pose, intrinsic, xy):
    """
    pose: [B, 4, 4]
    intrinsic: [B, 3, 3]
    xy: [B, R, 2]

    x_p = K [R | t] x_w
    => x_p = K (R x_w + t)
    => x_w = R_inv K_inv x_p - R_inv t
    => x_w = normalize(x_w)
    """
    B, R, _ = xy.shape

    # Align dimensions
    R_c2w = pose[:, np.newaxis, :3, :3]
    camloc = pose[:, np.newaxis, :3, 3:4]
    K_inv = np.linalg.inv(intrinsic[:, np.newaxis, :, :])

    # Transform pixel --> camera --> world
    z = np.ones([B, R, 1])
    xyz_pixel = np.concatenate([xy, z], axis=-1)[:, :, :, np.newaxis]
    xyz_camera = np.matmul(K_inv, xyz_pixel)

    # Note: we do not need to subtract camloc since R_c2w does not contain it.
    xyz_world = np.matmul(R_c2w, xyz_camera)

    # Normalize
    xyz_world = xyz_world.reshape((B, R, 3))
    raydir = xyz_world / \
        np.sqrt(np.sum(xyz_world ** 2, axis=-1, keepdims=True))

    return raydir, camloc.reshape((B, 3))


def generate_all_pixels(W, H):
    x = np.arange(0, W)
    y = np.arange(0, H)
    xx, yy = np.meshgrid(x, y)
    xy = np.asarray([xx.flatten(), yy.flatten()]).T
    return xy


def create_monitor_path(data_path, monitor_path=""):
    base_path = data_path.rstrip("/").split("/")[-1]
    monitor_path = f"{monitor_path}_{base_path}"
    return monitor_path


if __name__ == '__main__':
    rng = np.random.RandomState(412)
    pose = rng.randn(1, 3, 4).astype(np.float32)
    intrinsic = np.asarray([[2.0, 0, 0.5], [0, 3.0, 0.5], [0, 0, 1]])[
                           np.newaxis].astype(np.float32)
    x = np.arange(0, 3)
    y = np.arange(0, 2)
    xx, yy = np.meshgrid(x, y)
    xy = (np.asarray([xx.flatten(), yy.flatten()]).T)[np.newaxis, ...]

    ray_dirs, cam_loc = generate_raydir_camloc(pose, intrinsic, xy)

    print(ray_dirs)
    print(cam_loc)


# shape-reconstruction/implicit-differentiable-renderer/network.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
import nnabla.initializer as I
import numpy as np
from nnabla.ext_utils import get_extension_context
from nnabla.random import prng
from nnabla.parametric_functions import parametric_function_api
from nnabla.parameter import get_parameter_or_create

import matplotlib.pyplot as plt
from functools import partial

from ray_tracer import ray_trace


class WeightNormalizationInitializer(I.BaseInitializer):

    def __init__(self, w, dim=0, eps=1e-12):
        self.w = w
        self.dim = dim
        self.eps = eps

    def __repr__(self):
        return '{}({})'.format(self.__class__.__name__)

    def __call__(self, shape):
        axis = tuple([a for a in range(len(self.w.shape)) if a != self.dim])
        w_norm_data = np.sqrt(np.sum(self.w.d ** 2, axis=axis) + self.eps)
        return np.float32(w_norm_data)


@parametric_function_api("wn", [
    ('g', 'Weight Normalization adaptive scale scalar.', 'w.shape[dim]', True),
])
def weight_normalization(w, dim=1, eps=1e-12, fix_parameters=False):
    """
    """
    outmaps = w.shape[dim]
    initializer = WeightNormalizationInitializer(w, dim, eps)
    g = get_parameter_or_create("g", (outmaps, ),
                                initializer=initializer, need_grad=True,
                                as_need_grad=not fix_parameters)
    w_wn = F.weight_normalization(w, g, dim, eps=eps)
    return w_wn


def weight_normalization_backward(inputs, dim=1, eps=1e-12):
    # We do not need the gradients wrt weight and scale by nn.grad,
    # since nn.grad is performed wrt the intersection point.
    ## dw_wn = inputs[0]
    ## w0 = inputs[1]
    ## g0 = inputs[2]
    return None, None


nn.backward_functions.register(
    "WeightNormalization", weight_normalization_backward)


def affine(h, D, w_init=None, b_init=None, with_bias=True, name=None):
    with nn.parameter_scope(name):
        apply_w = partial(weight_normalization, dim=1)
        h = PF.affine(h, D, base_axis=h.ndim - 1, apply_w=apply_w,
                      w_init=w_init, b_init=b_init, with_bias=with_bias)
    return h


def positional_encoding(x, N=6, include_input=True):
    """
    Args:
      x: Input (B, R, 3)
      N: Number of bands, N=6 for implicit network and N=4 for rendering network.
    """

    gamma = [x] if include_input else []
    bands = 2 ** np.arange(0, N + 1)
    data_holder = nn.Variable if isinstance(x, nn.Variable) else nn.NdArray
    bands = data_holder.from_numpy_array(bands)
    bands = F.reshape(bands, tuple([1] * x.ndim) + (N + 1, )) \
        * F.reshape(x, x.shape + (1, ))
    bands = F.reshape(bands, bands.shape[:-2] + (-1, ))
    cos_x = F.cos(bands)
    sin_x = F.sin(bands)

    gamma += [cos_x, sin_x]
    gamma = F.concatenate(*gamma, axis=-1)

    return gamma


class GeometricInitializer(I.BaseInitializer):
    """
    """

    def __init__(self, Di, Do, sigma, zero_start=None):
        self.Di = Di
        self.Do = Do
        self.sigma = sigma
        self.zero_start = zero_start

    def __repr__(self):
        return '{}({})'.format(self.__class__.__name__)

    def __call__(self, shape):
        w_init = np.sqrt(self.sigma) * prng.randn(*[self.Di, self.Do])
        if self.zero_start is not None:
            w_init[self.zero_start:, :] = 0.0
        return w_init


@nn.parameter_scope("implicit-network")
def implicit_network(x, D=512, feature_size=256, L=9, skip_in=[4], N=6,
                     act="softplus",
                     including_input=True,
                     initial_sphere_radius=0.75):
    """Implicit Network.

    Args:
      x: Position on a ray.
      D: Dimension of a network.
      feature_size: Feature dimension of the final output.
      L: Number of layers
      skip_in: Where the skip connection appears.
      N: Number of frequency of the positional encoding.
      act: Activation function.
      inclugin_input: Include input to the positional encoding (PE).
      initial_sphere_radius: the radius of the initial network sphere.

    Network architecture looks like:

    x --> [PE(x)] --> affine --> relu --> ... -->
    concate([h, x]) --> affine --> relu --> ... -->
    affine(h) --> [sdf, feature]
    """

    act_map = dict(relu=F.relu, softplus=partial(F.softplus, beta=100))
    Dx = x.shape[-1]
    act = act_map[act]

    h = positional_encoding(x, N, including_input)
    for l in range(L):
        # First
        if l == 0:
            Dh = h.shape[-1]
            Dx = x.shape[-1]
            w_init = GeometricInitializer(Dh, D, 2 / D, Dx)
            h = affine(h, D, w_init=w_init, name=f"affine-{l:02d}")
            h = act(h)
        # Skip
        elif l in skip_in:
            w_init = GeometricInitializer(D, D, 2 / (D - Dx), -Dx)
            h = affine(h, D, w_init=w_init, name=f"affine-{l:02d}")
            h = act(h)
        # Last (scalar + feature_size)
        elif l == L - 1:
            Do = 1 + feature_size
            w_init = np.sqrt(np.pi / D) * np.ones([D, Do])
            h = affine(h, Do, w_init=w_init, b_init=I.ConstantInitializer(-initial_sphere_radius),
                       name=f"affine-last")
        # Intermediate
        else:
            Do = D - Dx if l + 1 in skip_in else D
            w_init = GeometricInitializer(D, Do, 2 / Do)
            h = affine(h, Do, w_init=w_init, name=f"affine-{l:02d}")
            h = act(h)
            h = F.concatenate(*[h, x]) if l + 1 in skip_in else h
            # h = F.concatenate(*[h, x]) / np.sqrt(2) if l + 1 in skip_in else h # (the paper used this scale)
    return h


def sdf_feature_grad(implicit_network, x, conf):
    y = implicit_network(x, initial_sphere_radius=conf.initial_sphere_radius)
    sdf = y[..., 0:1]
    feature = y[..., 1:]
    grad = nn.grad([sdf], [x])[0]
    return sdf, feature, grad


def sample_network(x_curr, sdf_cur, raydir, grad_curr):
    """
    x_curr: Points (B, R, 3) either on surface or not
    sdf_cur: SDF on x_curr (B, R, 1)
    raydir: Ray direction (B, R, 3)
    grad_curr: Gradients on x_curr (B, R, 3)
    """

    # Denominator
    de = F.batch_matmul(grad_curr[..., np.newaxis, :],
                        raydir[..., np.newaxis, :], transpose_b=True)
    de = de.reshape(sdf_cur.shape)
    de_inv = (1.0 / de).apply(need_grad=False)
    de_inv = F.minimum_scalar(de_inv, 1e30).apply(
        need_grad=False)  # (numerical issue de = cos(x, y) = 0)
    # Differentiable intersection point (discrete update of implicit differentiation)
    sdf_cur0 = sdf_cur.get_unlinked_variable(need_grad=False)
    x_hat = x_curr - (sdf_cur - sdf_cur0) * de_inv * raydir
    return x_hat


@nn.parameter_scope("lighting-network")
def lighting_network(x_hat, normal, feature, view, D=512, L=4, N=4, including_input=True):
    """
    Args
      x_hat: Differentiable intersection point (B, R, 3)
      normal: Normal on x_hat (B, R, 3) (should be normalized before).
      feature: Intermediate output of the implicit network (B, R, feature_size).
      view: View direction (B, R, 3)
      D: Dimension of a network.
      L: Number of layers.
      N: Number of frequency of the positional encoding.
      inclugin_input: Include input to the positional encoding (PE).
    """
    pe_view = positional_encoding(view, N, including_input)
    h = F.concatenate(*[x_hat, normal, feature, pe_view], axis=-1)
    for l in range(L - 1):
        h = affine(h, D, name=f"affine-{l:02d}")
        h = F.relu(h)
    h = affine(h, 3, name=f"affine-{L - 1:02d}")
    h = F.tanh(h)
    return h


def sdf_net(x, conf):
    out = implicit_network(x, D=conf.depth, feature_size=conf.feature_size,
                           initial_sphere_radius=conf.initial_sphere_radius)
    sdf = out[..., 0:1]
    return sdf


def norm(x, axis, eps=1e-24):
    return F.sum(x ** 2 + eps, axis, keepdims=True) ** 0.5


def idr_loss(camloc, raydir, alpha, color_gt, mask_obj, conf):
    # Setting
    B, R, _ = raydir.shape
    L = conf.layers
    D = conf.depth
    feature_size = conf.feature_size

    # Ray trace (visibility)
    x_hit, mask_hit, dists, mask_pin, mask_pout = \
        ray_trace(partial(sdf_net, conf=conf),
                  camloc, raydir, mask_obj, t_near=conf.t_near, t_far=conf.t_far,
                  sphere_trace_itr=conf.sphere_trace_itr,
                  ray_march_points=conf.ray_march_points,
                  n_chunks=conf.n_chunks,
                  max_post_itr=conf.max_post_itr,
                  post_method=conf.post_method, eps=conf.eps)

    x_hit = x_hit.apply(need_grad=False)
    mask_hit = mask_hit.apply(need_grad=False, persistent=True)
    dists = dists.apply(need_grad=False)
    mask_pin = mask_pin.apply(need_grad=False)
    mask_pout = mask_pout.apply(need_grad=False)
    mask_us = mask_pin + mask_pout
    P = F.sum(mask_us)

    # Current points
    x_curr = (camloc.reshape((B, 1, 3)) + dists * raydir).apply(need_grad=True)

    # Eikonal loss
    bounding_box_size = conf.bounding_box_size
    x_free = F.rand(-bounding_box_size, bounding_box_size,
                    shape=(B, R // 2, 3))
    x_point = F.concatenate(*[x_curr, x_free], axis=1)
    sdf_xp, _, grad_xp = sdf_feature_grad(implicit_network, x_point, conf)
    gp = (F.norm(grad_xp, axis=[grad_xp.ndim - 1], keepdims=True) - 1.0) ** 2.0
    loss_eikonal = F.sum(gp[:, :R, :] * mask_us) + F.sum(gp[:, R:, :])
    loss_eikonal = loss_eikonal / (P + B * R // 2)
    loss_eikonal = loss_eikonal.apply(persistent=True)

    sdf_curr = sdf_xp[:, :R, :]
    grad_curr = grad_xp[:, :R, :]

    # Mask loss
    logit = - alpha.reshape([1 for _ in range(sdf_curr.ndim)]) * sdf_curr
    loss_mask = F.sigmoid_cross_entropy(logit, mask_obj)
    loss_mask = loss_mask * mask_pout
    loss_mask = F.sum(loss_mask) / P / alpha
    loss_mask = loss_mask.apply(persistent=True)

    # Lighting
    x_hat = sample_network(x_curr, sdf_curr, raydir, grad_curr)
    _, feature, grad = sdf_feature_grad(implicit_network, x_hat, conf)
    normal = grad
    color_pred = lighting_network(x_hat, normal, feature, -raydir, D)

    # Color loss
    loss_color = F.absolute_error(color_gt, color_pred)
    loss_color = loss_color * mask_pin
    loss_color = F.sum(loss_color) / P
    loss_color = loss_color.apply(persistent=True)

    # Total loss
    loss = loss_color + conf.mask_weight * \
        loss_mask + conf.eikonal_weight * loss_eikonal

    return loss, loss_color, loss_mask, loss_eikonal, mask_hit


def render(camloc, raydir, conf):
    # Setting
    B, R, _ = raydir.shape
    L = conf.layers
    D = conf.depth
    feature_size = conf.feature_size

    # Sphere trace
    x_hit, mask_hit, _, _, _ = \
        ray_trace(partial(sdf_net, conf=conf),
                  camloc, raydir, test=True,
                  t_near=conf.t_near, t_far=conf.t_far,
                  sphere_trace_itr=conf.sphere_trace_itr,
                  ray_march_points=conf.ray_march_points,
                  n_chunks=conf.n_chunks,
                  max_post_itr=conf.max_post_itr,
                  post_method=conf.post_method, eps=conf.eps)
    x_hit = x_hit.apply(need_grad=True)
    mask_hit = mask_hit.apply(need_grad=False)

    # Intersection points and current points
    sdf_hit, feature, grad_xhit = sdf_feature_grad(
        implicit_network, x_hit, conf)

    # Lighting
    normal = grad_xhit
    raydir = raydir.reshape((B, R, 3))
    color_pred = lighting_network(x_hit, normal, feature, -raydir, D)
    color_pred = color_pred * mask_hit
    color_pred = color_pred.reshape((B, R, 3))

    return color_pred


def main(args):
    ctx = get_extension_context('cudnn', device_id=args.device_id)
    nn.set_default_context(ctx)

    grid_size = 50
    L = args.layers
    initial_sphere_radius = args.initial_sphere_radius

    mae_list = []
    for d in args.depth_list:
        mae = compute(grid_size, d, L, initial_sphere_radius)
        print(f"MAE(D={d})={mae}")
        mae_list.append(mae)

    # Check the law of large numbers over larger depths
    if len(args.depth_list) > 1:
        plt.plot(args.depth_list, mae_list)
        plt.xlabel("Depth")
        plt.ylabel("MAE")
        plt.title(f"Mean Absolute Error with various depth and L={L:03d}")
        plt.savefig(f"MAEL_{L:03d}.png")


def compute(grid_size, D, L, initial_sphere_radius):
    nn.clear_parameters()

    # 2D case
    x = np.linspace(-1, 1, grid_size)
    y = np.linspace(-1, 1, grid_size)
    xx, yy = np.meshgrid(x, y)
    xy = np.asarray([xx.flatten(), yy.flatten()]).T

    x = nn.Variable.from_numpy_array(xy)
    y = implicit_network(x, D, feature_size=256, L=L,
                         initial_sphere_radius=initial_sphere_radius)
    y = y[:, 0]
    y.forward()

    # Plot
    plt.contourf(xx, yy, y.d.reshape(grid_size, grid_size))
    plt.colorbar()
    plt.contourf(xx, yy, y.d.reshape(grid_size, grid_size),
                 colors=('r',), levels=[-0.01, 0, 0.01])
    plt.title(f"Contour of SDF(x, y) = 0 (D={D:03d} and L={L:03d})")
    plt.savefig(f"sdf_contour_D{D:03d}_L{L:03d}.png")
    plt.clf()

    # error
    gt = np.sqrt(np.sum(xy ** 2, axis=1)) - 1
    mae = np.sum(np.abs(gt - y.d.flatten())) / len(gt)

    return mae


if __name__ == '__main__':
    import argparse

    nn.random.seed(414)

    description = "Implicit Network. This script is intended for checking" \
        "1) the geometric initialization by plotting contour." \
        "2) the law of large numbers."
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument('--device-id', '-d', type=int, default="0")
    parser.add_argument('--depth-list', '-Ds', nargs="+",
                        type=int, default=[512])
    parser.add_argument('--layers', '-L', type=int, default=9)
    parser.add_argument('--initial-sphere-radius',
                        '-R', type=float, default=0.75)

    args = parser.parse_args()
    main(args)


# shape-reconstruction/implicit-differentiable-renderer/ray_tracer.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 numpy as np
from nnabla.ext_utils import get_extension_context
from nnabla.function import PythonFunction

import cv2
import time
from functools import partial

from helper import Camera, look_at, DistantLight, lambert


def bisection(x0, x1, implicit_function, max_post_itr):

    for i in range(max_post_itr):
        xm = (x0 + x1) * 0.5
        fm = implicit_function(xm)
        mp = F.greater_equal_scalar(fm, 0)
        mn = 1 - mp
        x0 = mp * xm + mn * x0  # f(x0) > 0
        x1 = mn * xm + mp * x1  # f(x1) < 0
    return x0, x1


def secant(x0, x1, implicit_function, max_post_itr, eps=1e-16):
    f0 = implicit_function(x0)  # > 0
    f1 = implicit_function(x1)  # < 0

    for i in range(max_post_itr):
        nu = f0 * (x1 - x0)
        de = f1 - f0
        mask0 = F.greater_scalar(F.abs(de), eps)
        mask1 = 1 - mask0
        nu = mask0 * nu + mask1 * 0
        de = mask0 * de + mask1 * 1

        xm = x0 - nu / de
        fm = implicit_function(xm)

        mp = F.greater_equal_scalar(fm, 0)
        mn = 1 - mp
        x0 = mp * xm + mn * x0
        f0 = mp * fm + mn * f0
        x1 = mn * xm + mp * x1
        f1 = mn * fm + mp * f1
    return x0, x1


class RayTrace(PythonFunction):

    def __init__(self, ctx, sdf_network, test=False, t_near=0, t_far=10,
                 sphere_trace_itr=50, ray_march_points=100, n_chunks=10,
                 max_post_itr=10, post_method="secant", eps=5e-5):

        super(RayTrace, self).__init__(ctx)

        self.sdf = sdf_network
        self.test = test
        self.t_init = t_near
        self.t_near = t_near
        self.t_far = t_far
        self.sphere_trace_itr = sphere_trace_itr
        self.ray_march_points = ray_march_points
        self.n_chunks = n_chunks
        self.max_post_itr = max_post_itr
        self.eps = eps
        if post_method == "secant":
            self.post_method = partial(
                secant, implicit_function=sdf_network, max_post_itr=max_post_itr)
        elif post_method == "bisection":
            self.post_method = partial(
                bisection, implicit_function=sdf_network, max_post_itr=max_post_itr)
        else:
            self.post_method = lambda x0, x1: (x0, x1)

    @property
    def name(self):
        return "RayTrace"

    def min_outputs(self):
        return 5

    def setup_impl(self, inputs, outputs):
        camloc = inputs[0]
        raydir = inputs[1]
        x_hit = outputs[0]
        mask_hit = outputs[1]
        dists = outputs[2]
        mask_pin = outputs[3]
        mask_pout = outputs[4]

        B, R, _ = raydir.shape

        x_hit.reset_shape((B, R, 3), True)
        mask_hit.reset_shape((B, R, 1), True)
        dists.reset_shape((B, R, 1), True)
        mask_pin.reset_shape((B, R, 1), True)
        mask_pout.reset_shape((B, R, 1), True)

    def forward_impl(self, inputs, outputs):
        # This auto forward is needed for weight normalization auto forward
        with nn.auto_forward():
            self._forward_impl(inputs, outputs)

    def _forward_impl(self, inputs, outputs):
        # inputs[0]: camera location (B, 3)
        # inputs[1]: ray direction (B, R, 3)
        # inputs[2]: mask_obj (B, R, 1)
        # outputs[0]: intersection points (B, R, 3)
        # outputs[1]: mask of intersection
        # outputs[2]: distance after unit sphere intersection (B, R, 1)
        #   - 1. between x_hit and camloc for P_in (= hit and mask_obj)
        #   - 2. as t_argmin for P_out (= ~P_in)
        # outputs[3]: mask of P_in (B, R, 1)
        # outputs[3]: mask of P_out (B, R, 1)

        # All variables becomes [B, R, 3] --> [B * R, {1, 3}] at the beginning
        # sdf_network: [B * R, 3] --> [B * R, 1]

        # Settings
        test = self.test
        t_near = self.t_near
        t_far = self.t_far
        sdf = self.sdf
        sphere_trace_itr = self.sphere_trace_itr
        ray_march_points = self.ray_march_points
        n_chunks = self.n_chunks
        max_post_itr = self.max_post_itr
        eps = self.eps

        # Ndarray
        camloc = inputs[0].data
        raydir = inputs[1].data
        mask_obj = inputs[2].data
        B, R, _ = raydir.shape
        N = ray_march_points
        camloc = F.broadcast(camloc[:, np.newaxis, :], (B, R, 3))
        camloc = F.reshape(camloc, (B * R, 3))
        raydir = F.reshape(raydir, (B * R, 3))
        mask_obj = F.reshape(mask_obj, (B * R, 1)) if not test else mask_obj

        # Unit sphare intersection
        t_start, t_finish, mask_us = \
            self.unit_sphere_intersection(camloc, raydir)

        # Bidirectional sphere tracing
        x_hit_st0, t_f, t_b, mask_st = \
            self.bidirectional_sphere_trace(camloc, raydir,
                                            t_start, t_finish)

        mask_obj = 1 if test else mask_obj

        # Ray marching
        x_hit_rm0, x_hit_rm1, mask_rm = self.ray_march(
            camloc, raydir, t_f, t_b, N, n_chunks)
        x_hit_rm, _ = self.post_method(x_hit_rm0, x_hit_rm1)

        x_hit = mask_st * x_hit_st0 + (1 - mask_st) * mask_rm * x_hit_rm
        mask_hit = mask_st + (1 - mask_st) * mask_rm

        if test:
            x_hit = F.reshape(x_hit, (B, R, 3))
            mask_hit = F.reshape(mask_hit, (B, R, 1))
            outputs[0].data.copy_from(x_hit)
            outputs[1].data.copy_from(mask_hit)
            return

        # Mask pin/pout
        mask_pin = mask_us * mask_hit * mask_obj
        mask_pout = mask_us * (1 - mask_hit * mask_obj)

        # Dists
        t_argmin = self.ray_march(
            camloc, raydir, t_start, t_finish, N, n_chunks, True)
        dists = F.norm(camloc - x_hit, axis=x_hit.ndim - 1, keepdims=True)
        dists = mask_pin * dists + mask_pout * t_argmin

        # Outputs
        x_hit = F.reshape(x_hit, (B, R, 3))
        mask_hit = F.reshape(mask_hit, (B, R, 1))
        dists = F.reshape(dists, (B, R, 1))
        mask_pin = F.reshape(mask_pin, (B, R, 1))
        mask_pout = F.reshape(mask_pout, (B, R, 1))
        outputs[0].data.copy_from(x_hit)
        outputs[1].data.copy_from(mask_hit)
        outputs[2].data.copy_from(dists)
        outputs[3].data.copy_from(mask_pin)
        outputs[4].data.copy_from(mask_pout)

    def unit_sphere_intersection(self, camloc, raydir):
        BR, _ = raydir.shape
        a = 1.0  # raydir is already normalized
        b = 2.0 * F.batch_matmul(F.reshape(camloc, (BR, 1, 3)),
                                 F.reshape(raydir, (BR, 3, 1)))
        c = F.batch_matmul(F.reshape(camloc, (BR, 1, 3)),
                           F.reshape(camloc, (BR, 3, 1))) - 1.0
        D = b ** 2 - 4 * a * c
        mask = F.reshape(F.greater_scalar(D, 0.0), (BR, 1))

        b = F.reshape(b, (BR, 1))
        D = F.reshape(D, (BR, 1))

        D = mask * D
        D_sqrt = D ** 0.5
        t_start = - (b + D_sqrt) / (2 * a)
        t_finish = - (b - D_sqrt) / (2 * a)

        t_start = t_start * mask + self.t_near * (1 - mask)
        t_finish = t_finish * mask + self.t_far * (1 - mask)

        return t_start, t_finish, mask

    def bidirectional_sphere_trace(self, camloc, raydir, t_start, t_finish):
        t_f = F.identity(t_start)
        x_f = camloc + t_f * raydir
        s_f = self.sdf(x_f)
        mask_hit_eps_f = 0 * F.identity(t_f)

        t_b = F.identity(t_finish)
        x_b = camloc + t_b * raydir
        s_b = self.sdf(x_b)
        mask_hit_eps_b = 0 * F.identity(t_b)

        for i in range(self.sphere_trace_itr - 1):
            # Forward direction
            mask_hit_eps_f_i = F.less_equal_scalar(F.abs(s_f), self.eps)
            mask_hit_eps_f += (1 - mask_hit_eps_f) * mask_hit_eps_f_i
            t_f += (1 - mask_hit_eps_f) * s_f
            x_f = camloc + t_f * raydir

            s_f_prev = F.identity(s_f)
            s_f = self.sdf(x_f)
            mask_pos_f_prev = (1 - mask_hit_eps_f) * \
                F.greater_scalar(s_f_prev, 0)
            mask_neg_f = (1 - mask_hit_eps_f) * F.less_scalar(s_f, 0)
            mask_revert_f = mask_pos_f_prev * mask_neg_f
            t_f -= mask_revert_f * s_f_prev
            s_f = mask_revert_f * s_f_prev + (1 - mask_revert_f) * s_f

            # Backward direction
            mask_hit_eps_b_i = F.less_equal_scalar(F.abs(s_b), self.eps)
            mask_hit_eps_b += (1 - mask_hit_eps_b) * mask_hit_eps_b_i
            t_b -= (1 - mask_hit_eps_b) * s_b
            x_b = camloc + t_b * raydir

            s_b_prev = F.identity(s_b)
            s_b = self.sdf(x_b)
            mask_pos_b_prev = (1 - mask_hit_eps_b) * \
                F.greater_scalar(s_b_prev, 0)
            mask_neg_b = (1 - mask_hit_eps_b) * F.less_scalar(s_b, 0)
            mask_revert_b = mask_pos_b_prev * mask_neg_b
            t_b += mask_revert_b * s_b_prev
            s_b = mask_revert_b * s_b_prev + (1 - mask_revert_b) * s_b

            ## print("s_f neg", np.sum(s_f.data < 0))
            ## print("s_b neg", np.sum(s_b.data < 0))

        # Fine grained start/finish points
        t_f0 = t_f
        t_f1 = t_f + mask_revert_f * s_f_prev
        x_hit_st0 = camloc + t_f0 * raydir
        ## x0, x1 = self.post_method(x_hit_st0, camloc + t_f1 * raydir)
        ## t_f0 = F.norm((x0 - camloc), axis=(x0.ndim - 1), keepdims=True)
        ## t_f1 = F.norm((x1 - camloc), axis=(x1.ndim - 1), keepdims=True)
        mask_hit_f1b = mask_revert_f * F.less(t_f1, t_b)
        t_b = t_f1 * mask_hit_f1b + t_b * (1 - mask_hit_f1b)

        # Reverse the opposite case
        mask_fb = F.less(t_f, t_b)
        t_f = t_f * mask_fb + t_start * (1 - mask_fb)
        t_b = t_b * mask_fb + t_finish * (1 - mask_fb)

        return x_hit_st0, t_f, t_b, mask_hit_eps_f

    def ray_march(self, camloc, raydir, t0, t1, N, n_chunks, t_argmin=False):
        # Points computation
        BR, _ = t0.shape
        t0 = F.reshape(t0, (BR, 1, 1))
        t1 = F.reshape(t1, (BR, 1, 1))
        camloc = F.reshape(camloc, (BR, 1, 3))
        raydir = F.reshape(raydir, (BR, 1, 3))
        step = (t1 - t0) / (N-1)
        intervals = F.reshape(F.arange(0, N), (1, N, 1))
        ts = t0 + step * intervals
        points = camloc + ts * raydir
        points = F.reshape(points, (BR * N, 3))

        # SDF computation
        sdf_points = []
        batch = (BR * N) // n_chunks
        for r in range(0, BR * N, batch):
            sdf_points.append(self.sdf(points[r:r+batch, :]))
        sdf_points = F.reshape(F.concatenate(*sdf_points, axis=0), (BR, N, 1)) if n_chunks != 1 else \
            F.reshape(sdf_points[0], (BR, N, 1))

        # t_argmin computation
        if t_argmin:
            idx_min = F.min(sdf_points, axis=1, keepdims=True, only_index=True)
            t_argmin = F.reshape(
                F.gather(ts, idx_min, axis=1, batch_dims=1), (BR, 1))
            return t_argmin

        # Intersection check
        points = F.reshape(points, (BR, N, 3))
        sdf_pos = F.greater_equal_scalar(sdf_points[:, :-1, :], 0)
        sdf_neg = F.less_equal_scalar(sdf_points[:, 1:, :], 0)
        mask_hit = sdf_pos * sdf_neg

        decreasing_consts = F.reshape(F.arange(N, 1, -1), (1, N - 1, 1))
        vals = mask_hit * decreasing_consts
        idx_max = F.max(vals, axis=1, only_index=True)

        points = points[:, :-1, :]
        x_hit = F.gather(points, idx_max, axis=1, batch_dims=1)
        x_hit = F.reshape(x_hit, (BR, 3))
        mask_hit = F.greater_scalar(F.sum(mask_hit, axis=1), 0)
        mask_hit = F.reshape(mask_hit, (BR, 1))

        x_hit_rm0 = x_hit
        step = F.reshape(step, (BR, 1))
        raydir = F.reshape(raydir, (BR, 3))
        x_hit_rm1 = x_hit_rm0 + step * raydir

        return x_hit_rm0, x_hit_rm1, mask_hit

    def backward_impl(self, inputs, outputs, propagate_down, accum):
        if not propagate_down[0] and not propagate_down[1]:
            return

        # Accum is zero addition, thus do not care.
        if propagate_down[0]:
            if not accum[0]:
                inputs[0].grad.fill(0)

        if propagate_down[1]:
            if not accum[1]:
                inputs[1].grad.fill(1)


def ray_trace(sdf_network, camloc, raydir, mask_obj=None, test=False, t_near=0, t_far=4,
              sphere_trace_itr=50, ray_march_points=100, n_chunks=10,
              max_post_itr=10, post_method="secant", eps=5e-5,
              ctx=None):
    func = RayTrace(ctx, sdf_network, test, t_near=t_near, t_far=t_far,
                    sphere_trace_itr=sphere_trace_itr, ray_march_points=ray_march_points,
                    n_chunks=n_chunks,
                    max_post_itr=max_post_itr,
                    post_method=post_method,
                    eps=eps)
    mask_obj = mask_obj if mask_obj is not None else nn.Variable()  # dummy
    return func(camloc, raydir, mask_obj)


def main(args):
    from network import implicit_network

    # Setting
    # nn.set_auto_forward(True)
    ctx = get_extension_context('cudnn', device_id=args.device_id)
    nn.set_default_context(ctx)
    D = args.depth
    L = args.layers
    W = args.width
    H = args.height
    R = H * W
    z_orientation = 1

    # Camera parameters
    camera = Camera(image_width=W, image_height=H, z_orientation=z_orientation)
    camloc = np.array([0.75, 0.5, 1])
    camloc = (camloc / np.sum(camloc ** 2) ** 0.5) * 2
    to = np.array([0, 0, 0])
    Rt_inv = look_at(camloc, to, z_orientation=z_orientation)
    R_inv = Rt_inv[:3, :3]
    fov = 90
    K_inv = camera.compute_intrinsic_inv(fov)

    # Rays
    x, y = np.meshgrid(np.arange(W), np.arange(H), indexing="xy")
    xy = np.asarray([x.flatten(), y.flatten()])
    xy1 = np.concatenate([xy, np.ones(R)[np.newaxis, :]])
    raydir = R_inv.dot(K_inv.dot(xy1))
    raydir = raydir / np.sum(raydir ** 2, axis=0) ** 0.5
    raydir = raydir.transpose((1, 0))

    # Network
    camloc = nn.Variable.from_numpy_array(camloc[np.newaxis, ...])
    raydir = nn.Variable.from_numpy_array(raydir[np.newaxis, ...])
    sdf_net = partial(implicit_network, D=D, L=L,
                      initial_sphere_radius=args.initial_sphere_radius)
    sdf_net0 = sdf_net

    def sdf_net0(x):
        out = sdf_net(x)
        sdf = out[..., 0][..., np.newaxis]
        return sdf

    # Sphere trace
    t_near = args.t_near
    t_far = args.t_far
    sphere_trace_itr = args.sphere_trace_itr
    ray_march_points = args.ray_march_points
    n_chunks = args.n_chunks
    max_post_itr = args.max_post_itr
    post_method = args.post_method
    eps = args.eps
    st = time.time()
    x_hit, mask_hit, dists, _, _ = ray_trace(sdf_net0, camloc, raydir, test=True,
                                             t_near=t_near, t_far=t_far,
                                             sphere_trace_itr=sphere_trace_itr,
                                             ray_march_points=ray_march_points,
                                             n_chunks=n_chunks,
                                             max_post_itr=max_post_itr,
                                             post_method=post_method, eps=eps)

    x_hit.need_grad = False
    dists.need_grad = False
    mask_hit.need_grad = False

    x_curr = x_hit
    F.sink(*[x_curr, mask_hit]).forward(clear_buffer=False)
    # Lighting
    x_curr = x_curr.get_unlinked_variable(need_grad=True)
    sdf = sdf_net0(x_curr)
    normal = nn.grad([sdf], [x_curr])[0]
    normal = F.norm_normalization(normal, axes=normal.ndim - 1, eps=1e-24)
    dlight = DistantLight()
    cos = lambert(normal, dlight.direction.reshape([3, 1])).reshape((1, H, W))
    mask_hit = mask_hit.get_unlinked_variable(need_grad=False)
    mask_hit = F.reshape(mask_hit, (1, H, W))
    mask_hit = F.broadcast(mask_hit, (3, H, W))
    image = mask_hit * 255.0 * cos
    image.forward(clear_buffer=True)

    cv2.imwrite(f"sphere_{W}x{H}_sti{sphere_trace_itr:03d}_mpi{max_post_itr:03d}_{args.post_method}.png",
                image.d.transpose(1, 2, 0))
    print(
        f"Bidirectional sphere trace/ray march (W={W}, H={H}): {time.time() - st} [s]")


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(description="Sphere Tracinig.")
    parser.add_argument('--device-id', '-d', type=int, default="0")
    parser.add_argument('--depth', '-D', type=int, default=200)
    parser.add_argument('--layers', '-L', type=int, default=9)
    parser.add_argument('--width', '-W', type=int, default=240)
    parser.add_argument('--height', '-H', type=int, default=180)
    parser.add_argument('--t-near', '-tf', type=int, default=0)
    parser.add_argument('--t-far', '-tn', type=int, default=4)
    parser.add_argument('--sphere-trace-itr', '-sti', type=int, default=10)
    parser.add_argument('--ray-march-points', '-N', type=int, default=100)
    parser.add_argument('--n-chunks', '-NC', type=int, default=10)
    parser.add_argument('--max-post-itr', '-mpi', type=int, default=8)
    parser.add_argument('--eps', type=float, default=5e-5)
    parser.add_argument('--initial-sphere-radius', type=float, default=0.75)
    parser.add_argument('--post-method', type=str, default="bisection")
    parser.add_argument('--fov', type=float, default=90)
    args = parser.parse_args()
    main(args)


# shape-reconstruction/implicit-differentiable-renderer/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 nnabla as nn
import numpy as np
from nnabla.ext_utils import get_extension_context
from nnabla.monitor import Monitor, MonitorImage

from tqdm import tqdm

from dataset import DTUMVSDataSource
from helper import generate_raydir_camloc, generate_all_pixels
from network import render as _render


def render(pose, intrinsic, mask_obj, conf):
    assert conf.height % conf.batch_height == 0, \
      f"conf.height ({conf.height}) % conf.batch_height ({conf.batch_height}) != 0"

    W, H = conf.width, conf.height
    bh = conf.batch_height

    xy = generate_all_pixels(W, H)
    xy = xy.reshape((1, H, W, 2))

    camloc = nn.Variable([1, 3])
    raydir = nn.Variable([1, bh * W, 3])

    with nn.auto_forward(False):
        color_pred = _render(camloc, raydir, conf).reshape((1, bh, W, 3))

    rimage = np.ndarray([1, H, W, 3])
    for h in tqdm(range(0, H, bh), desc="Rendering"):
        xy_h = xy[:, h:h+bh, :, :].reshape((1, bh * W, 2))
        raydir.d, camloc.d = generate_raydir_camloc(pose, intrinsic, xy_h)

        color_pred.forward(clear_buffer=True)
        rimage[0, h:h+bh, :, :] = color_pred.d.copy()

    rimage = rimage * mask_obj
    return rimage.transpose((0, 3, 1, 2))  # NCHW


def main(args):
    # Setting
    device_id = args.device_id
    conf = args.conf
    path = conf.data_path
    B = conf.batch_size
    R = conf.n_rays
    L = conf.layers
    D = conf.depth
    feature_size = conf.feature_size

    ctx = get_extension_context('cudnn', device_id=device_id)
    nn.set_default_context(ctx)

    # Dataset
    ds = DTUMVSDataSource(path, R, shuffle=True)

    # Monitor
    monitor_path = "/".join(args.model_load_path.split("/")[0:-1])
    monitor = Monitor(monitor_path)
    monitor_image = MonitorImage(
        f"Rendered image synthesis", monitor, interval=1)

    # Load model
    nn.load_parameters(args.model_load_path)

    # Render
    pose = ds.poses[conf.valid_index:conf.valid_index+1, ...]
    intrinsic = ds.intrinsics[conf.valid_index:conf.valid_index+1, ...]
    mask_obj = ds.masks[conf.valid_index:conf.valid_index+1, ...]
    image = render(pose, intrinsic, mask_obj, conf)
    monitor_image.add(conf.valid_index, image)


if __name__ == '__main__':
    import argparse
    from ruamel.yaml import YAML
    from collections import namedtuple

    parser = argparse.ArgumentParser(
        description="Implicit Differentiable Renderer Training.")
    parser.add_argument('--device-id', '-d', type=int, default="0")
    parser.add_argument('--model-load-path', type=str, required=True)
    parser.add_argument('--config', type=str,
                        default="conf/default.yaml", required=True)

    args = parser.parse_args()
    with open(args.config, "r") as f:
        conf = YAML(typ='safe').load(f)
        conf = namedtuple("Conf", conf)(**conf)
    args.conf = conf

    main(args)


# shape-reconstruction/implicit-differentiable-renderer/test_postmethod.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 numpy as np

from ray_tracer import bisection, secant

import pytest


def f(x, a=1, alpha=0, beta=5, gamma=10):
    return a * (x - alpha) * (x - beta) * (x - gamma)


@pytest.mark.parametrize("seed", [313])
@pytest.mark.parametrize("method", [bisection, secant])
@pytest.mark.parametrize("max_post_itr", [10])
@pytest.mark.parametrize("a, alpha, beta, gamma, eps0, eps1",
                         [(1, 0, 5, 10, 1, 2)])
@pytest.mark.parametrize("same", [False, True])
def test_secant(seed, method, max_post_itr,
                a, alpha, beta, gamma, eps0, eps1, same):

    with nn.auto_forward():
        start = alpha + eps0
        finish = gamma - eps1 if not same else start
        x0 = nn.Variable.from_numpy_array(start).data
        x1 = nn.Variable.from_numpy_array(finish).data
        x0, x1 = method(x0, x1, f, max_post_itr)

    expected = beta if not same else start
    np.testing.assert_allclose(x0.data, expected, atol=1e-2)


# shape-reconstruction/implicit-differentiable-renderer/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 nnabla as nn
import nnabla.solvers as S
import numpy as np
import nnabla_ext
from nnabla.ext_utils import get_extension_context
from nnabla.monitor import Monitor, MonitorSeries, MonitorTimeElapsed, MonitorImage

from dataset import data_iterator_dtumvs, DTUMVSDataSource
from helper import generate_raydir_camloc, create_monitor_path
from network import idr_loss
from synthesize import render

# TODO: impl camera parameter optimization


def main(args):
    # Setting
    device_id = args.device_id
    conf = args.conf
    path = args.conf.data_path
    B = conf.batch_size
    R = conf.n_rays
    L = conf.layers
    D = conf.depth
    feature_size = conf.feature_size

    # Dataset
    ds = DTUMVSDataSource(path, R, shuffle=True)
    di = data_iterator_dtumvs(ds, B)

    camloc = nn.Variable([B, 3])
    raydir = nn.Variable([B, R, 3])
    alpha = nn.Variable.from_numpy_array(conf.alpha)
    color_gt = nn.Variable([B, R, 3])
    mask_obj = nn.Variable([B, R, 1])

    # Monitor
    interval = di.size
    monitor_path = create_monitor_path(conf.data_path, args.monitor_path)
    monitor = Monitor(monitor_path)
    monitor_loss = MonitorSeries("Training loss", monitor, interval=interval)
    monitor_mhit = MonitorSeries("Hit count", monitor, interval=1)
    monitor_color_loss = MonitorSeries(
        "Training color loss", monitor, interval=interval)
    monitor_mask_loss = MonitorSeries(
        "Training mask loss", monitor, interval=interval)
    monitor_eikonal_loss = MonitorSeries(
        "Training eikonal loss", monitor, interval=interval)
    monitor_time = MonitorTimeElapsed(
        "Training time", monitor, interval=interval)
    monitor_image = MonitorImage("Rendered image", monitor, interval=1)

    # Solver
    solver = S.Adam(conf.learning_rate)
    loss, color_loss, mask_loss, eikonal_loss, mask_hit = \
        idr_loss(camloc, raydir, alpha, color_gt, mask_obj, conf)
    solver.set_parameters(nn.get_parameters())

    # Training loop
    for i in range(conf.train_epoch):

        ds.change_sampling_idx()

        # Validate
        if i % conf.valid_epoch_interval == 0 and not args.skip_val:
            def validate(i):
                pose_ = ds.poses[conf.valid_index:conf.valid_index+1, ...]
                intrinsic_ = ds.intrinsics[conf.valid_index:conf.valid_index+1, ...]
                mask_obj_ = ds.masks[conf.valid_index:conf.valid_index+1, ...]
                image = render(pose_, intrinsic_, mask_obj_, conf)
                monitor_image.add(i, image)
                nn.save_parameters(f"{monitor_path}/model_{i:05d}.h5")
            validate(i)

        # Train
        for j in range(di.size):
            # Feed data
            color_, mask_, intrinsic_, pose_, xy_ = di.next()
            color_gt.d = color_
            mask_obj.d = mask_
            raydir_, camloc_ = generate_raydir_camloc(pose_, intrinsic_, xy_)
            raydir.d = raydir_
            camloc.d = camloc_

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

            # Monitor
            t = i * di.size + j
            monitor_mhit.add(t, np.sum(mask_hit.d))
            monitor_loss.add(t, loss.d)
            monitor_color_loss.add(t, color_loss.d)
            monitor_mask_loss.add(t, mask_loss.d)
            monitor_eikonal_loss.add(t, eikonal_loss.d)
            monitor_time.add(t)

        # Decay
        if i in conf.alpha_decay:
            alpha.d = alpha.d * 2.0
        if i in conf.lr_decay:
            solver.set_learning_rate(solver.learning_rate() * 0.5)

    validate(i)


if __name__ == '__main__':
    import argparse
    import shutil
    import os
    import glob
    from ruamel.yaml import YAML
    from collections import namedtuple

    parser = argparse.ArgumentParser(
        description="Implicit Differentiable Renderer Training.")
    parser.add_argument('--device-id', '-d', type=int, default="0")
    parser.add_argument('--monitor-path', type=str, default="results")
    parser.add_argument('--config', type=str,
                        default="conf/default.yaml", required=True)
    parser.add_argument('--skip-val', action="store_true")

    args = parser.parse_args()
    with open(args.config, "r") as f:
        conf = YAML(typ='safe').load(f)
        conf = namedtuple("Conf", conf)(**conf)
    args.conf = conf

    monitor_path = create_monitor_path(conf.data_path, args.monitor_path)
    os.makedirs(monitor_path, exist_ok=True)
    shutil.copy(args.config, monitor_path)
    [shutil.copy(f, monitor_path) for f in glob.glob("*.py")]

    ctx = get_extension_context('cudnn', device_id=args.device_id)
    nn.set_default_context(ctx)
    # nn.set_auto_forward(True)

    main(args)


# shape-reconstruction/implicit-geometric-regularization/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.

from __future__ import division

import numpy as np

import nnabla as nn
import nnabla.logger as logger
import nnabla.monitor as M
import nnabla.functions as F
import nnabla.parametric_functions as PF


def get_args(mode="train"):
    import argparse
    import os
    parser = argparse.ArgumentParser(
        description='''Implicit Geometric Regularization.
        ''')
    parser.add_argument("-f", "--fpath", type=str,
                        help='File path to a mesh or point cloud data.')
    parser.add_argument("-b", "--batch-size", type=int, default=512,
                        help='Batch size.')
    parser.add_argument("-sb", "--sub-batch-size", type=int, default=512,
                        help='Sub batch size for computing the pair-wise distance.')
    parser.add_argument("-d", "--device-id", type=str, default="0",
                        help='Device ID.')
    parser.add_argument("-k", "--knn", type=int, default=50,
                        help='K-Nearest Neighbors.')
    parser.add_argument("-vf", "--volume-factor", type=float, default=1.5,
                        help='Volume factors to a cubic space [-1, 1]^3.')
    parser.add_argument("-tr", "--test-rate", type=float, default=0.25,
                        help='Test sample rate for the entire dataset.')
    parser.add_argument("-te", "--train-epoch", type=int, default=1500,
                        help='Training Epoch.')
    parser.add_argument("-vi", "--valid-iter", type=int, default=100,
                        help='Validation Epoch.')
    parser.add_argument("-sie", "--save-interval-epoch", type=int, default=100,
                        help='Save interval for epoch. Model is saved at every this epochs.')
    parser.add_argument("-m", "--monitor-path", type=str, default="tmp.monitor",
                        help='Monitor path.')
    parser.add_argument("-mlp", "--model-load-path", type=str, default="",
                        help='MOdel load path.')
    parser.add_argument("-mdp", "--mesh-data-path", type=str, default="",
                        help='Mesh data path.')
    parser.add_argument("--lam", type=float, default=0.1,
                        help='Coefficient to the impcilit geometric regulartization.')
    parser.add_argument("--tau", type=float, default=1.0,
                        help='Coefficient to the normals loss.')
    parser.add_argument("-lr", "--learning-rate", type=float, default=1e-4,
                        help='Learning rate.')
    parser.add_argument("--beta0", type=float, default=0.9,
                        help='First-order momentum of Adam.')
    parser.add_argument("--beta1", type=float, default=0.999,
                        help='Second-order momentum of Adam.')
    parser.add_argument("--dims", type=int, default=512,
                        help='Dimensions of MLP.')
    parser.add_argument("--ldims", type=int, default=0,
                        help='Latent dimensions.')
    parser.add_argument("--grid-size", type=int, default=256,
                        help='Grid size for computing the volume, SDF is computed for the grid-size^3'
                        'The grid is first normalized in [-1, 1], then the volume factor is multiplied.')
    parser.add_argument('--mesh-colors', type=float, nargs='+',
                        help='Mesh colors')
    parser.add_argument('--with-normals', action="store_true",
                        help='Use the normals loss.')
    parser.add_argument('-gd', '--gradient-direction', type=str, default="ascent",
                        choices=["ascent", "descent"],
                        help='Gradient direction for the marching_cubes_lewiner.')

    # Subcommands
    subparsers = parser.add_subparsers(dest="command")
    parser_mesh = subparsers.add_parser("mesh",
                                        help='')
    parser_mesh.add_argument("-f", "--fpath", type=str,
                             help='File path to a mesh data.')
    parser_mesh.add_argument("-c", "--cpath", type=str,
                             help='Path to a camera parameters.')
    parser_pcd = subparsers.add_parser("pcd",
                                       help='')
    parser_pcd.add_argument("-f", "--fpath", type=str,
                            help='File path to a point cloud data.')
    parser_pcd.add_argument("-c", "--cpath", type=str,
                            help='Path to a camera parameters.')
    parser_images = subparsers.add_parser("images",
                                          help='Create a video based on the mesh file in the same directory.')
    parser_images.add_argument("-f", "--fpath", type=str,
                               help='File path to a mesh data.')
    parser_pcdmesh = subparsers.add_parser("pcdmesh",
                                           help='Synchronized two views of a pcd sampled from a mesh and the mesh.')
    parser_pcdmesh.add_argument("-f", "--fpath", type=str,
                                help='File path to a mesh data.')
    parser_pcdmesh.add_argument(
        "-r", "--rate", help="Sampling rate.", type=float, default=0.1)

    parser_create = subparsers.add_parser("create",
                                          help='Create dataset.')
    parser_create.add_argument("-f", "--mesh-data-path", type=str,
                               help='File path to a mesh data.')

    args = parser.parse_args()

    if not os.path.exists(args.monitor_path):
        os.makedirs(args.monitor_path)
    return args


# shape-reconstruction/implicit-geometric-regularization/datasets.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 division

import glob
import os
import numpy as np
from scipy import spatial
import open3d as o3d

import nnabla as nn
import nnabla.logger as logger
import nnabla.monitor as M
import nnabla.functions as F
import nnabla.parametric_functions as PF
from nnabla.utils.data_iterator import data_iterator, data_iterator_simple
from nnabla.utils.data_source import DataSource
import utils

from args import get_args


class PointCloudDataSource(DataSource):

    def __init__(self, fpath, knn=50, test_rate=0.25, test=False, shuffle=True, rng=None):
        super(PointCloudDataSource, self).__init__(shuffle=shuffle)
        self.knn = knn
        self.test_rate = 0.25
        self.rng = np.random.RandomState(313) if rng is None else rng

        # Split info
        pcd = self._read_dataset(fpath)
        total_size = len(pcd.points)
        test_size = int(total_size * test_rate)
        indices = self.rng.permutation(total_size)
        test_indices = indices[:test_size]
        train_indices = indices[test_size:]
        indices = test_indices if test else train_indices
        self._size = test_size if test else total_size - test_size
        # Points
        points = np.asarray(pcd.points)
        self._points = self._preprocess(points)[indices]
        # Normals
        normals = np.asarray(pcd.normals)
        self._normals = normals[indices] if self.has_normals(
            normals) else normals
        # Radius
        self._radius = self._compute_radius(self._points, self.knn)
        self._variables = ('points', 'normals', 'radius')
        self.reset()

        logger.info("Data size = {}".format(self._size))

    def has_normals(self, normals):
        return False if normals.shape[0] == 0 else True

    def _preprocess(self, points):
        return utils.normalize(points)

    def _compute_radius(self, points, knn):
        if knn < 0:
            logger.info("Radius is not computed.")
            return
        # KDTree
        logger.info(
            "Constructing KDTree and querying {}-nearest neighbors".format(self.knn))
        tree = spatial.cKDTree(points, compact_nodes=True)
        # exclude self by adding 1
        dists, indices = tree.query(points, k=knn + 1)
        return dists[:, -1].reshape(dists.shape[0], 1)

    def _read_dataset(self, fpath):
        pcd = utils.read_pcd(fpath)
        return pcd

    def _get_data(self, position):
        points = self._points[self._indices[position]]
        normals = self._normals[self._indices[position]
                                ] if self.has_normals(self._normals) else [0.0]
        radius = self._radius[self._indices[position]]
        return points, normals, radius

    @property
    def points(self):
        return self._points

    @property
    def normals(self):
        return self._normals

    @property
    def radius(self):
        return self._radius

    def reset(self):
        self._indices = self.rng.permutation(self._size) \
            if self._shuffle else np.arange(self._size)
        return super(PointCloudDataSource, self).reset()


def point_cloud_data_source(fpath, knn=50, test_rate=0.25, test=False, shuffle=True, rng=None):
    return PointCloudDataSource(fpath, knn, test_rate, test, shuffle, rng)


def point_cloud_data_iterator(data_source, batch_size):
    return data_iterator(data_source, batch_size=batch_size,
                         with_memory_cache=False,
                         with_file_cache=False)


def create_pcd_dataset_from_mesh(fpath):
    mesh = utils.read_mesh(fpath)
    pcd = mesh.sample_points_poisson_disk(len(mesh.vertices),
                                          use_triangle_normal=True,
                                          seed=412)
    dpath = "/".join(fpath.split("/")[:-1])
    fname = fpath.split("/")[-1]
    fname = "{}_pcd.ply".format(os.path.splitext(fname)[0])
    fpath = os.path.join(dpath, fname)
    logger.info("PointCloud data ({}) is being created.".format(fpath))
    utils.write_pcd(fpath, pcd)


def main():
    args = get_args()
    if args.command == "create":
        create_pcd_dataset_from_mesh(args.mesh_data_path)


if __name__ == '__main__':
    main()


# shape-reconstruction/implicit-geometric-regularization/evaluate.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 division

import os
import tqdm
import numpy as np
from skimage import measure
import open3d as o3d

import nnabla as nn
import nnabla.logger as logger
import nnabla.monitor as M
from nnabla.ext_utils import get_extension_context

from args import get_args
from models import MLP
from datasets import point_cloud_data_source, point_cloud_data_iterator
import utils


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

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

    # Sample from mesh (unnormalized)
    mesh = utils.read_mesh(args.mesh_data_path)
    pcd = mesh.sample_points_poisson_disk(ds.size, 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)

    # Chamfer distance
    print("----- Chamfer distance -----")
    log = """
    One-sided Chamfer distance (Pred, True):   {}
    One-sided Chamfer distance (True, Pred):   {}
    Chamfer distance:                          {}
    """.format(cd0, cd1, cd)
    print(log)

    # Hausdorff distance
    print("----- Hausdorff distance -----")
    log = """
    One-sided Hausdorff distance (Pred, True): {}
    One-sided Hausdorff distance (True, Pred): {}
    Hausdorff distance:                        {}
    """.format(hd0, hd1, hd)
    print(log)


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


# shape-reconstruction/implicit-geometric-regularization/infer.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 division

import os
import tqdm
import numpy as np
import open3d as o3d

import nnabla as nn
import nnabla.logger as logger
import nnabla.monitor as M
from nnabla.ext_utils import get_extension_context

from args import get_args
from models import MLP
import utils


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

    # Network
    nn.load_parameters(args.model_load_path)
    model = MLP(args.dims, args.ldims, test=True)

    # Compute points and values
    pts, vol = utils.compute_pts_vol(model, args.grid_size, args.volume_factor,
                                     args.sub_batch_size)
    mesh = utils.create_mesh_from_volume(vol, args.gradient_direction)

    # Save as mesh
    dirname, pname = args.model_load_path.split("/")
    fname = os.path.splitext(pname)[0]
    fpath = "{}/{}_mesh.ply".format(dirname, fname)
    logger.info("Saving the mesh file to {}".format(fpath))
    utils.write_mesh(fpath, mesh)


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


# shape-reconstruction/implicit-geometric-regularization/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.

from __future__ import division

import numpy as np

import nnabla as nn
import nnabla.logger as logger
import nnabla.monitor as M
import nnabla.initializer as I
import nnabla.functions as F
import nnabla.parametric_functions as PF

from functools import partial


class MLP():

    def __init__(self, dims=512, ldims=0, test=False, act="leaky_relu"):
        self.dims = dims
        self.ldims = ldims
        self.test = test
        self.act_map = dict(
            softplus=partial(self.softplus, beta=100),
            relu=F.relu,  # got nan
            leaky_relu=F.leaky_relu)
        self.act = self.act_map[act]

    def softplus(self, x, beta=1):
        return (1.0 / beta) * F.log(1.0 + F.exp(beta * x))

    def affine_act(self, x, dims, name):
        c = x.shape[1]
        s = I.calc_normal_std_he_forward(c, dims)
        w_init = I.NormalInitializer(s, )
        return self.act(PF.affine(x, dims, w_init=w_init, name=name))

    def last_affine(self, x, dims, name):
        c = x.shape[1]
        l, u = I.calc_uniform_lim_glorot(c, 1)
        w_init = I.UniformInitializer((l, u))
        return PF.affine(x, 1, w_init=w_init, name=name)

    def __call__(self, x, z=None):
        b, c = x.shape[0:2]
        h = x
        h = self.affine_act(h, self.dims, name="fc0")
        h = self.affine_act(h, self.dims, name="fc1")
        h = self.affine_act(h, self.dims, name="fc2")
        h = self.affine_act(h, self.dims - (self.ldims + c), name="fc3")
        h = F.concatenate(
            *[x, z, h], axis=1) if z is not None else F.concatenate(*[x, h], axis=1)
        h = self.affine_act(h, self.dims, name="fc4")
        h = self.affine_act(h, self.dims, name="fc5")
        h = self.affine_act(h, self.dims, name="fc6")
        y = self.last_affine(h, 1, name="fc7")
        return y


if __name__ == '__main__':
    x = nn.Variable([1, 3])
    model = MLP()
    y = model(x)

    class PrintFunc():
        def __call__(sef, func):
            print(func.info.type_name)
            print("\tinp: {}".format([i.shape for i in func.inputs]))
            print("\tout: {}".format([o.shape for o in func.outputs]))
    y.visit(PrintFunc())


# shape-reconstruction/implicit-geometric-regularization/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 __future__ import division

import tqdm
import os
import numpy as np

import nnabla as nn
import nnabla.logger as logger
import nnabla.monitor as M
import nnabla.functions as F
import nnabla.parametric_functions as PF
import nnabla.solvers as S
from nnabla.ext_utils import get_extension_context

from args import get_args
from models import MLP
from datasets import point_cloud_data_source, point_cloud_data_iterator

import utils


def points_loss(y):
    return F.sum(y ** 2.0).apply(persistent=True)


def normals_loss(y, x, n, with_normals):
    loss_normals = 0.0
    if args.with_normals:
        grads = nn.grad([y], [x])
        g = grads[0]
        loss_normals = F.sum((g - n) ** 2.0, axis=[1]) ** 0.5
        loss_normals = F.sum(loss_normals).apply(persistent=True)
    return loss_normals


def eikonal_reg(x, r, model, volume_factor):
    # volumetric space and neighbor points
    xu = (2 * args.volume_factor) * F.rand(shape=x.shape) - args.volume_factor
    xn = x + 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):
    """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).

        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}

    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)
        partial_info["states_path"] = states_path

        # 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

        checkpoint_info[solvername] = partial_info

    # save parameters.
    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)
    checkpoint_info["params_path"] = params_path
    checkpoint_info["current_iter"] = current_iter

    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 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
                    }

    """

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

    with open(path, 'r') as f:
        checkpoint_info = json.load(f)

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

    logger.info("Checkpoint load (.json): {}".format(path))

    # load parameters (stored in global).
    params_path = checkpoint_info["params_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 = partial_info["states_path"]
        assert os.path.isfile(states_path), "states file not found."

        # set solver states.
        solver_obj.load_states(states_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/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 nnabla as nn
import numpy as np
from nnabla.logger import logger


def set_random_pseed(comm):
    x = nn.Variable.from_numpy_array(np.random.randint(low=0, high=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('%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/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)
        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):
        if self.is_synced:
            return

        for name, logger in sorted(self.name2logger.items(), key=lambda x: x):
            # 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], division=False, inplace=True)
                    self.comm.all_reduce(
                        [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"))

        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):
        if sync:
            self.sync_all(reset=reset)

        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):
        if sync:
            self.sync_all(reset=reset)

        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/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/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()


# 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 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


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:
            tar.extractall(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


# 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/config.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


def _check_args_has(args, attr):
    assert isinstance(attr, str), f"attr must be str not {type(attr)}."
    assert hasattr(args, attr), f"args must have `{attr}` as an attribute."


def _refine_if_none(val, default):
    return default if val is None else val


def _refine_image_shape(args, default_resolution):
    args.image_size = _refine_if_none(args.image_size, default_resolution)

    args.image_shape = (args.batch_size, 3, args.image_size, args.image_size)


def _refine_lr(args):
    _check_args_has(args, "image_shape")
    _check_args_has(args, "loss_scaling")

    if max(*args.image_shape[-2:]) <= 64:
        args.lr = _refine_if_none(args.lr, 1e-4)
        args.grad_clip = -1  # disable gradient clipping
    else:
        # bigger than 64
        args.lr = _refine_if_none(args.lr, 2e-5)
        args.grad_clip = 1.

    if args.loss_scaling > (1.0 + 1e-5):
        args.lr /= args.loss_scaling


def _celebahq(args):
    # default resolution is 256
    _refine_image_shape(args, 256)

    args.dataset_root_dir = os.path.join(args.data_dir, "images")


def _cifar10(args):
    # default resolution is 32
    _refine_image_shape(args, 32)

    # cifar10 data iterator will automatically download and manage dataset.
    args.dataset_root_dir = None


def _imagenet(args):
    # default resolution is 64
    _refine_image_shape(args, 64)

    if args.image_size <= 64:
        args.dataset_root_dir = os.path.join(args.data_dir, "train_64x64")
    else:
        args.dataset_root_dir = os.path.join(args.data_dir, "ilsvrc2012")


def _custum_dataset(args):
    # default resolution is 256
    _refine_image_shape(args, 256)

    _check_args_has(args, "dataset_root_dir")


def _common(args):
    _check_args_has(args, "image_shape")

    _refine_lr(args)

    B, C, H, W = args.image_shape

    if H <= 32 and W <= 32:
        args.channel_mult = _refine_if_none(args.channel_mult, (1, 2, 2, 2))
        args.base_channels = _refine_if_none(args.base_channels, 128)
        args.num_res_blocks = _refine_if_none(args.num_res_blocks, 3)
    elif H <= 64 and W <= 64:
        args.channel_mult = _refine_if_none(args.channel_mult, (1, 2, 3, 4))
        args.base_channels = _refine_if_none(args.base_channels, 128)
        args.num_res_blocks = _refine_if_none(args.num_res_blocks, 3)
    else:
        args.channel_mult = _refine_if_none(
            args.channel_mult, (1, 1, 2, 2, 4, 4))
        args.base_channels = _refine_if_none(args.base_channels, 192)
        args.num_res_blocks = _refine_if_none(args.num_res_blocks, 2)

    # todo: more larger model


def refine_args_by_dataset(args):
    _check_args_has(args, "dataset")
    _check_args_has(args, "batch_size")
    _check_args_has(args, "data_dir")

    if args.dataset == "celebahq":
        _celebahq(args)
    elif args.dataset == "cifar10":
        _cifar10(args)
    elif args.dataset.startswith("imagenet"):
        _imagenet(args)
    else:
        _custum_dataset(args)

    _common(args)


__all__ = [refine_args_by_dataset]


# diffusion-models/diffusion.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 enum
import math
import numpy as np
import nnabla as nn
import nnabla.functions as F

from layers import chunk
from neu.losses import kl_normal, gaussian_log_likelihood
from neu.misc import AttrDict

# Better to implement ModelVarType as a module?


class ModelVarType(enum.Enum):
    """
    Enum class for the type of model variance.

    Each entry represents
    - FIXED_SMALL: variance of posterior q(x_{t-1}|x_t, x_0) corresponding to beta_t in the paper.
    - FIXED_LARGE: variance of forward noising process q(x_t | x_{t-1}) corresponding to beta_tilde_t in the paper
    - LEARNED_RANGE: interpolation between FIXED_SMALL and FIXED_LARGE proposed by "Improved Denoising Diffusion Probabilistic Models".
    """
    FIXED_SMALL = enum.auto()
    FIXED_LARGE = enum.auto()
    LEARNED_RANGE = enum.auto()

    @staticmethod
    def get_supported_keys():
        return [x.name for x in ModelVarType]

    @staticmethod
    def get_vartype_from_key(key):
        for elem in ModelVarType:
            if elem.name.lower() == key.lower():
                return elem

        raise ValueError(
            f"key '{key}' is not supported. Key must be one of {ModelVarType.get_supported_keys()}.")


def is_learn_sigma(model_var_type: ModelVarType):
    # If we add model variance type, modify this also.
    return model_var_type == ModelVarType.LEARNED_RANGE


def _get_alpha_bar_from_time(t):
    """
    Noise scheduling method proposed by Nichol et. al to avoid too noisy image especially for smaller resolution.

    This strategy creates beta as follows:
        alpha_bar(t) = f(t) / f(0)
        f(t) = cos((t / T + s) / (1 + s) * PI / 2) ** 2
        beta(t) = 1 - alpha_bar(t) / alpha_bar(t-1)

    where s = 0.008 (~= 1 / 127.5), which ensures sqrt(beta_0) is slightly smaller than the pixel bin size.
    """

    assert 0 <= t <= 1, "t must be normalized by max diffusion timestep (T)."

    return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2


def get_beta_schedule(strategy, num_timesteps):
    """
    Get a pre-defined beta schedule for the given strategy.
    Based on a scheduling proposed by Ho et. al, scale it to any timesteps.

    Args:
        strategy (string): Strategy to create noise schedule. Should be one of {"linear", "cosine"}.
        num_timesteps (int): Max timestep for the diffusion process.
    """
    if strategy == "linear":
        scale = 1000 / num_timesteps
        beta_start = scale * 0.0001
        beta_end = scale * 0.02
        betas = np.linspace(beta_start, beta_end,
                            num_timesteps, dtype=np.float64)
    elif strategy == "cosine":
        betas = []
        for i in range(num_timesteps):
            t1 = i / num_timesteps
            t2 = (i + 1) / num_timesteps
            max_beta = 0.999  # make beta lower than 1 to prevent singularities
            alpha_bar = 1 - \
                _get_alpha_bar_from_time(t2) / _get_alpha_bar_from_time(t1)
            betas.append(min(alpha_bar, max_beta))
            assert betas[-1] <= max_beta

        betas = np.array(betas)
    else:
        raise NotImplementedError(strategy)

    assert betas.shape == (num_timesteps,)
    return betas


def const_var(np_array):
    """
    Create constant nn.Variable from numpy array.
    """
    ret = nn.Variable.from_numpy_array(np_array, need_grad=False)
    ret.persistent = True
    return ret


def noise_like(shape, noise_function=F.randn, repeat=False):
    """
    Create noise with given shape.
    """
    if repeat:
        noise = noise_function(shape=(1, ) + shape[1:])
        return F.tile(noise, (shape[0], ) + (1, ) * (len(shape) - 1))

    return noise_function(shape=shape)


def mean_along_except_batch(x, batch_axis=0):
    return F.mean(x, axis=list(range(1, len(x.shape))))


class GaussianDiffusion(object):
    """
    An interface for gaussian diffusion process.

    Args:
        betas (numpy.ndarray): 
            A 1-D tensor of noise scales at each timestep. 
            Assume this is created by get_beta_schedule() defined above.
    """

    def __init__(self,
                 betas,
                 model_var_type=ModelVarType.FIXED_SMALL):
        self.model_var_type = model_var_type

        # check input betas
        assert (isinstance(betas, np.ndarray))
        assert (betas > 0).all() and (betas <= 1).all()

        betas = betas.astype(np.float64)
        self.np_betas = betas
        T, = betas.shape
        self.num_timesteps = int(T)

        alphas = 1. - betas
        alphas_cumprod = np.cumprod(alphas, axis=0)
        alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
        alphas_cumprod_next = np.append(alphas_cumprod[1:], 0.0)
        assert alphas_cumprod_prev.shape == (T, )

        self.betas = const_var(betas)
        self.alphas_cumprod = const_var(alphas_cumprod)
        self.alphas_cumprod_prev = const_var(alphas_cumprod_prev)
        self.alphas_cumprod_next = const_var(alphas_cumprod_next)

        # for q(x_t | x{t-1})
        self.sqrt_alphas_cumprod = const_var(np.sqrt(alphas_cumprod))
        self.sqrt_one_minus_alphas_cumprod = const_var(
            np.sqrt(1. - alphas_cumprod))
        self.log_one_minus_alphas_cumprod = const_var(
            np.log(1. - alphas_cumprod))
        self.sqrt_recip_alphas_cumprod = const_var(
            np.sqrt(1. / alphas_cumprod))
        self.sqrt_recipm1_alphas_cumprod = const_var(
            np.sqrt(1. / alphas_cumprod - 1.))

        # for q(x_{t-1} | x_t, x_0)
        posterior_var = betas * \
            (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
        self.posterior_var = const_var(posterior_var)
        self.posterior_log_var_clipped = const_var(
            np.log(np.append(posterior_var[1], posterior_var[1:])))
        self.posterior_mean_coef1 = const_var(
            betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod))
        self.posterior_mean_coef2 = const_var(
            (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod))

        # for FIXED_LARGE
        # According to the origianl author's implementation,
        # initial (log-)variance should be the posterior's one to get a better decoder log likelihood.
        np_betas_clipped = np.append(posterior_var[1], betas[1:])
        self.betas_clipped = const_var(np_betas_clipped)
        self.log_betas_clipped = const_var(np.log(np_betas_clipped))

    @staticmethod
    def _extract(a, t, x_shape=None):
        """
        Extract some coefficients at specified timesteps.
        If x_shape is not None, 
        reshape the extracted one to [batch_size, 1, 1, 1, 1, ...] corrensponding to x_shape for broadcasting purposes.
        """

        B, = t.shape
        out = F.gather(a, t, axis=0)

        if x_shape is None:
            return out

        assert x_shape[0] == B
        assert out.shape == (B, )
        out_shape = np.array(list(x_shape))
        out_shape[1:] = 1
        return F.reshape(out, out_shape)

    def q_sample(self, x_start, t, noise=None):
        """
        Diffuse the data (t == 0 means diffused for 1 step), which samples from q(x_t | x_0).
        xt = sqrt(cumprod(alpha_0, ..., alpha_t)) * x_0 + sqrt(1 - cumprod(alpha_0, ..., alpha_t)) * epsilon

        Args:
            x_start (nn.Variable): The (B, C, ...) tensor of x_0.
            t (nn.Variable): A 1-D tensor of timesteps.

        Return:
            x_t (nn.Variable): 
                The (B, C, ...) tensor of x_t.
                Each sample x_t[i] corresponds to the noisy image at timestep t[i] constructed from x_start[i].
        """
        if noise is None:
            noise = F.randn(shape=x_start.shape)
        assert noise.shape == x_start.shape
        return (
            self._extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
            self._extract(self.sqrt_one_minus_alphas_cumprod,
                          t, x_start.shape) * noise
        )

    def predict_xstart_from_noise(self, x_t, t, noise):
        """
        Predict x_0 from x_t.
        x0 = x_t / sqrt(cumprod(alpha_0, ..., aplha_t)) - sqrt(1 / cumprod(alpha_0, ..., alpha_t) - 1) * epsilon

        Args:
            x_t (nn.Variable): The (B, C, ...) tensor whose i-th element is an noisy data at timestep t[i].
            t (nn.Variable): A 1-D tensor of timesteps.
            noise (nn.Variable): The (B, C, ...) tensor of predicted noise in x_t.

        Return:
            x_0 (nn.Variable): The (B, C, ...) tensor of predicted x_0.
        """
        assert x_t.shape == noise.shape
        return (
            self._extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
            self._extract(self.sqrt_recipm1_alphas_cumprod,
                          t, x_t.shape) * noise
        )

    def predict_noise_from_xstart(self, x_t, t, pred_xstart):
        """
        Predict epsilon from x_t and predicted x_0.
        epsilon = (x_t / sqrt(cumprod(alpha_0, ..., alpha_t)) - x_0) / (sqrt(1 / cumprod(alpha_0, ..., alpha_t) - 1))

        Args:
            x_t (nn.Variable): The (B, C, ...) tensor whose i-th element is an noisy data at timestep t[i].
            t (nn.Variable): A 1-D tensor of timesteps.
            noise (nn.Variable): The (B, C, ...) tensor of predicted noise in x_t.

        Return:
            x_0 (nn.Variable): The (B, C, ...) tensor of predicted x_0.
        """
        assert x_t.shape == pred_xstart.shape

        return (
            self._extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
            - pred_xstart
        ) / self._extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)

    def q_posterior(self, x_start, x_t, t):
        """
        Compute the mean and variance of the diffusion posterior q(x_{t-1} | x_t, x_0).
        """
        assert x_start.shape == x_t.shape
        posterior_mean = (
            self._extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +
            self._extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
        )
        posterior_var = self._extract(self.posterior_var, t, x_t.shape)
        posterior_log_var_clipped = self._extract(
            self.posterior_log_var_clipped, t, x_t.shape)
        assert (posterior_mean.shape[0] == posterior_var.shape[0] == posterior_log_var_clipped.shape[0] ==
                x_start.shape[0])
        return posterior_mean, posterior_var, posterior_log_var_clipped

    def _vlb_in_bits_per_dims(self, model, x_start, x_t, t,
                              clip_denoised=True):
        """
        Calculate variational lower bound in bits/dims.
        """
        B, C, H, W = x_start.shape
        assert x_start.shape == x_t.shape
        assert t.shape == (B, )

        # true parameters
        mean, _, log_var_clipped = self.q_posterior(x_start, x_t, t)

        # pred parameters
        preds = self.p_mean_var(model, x_t, t, clip_denoised=clip_denoised)

        # Negative log-likelihood
        nll = -gaussian_log_likelihood(x_start,
                                       mean=preds.mean, logstd=0.5 * preds.log_var)
        nll_bits = mean_along_except_batch(nll) / np.log(2.0)
        assert nll.shape == x_start.shape
        assert nll_bits.shape == (B, )

        # kl between true and pred in bits
        kl = kl_normal(mean, log_var_clipped, preds.mean, preds.log_var)
        kl_bits = mean_along_except_batch(kl) / np.log(2.0)
        assert kl.shape == x_start.shape
        assert kl_bits.shape == (B, )

        # Return nll at t = 0, otherwise KL(q(x_{t-1}|x_t,x_0)||p(x_{t-1}|x_t))
        return F.where(F.equal_scalar(t, 0), nll_bits, kl_bits)

    def train_loss(self, model, x_start, t, noise=None):
        """
        Calculate training loss for given data and model.

        Args:
            model (callable): 
                A trainable model to predict noise in data conditioned by timestep.
                This function should perform like pred_noise = model(x_noisy, t).
                If self.model_var_type is the one that requires prediction for sigma, model has to output them as well.
            x_start (nn.Variable): The (B, C, ...) tensor of x_0.
            t (nn.Variable): A 1-D tensor of timesteps.
            noise (callable or None): A noise generator. If None, F.randn(shape=x_start.shape) will be used.

        Return:
            loss (dict of {string: nn.Variable}): 
                Return dict that has losses to train the `model`.
                You can access each loss by a name that will be:
                    - `vlb`: Variational Lower Bound for learning sigma. 
                             This will be included only if self.model_var_type requires to learn sigma.
                    - `mse`: MSE between actual and predicted noise.
                Each entry is the (B, ) tensor of batched loss computed from given inputs.
                Note that this function doesn't reduce batch dim
                in order to make it easy to trace the loss value at each timestep.
                Therefore, you should take average for returned Variable over batch dim to train the model.
        """
        B, C, H, W = x_start.shape
        assert t.shape == (B, )

        if noise is None:
            noise = F.randn(shape=x_start.shape)
        assert noise.shape == x_start.shape

        # Calculate x_t from x_start, t, and noise.
        x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
        assert x_noisy.shape == x_start.shape

        # Predict noise.
        # According to the original DDPM, this is superior than reconstructing x_0.
        # If model_var_type requires to learn sigma, model must output pred_sigma as well.
        pred = model(x_noisy, t)

        # Calculate losses
        ret = AttrDict()

        if is_learn_sigma(self.model_var_type):
            # split pred into 2 variables along channel axis.
            pred_noise, pred_sigma = chunk(pred, num_chunk=2, axis=1)
            assert pred_sigma.shape == x_start.shape, \
                f"Shape mismutch between pred_sigma {pred_sigma.shape} and x_start {x_start.shape}"

            # Variational lower bound for sigma
            # Use dummy function as model, since we already got prediction from model.
            var = F.concatenate(pred_noise.get_unlinked_variable(
                need_grad=False), pred_sigma, axis=1)
            ret.vlb = self._vlb_in_bits_per_dims(model=lambda x_t, t: var,
                                                 x_start=x_start,
                                                 x_t=x_noisy,
                                                 t=t)
        else:
            pred_noise = pred

        assert pred_noise.shape == x_start.shape, \
            f"Shape mismutch between pred_noise {pred_sigma.shape} and x_start {x_start.shape}"

        ret.mse = mean_along_except_batch(F.squared_error(noise, pred_noise))

        # shape check for all losses
        for name, loss in ret.items():
            assert loss.shape == (B, ), \
                 f"A Variabla for loss `{name}` has a wrong shape ({loss.shape} != {(B, )})"

        return ret

    def p_mean_var(self, model, x_t, t, clip_denoised=True):
        """
        Compute mean and var of p(x_{t-1}|x_t) from model.

        Args:
            model (Callable): A callbale that takes x_t and t and predict noise (and more).
            x_t (nn.Variable): The (B, C, ...) tensor at timestep t (x_t).
            t (nn.Variable): A 1-D tensor of timesteps. The first axis represents batchsize.
            clip_denoised (bool): If True, clip the denoised signal into [-1, 1].

        Returns:
            An AttrDict containing the following items:
                "mean": the mean predicted by model.
                "var": the variance predicted by model (or pre-defined variance).
                "log_var": the log of "var".
                "xstart": the x_0 predicted from x_t and t by model.
        """
        B, C, H, W = x_t.shape
        assert t.shape == (B, )
        pred = model(x_t, t)

        if self.model_var_type == ModelVarType.LEARNED_RANGE:
            assert pred.shape == (B, 2 * C, H, W)
            pred_noise, pred_var_coeff = chunk(pred, num_chunk=2, axis=1)

            min_log = self._extract(
                self.posterior_log_var_clipped, t, x_t.shape)
            max_log = F.log(self._extract(self.betas, t, x_t.shape))

            # pred_var_coeff should be [0, 1]
            v = F.sigmoid(pred_var_coeff)
            model_log_var = v * max_log + (1 - v) * min_log
            model_var = F.exp(model_log_var)
        else:
            # Model only predicts noise
            pred_noise = pred

            model_log_var, model_var = {
                ModelVarType.FIXED_LARGE: lambda: (
                    self._extract(self.log_betas_clipped, t, x_t.shape),
                    self._extract(self.betas_clipped, t, x_t.shape)
                ),
                ModelVarType.FIXED_SMALL: lambda: (
                    self._extract(
                        self.posterior_log_var_clipped, t, x_t.shape),
                    self._extract(self.posterior_var, t, x_t.shape)
                )
            }[self.model_var_type]()

        x_recon = self.predict_xstart_from_noise(
            x_t=x_t, t=t, noise=pred_noise)

        if clip_denoised:
            x_recon = F.clip_by_value(x_recon, -1, 1)

        model_mean, _, _ = self.q_posterior(x_start=x_recon, x_t=x_t, t=t)

        assert model_mean.shape == x_recon.shape == x_t.shape

        assert model_mean.shape == model_var.shape == model_log_var.shape or \
            (model_mean.shape[0] == model_var.shape[0] == model_log_var.shape[0] and model_var.shape[1:] == (
                1, 1, 1) and model_log_var.shape[1:] == (1, 1, 1))

        # returns
        ret = AttrDict()
        ret.mean = model_mean
        ret.var = model_var
        ret.log_var = model_log_var
        ret.xstart = x_recon

        return ret

    def p_sample(self,
                 model,
                 x_t,
                 t,
                 clip_denoised=True,
                 noise_function=F.randn,
                 repeat_noise=False,
                 no_noise=False):
        """
        Sample from the model for one step.
        Also return predicted x_start.
        """
        preds = self.p_mean_var(model=model, x_t=x_t,
                                t=t, clip_denoised=clip_denoised)

        # no noise when t == 0
        if no_noise:
            return preds.mean, preds.xstart

        noise = noise_like(x_t.shape, noise_function, repeat_noise)
        assert noise.shape == x_t.shape

        # sample from gaussian N(model_mean, )
        return preds.mean + F.exp(0.5 * preds.log_var) * noise, preds.xstart

    # DDIM sampler
    def ddim_sample(self,
                    model,
                    x_t,
                    t,
                    clip_denoised=True,
                    noise_function=F.randn,
                    repeat_noise=False,
                    no_noise=False,
                    eta=0.):
        """
        sample x_{t-1} from x_{t} by the model using DDIM sampler.
        Also return predicted x_start.
        """
        preds = self.p_mean_var(model, x_t, t, clip_denoised=clip_denoised)

        pred_noise = self.predict_noise_from_xstart(x_t, t, preds.xstart)

        from layers import sqrt
        alpha_bar = self._extract(self.alphas_cumprod, t, x_t.shape)
        alpha_bar_prev = self._extract(self.alphas_cumprod_prev, t, x_t.shape)
        sigma = (
            eta
            * sqrt((1 - alpha_bar_prev) / (1 - alpha_bar))
            * sqrt(1 - alpha_bar / alpha_bar_prev)
        )

        mean_pred = (
            preds.xstart * sqrt(alpha_bar_prev)
            + sqrt(1 - alpha_bar_prev - sigma ** 2) * pred_noise
        )

        if no_noise:
            return mean_pred, preds.xstart

        noise = noise_like(x_t.shape, noise_function, repeat_noise)
        return mean_pred + sigma * noise, preds.xstart

    def ddim_rev_sample(self, model, x_t, t, clip_denoised=True, eta=0.0):
        """
        sample x_{t+1} from x_{t} by the model using DDIM reverse ODE.
        """
        assert eta == 0.0, "ReverseODE only for deterministic path"
        preds = self.p_mean_var(model, x_t, t, clip_denoised=clip_denoised)

        pred_noise = self.predict_noise_from_xstart(x_t, t, preds.xstart)
        alpha_bar_next = self._extract(self.alphas_cumprod_next, t, x_t.shape)

        from layers import sqrt
        return (
            preds.xstart * sqrt(alpha_bar_next)
            + sqrt(1 - alpha_bar_next) * pred_noise
        )

    def sample_loop(self, model, shape, sampler,
                    noise=None,
                    dump_interval=-1,
                    progress=False,
                    without_auto_forward=False):
        """
        Iteratively Sample data from model from t=T to t=0.
        T is specified as the length of betas given to __init__().

        Args:
            model (collable): 
                A callable that takes x_t and t and predict noise (and sigma related parameters).
            shape (list like object): A data shape.
            sampler (callable): A function to sample x_{t-1} given x_{t} and t. Typically, self.p_sample or self.ddim_sample.
            noise (collable): A noise generator. If None, F.randn(shape) will be used.
            interval (int): 
                If > 0, all intermediate results at every `interval` step will be returned as a list.
                e.g. if interval = 10, the predicted results at {10, 20, 30, ...} will be returned.
            progress (bool): If True, tqdm will be used to show the sampling progress.

        Returns:
            - x_0 (nn.Variable): the final sampled result of x_0
            - samples (a list of nn.Variable): the sampled results at every `interval`
            - pred_x_starts (a list of nn.Variable): the predicted x_0 from each x_t at every `interval`: 
        """
        T = self.num_timesteps
        indices = list(range(T))[::-1]

        samples = []
        pred_x_starts = []

        if progress:
            from tqdm.auto import tqdm
            indices = tqdm(indices)

        if without_auto_forward:
            if noise is None:
                noise = np.random.randn(*shape)
            else:
                assert isinstance(noise, np.ndarray)
                assert noise.shape == shape

            x_t = nn.Variable.from_numpy_array(noise)
            t = nn.Variable.from_numpy_array([T - 1 for _ in range(shape[0])])

            # build graph
            y, pred_x_start = sampler(model, x_t, t)
            up_x_t = F.assign(x_t, y)
            up_t = F.assign(t, t - 1)
            update = F.sink(up_x_t, up_t)

            cnt = 0
            for step in indices:
                y.forward(clear_buffer=True)
                update.forward(clear_buffer=True)

                cnt += 1
                if dump_interval > 0 and cnt % dump_interval == 0:
                    samples.append((step, y.d.copy()))
                    pred_x_starts.append((step, pred_x_start.d.copy()))
        else:
            with nn.auto_forward():
                if noise is None:
                    x_t = F.randn(shape=shape)
                else:
                    assert isinstance(noise, np.ndarray)
                    assert noise.shape == shape
                    x_t = nn.Variable.from_numpy_array(noise)
                cnt = 0
                for step in indices:
                    t = F.constant(step, shape=(shape[0], ))
                    x_t, pred_x_start = sampler(
                        model, x_t, t, no_noise=step == 0)
                    cnt += 1
                    if dump_interval > 0 and cnt % dump_interval == 0:
                        samples.append((step, x_t.d.copy()))
                        pred_x_starts.append((step, pred_x_start.d.copy()))

        assert x_t.shape == shape
        return x_t.d.copy(), samples, pred_x_starts

    def p_sample_loop(self, *args, **kwargs):
        """
        Sample data from x_T ~ N(0, I) with p(x_{t-1}|x_{t}) proposed by "Denoising Diffusion Probabilistic Models".
        See self.sample_loop for more details about sampling process.

        """

        return self.sample_loop(*args, sampler=self.p_sample, **kwargs)

    def ddim_sample_loop(self, *args, **kwargs):
        """
        Sample data from x_T ~ N(0, I) with p(x_{t-1}|x_{t}, x_{0}) proposed by "Denoising Diffusion Implicit Models".
        See self.sample_loop for more details about sampling process.

        """
        from functools import partial

        return self.sample_loop(*args,
                                sampler=partial(self.ddim_sample, eta=0.),
                                **kwargs)


# 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
import moviepy.editor as mp

import click
import nnabla as nn
from nnabla.logger import logger
import numpy as np
from nnabla.utils.image_utils import imsave
from neu.misc import AttrDict, init_nnabla
from neu.reporter import save_tiled_image, get_tiled_image
from neu.yaml_wrapper import read_yaml

from model import Model
from diffusion import ModelVarType


@click.command()
# configs for generating process
@click.option("--device-id", default='0', help="Device id.", show_default=True)
@click.option("--samples", default=16, help="# of generating samples on each device.", show_default=True)
@click.option("--batch-size", default=32, help="# of generating samples for each inference.", show_default=True)
# model configs
@click.option("--config", required=True, type=str, help="A path for config file.")
@click.option("--h5", required=True, type=str, help="A path for parameter to load.")
@click.option("--ema/--no-ema", default=True, help="Use ema params or not.")
@click.option("--ddim/--no-ddim", default=False, help="Use ddim sampler to generate data.", show_default=True)
@click.option("--sampling-interval", "-s", default=None, type=int, help="A timestep interval for sampling.")
# configs for dumping
@click.option("--output-dir", default="./outs", help="output dir", show_default=True)
@click.option("--tiled/--no-tiled", default=True, help="If true, generated images will be saved as tiled image.")
@click.option("--save-xstart/--no-save-xstart", default=False, help="If true, predicted xstarts at each timestep will be saved as video.")
def main(**kwargs):
    # set training args
    args = AttrDict(kwargs)

    assert os.path.exists(
        args.config), f"{args.config} is not found. Please make sure the config file exists."
    conf = read_yaml(args.config)

    comm = init_nnabla(ext_name="cudnn", device_id=args.device_id,
                       type_config="float", random_pseed=True)
    if args.sampling_interval is None:
        args.sampling_interval = 1

    use_timesteps = list(
        range(0, conf.num_diffusion_timesteps, args.sampling_interval))
    if use_timesteps[-1] != conf.num_diffusion_timesteps - 1:
        # The last step should be included always.
        use_timesteps.append(conf.num_diffusion_timesteps - 1)

    # setup model variance type
    model_var_type = ModelVarType.FIXED_SMALL
    if "model_var_type" in conf:
        model_var_type = ModelVarType.get_vartype_from_key(conf.model_var_type)

    model = Model(beta_strategy=conf.beta_strategy,
                  use_timesteps=use_timesteps,
                  model_var_type=model_var_type,
                  num_diffusion_timesteps=conf.num_diffusion_timesteps,
                  attention_num_heads=conf.num_attention_heads,
                  attention_resolutions=conf.attention_resolutions,
                  scale_shift_norm=conf.ssn,
                  base_channels=conf.base_channels,
                  channel_mult=conf.channel_mult,
                  num_res_blocks=conf.num_res_blocks)

    # load parameters
    assert os.path.exists(
        args.h5), f"{args.h5} is not found. Please make sure the h5 file exists."
    nn.parameter.load_parameters(args.h5)

    # Generate
    # sampling
    B = args.batch_size
    num_samples_per_iter = B * comm.n_procs
    num_iter = (args.samples + num_samples_per_iter -
                1) // num_samples_per_iter

    local_saved_cnt = 0
    for i in range(num_iter):
        logger.info(f"Generate samples {i + 1} / {num_iter}.")
        sample_out, _, x_starts = model.sample(shape=(B, ) + conf.image_shape[1:],
                                               dump_interval=1,
                                               use_ema=args.ema,
                                               progress=comm.rank == 0,
                                               use_ddim=args.ddim)

        # scale back to [0, 255]
        sample_out = (sample_out + 1) * 127.5

        if args.tiled:
            save_path = os.path.join(
                args.output_dir, f"gen_{local_saved_cnt}_{comm.rank}.png")
            save_tiled_image(sample_out.astype(np.uint8), save_path)
            local_saved_cnt += 1
        else:
            for b in range(B):
                save_path = os.path.join(
                    args.output_dir, f"gen_{local_saved_cnt}_{comm.rank}.png")
                imsave(save_path, sample_out[b].astype(
                    np.uint8), channel_first=True)
                local_saved_cnt += 1

        # create video for x_starts
        if args.save_xstart:
            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(
                args.output_dir, f"pred_x0_along_time_{local_saved_cnt}_{comm.rank}.mp4"))


if __name__ == "__main__":
    main()


# diffusion-models/layers.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.functions as F
import nnabla.parametric_functions as PF
import nnabla.initializer as I


def sinusoidal_embedding(timesteps, embedding_dim):
    """
    Sinusoidal embeddings originally proposed in "Attention Is All You Need" (https://arxiv.org/abs/1706.03762).
    """
    assert len(timesteps.shape) == 1

    half_dim = embedding_dim // 2
    denominator = -np.log(10000) / half_dim
    emb = F.exp(denominator * F.arange(start=0, stop=half_dim))
    emb = F.reshape(timesteps, (-1, 1)) * F.reshape(emb, (1, -1))
    emb = F.concatenate(F.cos(emb), F.sin(emb), axis=1)

    if embedding_dim & 1:  # zero pad to be divisible by two
        emb = F.pad(emb, [[0, 0], [0, 1]])

    assert emb.shape == (timesteps.shape[0], embedding_dim)

    return emb


def nonlinearity(x):
    return F.swish(x)


def normalize(x, name):
    with nn.parameter_scope(name):
        return PF.group_normalization(x, num_groups=32)


def conv(x, c, name, kernel=(3, 3), pad=(1, 1), stride=(1, 1), zeroing_w=False):
    # init weight and bias with uniform, which is the same as pytorch
    lim = I.calc_normal_std_he_forward(x.shape[1] * 2, c, tuple(kernel))
    w_init = I.UniformInitializer(lim=(-lim, lim), rng=None)
    b_init = I.UniformInitializer(lim=(-lim, lim), rng=None)

    if zeroing_w:
        w_init = I.ConstantInitializer(0)
        b_init = I.ConstantInitializer(0)

    return PF.convolution(x, c, kernel,
                          pad=pad, stride=stride, name=name,
                          w_init=w_init, b_init=b_init)


def nin(x, c, name, zeroing_w=False):
    lim = np.sqrt(x.shape[1]) ** -1
    w_init = I.UniformInitializer(lim=(-lim, lim))  # same as pytorch's default
    b_init = I.UniformInitializer(lim=(-lim, lim))  # same as pytorch's default

    if zeroing_w:
        w_init = I.ConstantInitializer(0)
        b_init = I.ConstantInitializer(0)

    return PF.convolution(x, c,
                          kernel=(1, 1),
                          pad=(0, 0), stride=(1, 1), name=name,
                          w_init=w_init, b_init=b_init)


def upsample(x, name, with_conv):
    with nn.parameter_scope(name):
        B, C, H, W = x.shape
        x = F.interpolate(x, scale=(2, 2), mode="nearest", align_corners=True)
        assert x.shape == (B, C, H * 2, W * 2)
        if with_conv:
            x = conv(x, C, "upsample_conv")
            assert x.shape == (B, C, H * 2, W * 2)
        return x


def downsample(x, name, with_conv):
    with nn.parameter_scope(name):
        B, C, H, W = x.shape
        if with_conv:
            x = conv(x, C, "downsample_conv",
                     kernel=(3, 3), pad=(1, 1), stride=(2, 2))
        else:
            x = F.average_pooling(x, (2, 2), stride=(2, 2))

        assert x.shape == (B, C, H // 2, W // 2)
        return x


def chunk(x, num_chunk, axis):
    """
    Split `x` to `num_chunk` arrays along specified axis.
    """
    shape = x.shape
    C = shape[axis]
    num_elems = (C + num_chunk - 1) // num_chunk

    ret = []
    for i in range(num_chunk):
        start = [0 for _ in shape]
        stop = [s for s in shape]
        start[axis] = i * num_elems
        stop[axis] = start[axis] + num_elems

        segment = F.slice(x, start=start, stop=stop)
        assert len(segment.shape) == len(x.shape)
        ret.append(segment)

    return ret


def sqrt(x):
    assert isinstance(x, (nn.Variable, nn.NdArray))
    return x ** 0.5


# diffusion-models/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 numpy as np

import nnabla as nn
import nnabla.functions as F
from functools import partial

from diffusion import ModelVarType, is_learn_sigma, get_beta_schedule, const_var, GaussianDiffusion
from unet import UNet

from neu.misc import AttrDict


def respace_betas(betas, use_timesteps):
    """
    Given betas and use_timestpes, returns respaced betas.

    Args:
        betas (np.array): A T elements of array containing beta_t at index t.
        use_timesteps (tuple):
            A tuple indicates which timesteps are used for generation process.
            For example, if use_timestpes = (0, 100, 300, 500, 700, 900, 1000), only ts included in use_timestpes will be used.

    Returns:
        new_betas (np.array): A respaced betas.
        timestep_map (np.array): 
            A list indicating how to map a new index to original index.
            It will be the same as use_timesteps sorted in increasing order.
    """

    if isinstance(use_timesteps, list):
        use_timesteps = tuple(use_timesteps)

    assert isinstance(use_timesteps, tuple)

    T = len(betas)

    new_betas = []
    prev_alphas_cumprod = 1.
    alphas_cumprod = 1.
    timestep_map = []
    for t in range(T):
        alphas_cumprod *= 1. - betas[t]
        if t in use_timesteps:
            new_beta = 1 - alphas_cumprod / prev_alphas_cumprod
            new_betas.append(new_beta)
            timestep_map.append(t)
            prev_alphas_cumprod = alphas_cumprod

    return np.asarray(new_betas, dtype=np.float64), np.asarray(timestep_map, dtype=int)


class Model(object):
    def __init__(self,
                 beta_strategy,
                 num_diffusion_timesteps,
                 use_timesteps=None,
                 num_classes=1,
                 model_var_type: ModelVarType = ModelVarType.FIXED_SMALL,
                 randflip=True,
                 attention_num_heads=1,
                 attention_resolutions=(16, 8),
                 base_channels=None,
                 channel_mult=None,
                 num_res_blocks=None,
                 scale_shift_norm=True):

        betas = get_beta_schedule(
            beta_strategy, num_timesteps=num_diffusion_timesteps)
        self.original_timesteps = len(betas)
        self.timestep_map = None

        if use_timesteps is not None:
            betas, timestep_map = respace_betas(betas, use_timesteps)
            self.timestep_map = const_var(timestep_map)

        self.model_var_type = model_var_type
        self.diffusion: GaussianDiffusion = GaussianDiffusion(betas=betas,
                                                              model_var_type=model_var_type)
        self.num_classes = num_classes
        self.randflip = randflip
        self.attention_num_heads = attention_num_heads
        self.attention_resolutions = attention_resolutions
        self.base_channels = base_channels
        self.channel_mult = channel_mult
        self.num_res_blocks = num_res_blocks
        self.scale_shift_norm = scale_shift_norm

    def _define_model(self, input_shape, dropout):
        assert isinstance(input_shape, (tuple, list))
        assert len(input_shape) == 4

        B, C, H, W = input_shape

        output_channels = C
        if is_learn_sigma(self.model_var_type):
            output_channels *= 2

        unet = UNet(num_classes=self.num_classes,
                    model_channels=self.base_channels,
                    output_channels=output_channels,
                    num_res_blocks=self.num_res_blocks,
                    attention_resolutions=self.attention_resolutions,
                    attention_num_heads=self.attention_num_heads,
                    channel_mult=self.channel_mult,
                    dropout=dropout,
                    scale_shift_norm=self.scale_shift_norm,
                    conv_resample=True)

        return unet

    def rescale_timestep(self, t):
        if self.timestep_map is None:
            return t

        assert isinstance(self.timestep_map, nn.Variable)

        # revert t to the original timestep
        return self.diffusion._extract(self.timestep_map, t)

    def _denoise(self, x, t, dropout):
        assert x.shape[0] == t.shape[0]

        unet = self._define_model(x.shape, dropout)

        out = unet(x, self.rescale_timestep(t))

        if is_learn_sigma(self.model_var_type):
            B, C, H, W = x.shape
            assert out.shape == (B, 2 * C, H, W)
        else:
            assert out.shape == x.shape
        return out

    def get_denoise_net_intermediates(self, x, t, dropout):
        assert x.shape[0] == t.shape[0]

        unet = self._define_model(x.shape, dropout)

        return unet.get_intermediates(x, self.rescale_timestep(t))

    def build_denoise_graph(self, shape):
        """just create parameters for inference."""
        B, _, _, _ = shape
        x_dummy = nn.Variable(shape)
        t_dummy = nn.Variable((B, ))

        out = self._denoise(x_dummy, t_dummy, dropout=0)

        return out

    def build_train_graph(self, x, t=None, dropout=0, noise=None, loss_scaling=None):
        B, C, H, W = x.shape
        if self.randflip:
            x = F.random_flip(x)
            assert x.shape == (B, C, H, W)

        if t is None:
            t = F.randint(
                low=0, high=self.diffusion.num_timesteps, shape=(B, ))
            # F.randint could return high with very low prob. Workaround to avoid this.
            t = F.clip_by_value(t, min=0, max=self.diffusion.num_timesteps-0.5)

        loss_dict = self.diffusion.train_loss(model=partial(self._denoise, dropout=dropout),
                                              x_start=x, t=t, noise=noise)
        assert isinstance(loss_dict, AttrDict)

        # setup training loss
        loss_dict.batched_loss = loss_dict.mse
        if is_learn_sigma(self.model_var_type):
            assert "vlb" in loss_dict
            loss_dict.batched_loss += loss_dict.vlb * 1e-3
            # todo: implement loss aware sampler

        if loss_scaling is not None and loss_scaling > 1:
            loss_dict.batched_loss *= loss_scaling

        # setup flat training loss
        loss_dict.loss = F.mean(loss_dict.batched_loss)
        assert loss_dict.batched_loss.shape == t.shape == (B, )

        # Keep interval values to compute loss for each quantile
        t.persistent = True
        for v in loss_dict.values():
            v.persistent = True

        return loss_dict, t

    def sample(self, shape, dump_interval=-1, noise=None, use_ema=True, progress=False, use_ddim=False):
        if use_ema:
            with nn.parameter_scope("ema"):
                return self.sample(shape,
                                   dump_interval=dump_interval,
                                   noise=noise,
                                   use_ema=False,
                                   progress=progress,
                                   use_ddim=use_ddim)

        loop_func = self.diffusion.ddim_sample_loop if use_ddim else self.diffusion.p_sample_loop

        with nn.no_grad():
            return loop_func(
                model=partial(self._denoise, dropout=0),
                shape=shape,
                noise=noise,
                dump_interval=dump_interval,
                progress=progress
            )

    def sample_trajectory(self, shape, noise=None, use_ema=True, progress=False, use_ddim=False):
        return self.sample(shape,
                           dump_interval=100,
                           noise=noise,
                           use_ema=use_ema,
                           progress=progress,
                           use_ddim=use_ddim)


# 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 os
import sys

import fnmatch

import click
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 AttrDict, get_current_time, init_nnabla
from neu.reporter import KVReporter, save_tiled_image
from neu.yaml_wrapper import write_yaml

from dataset import get_dataset
from model import Model
from utils import get_warmup_lr, sum_grad_norm, create_ema_op
from config import refine_args_by_dataset
from diffusion import ModelVarType, is_learn_sigma


def get_output_dir_name(org, dataset):
    return os.path.join(org, f"{get_current_time()}_{dataset}")


def str_as_integer_list(ctx, param, value):
    if value is None or len(value) == 0:
        return None

    return tuple(int(x) for x in value.split(","))


@click.command()
# configs for training process
@click.option("--accum", default=1, help="# of gradient accumulation.", show_default=True)
@click.option("--device-id", default='0', help="Device id.", show_default=True)
@click.option("--batch-size", default=4, help="Batch size to train.", show_default=True)
@click.option("--n-iters", default=int(5e5), help="# of training iterations.", show_default=True)
@click.option("--progress/--no-progress", default=False, help="Use tqdm to show progress.")
@click.option("--resume/--no-resume", default=False, help="Resume training from the latest params saved at the same output_dir.")
@click.option("--loss-scaling", default=1.0, type=float, help="Loss scaling factor.", show_default=True)
@click.option("--lr", default=None, type=float, help="Learning rate.")
# model configs
@click.option("--beta-strategy", default="linear", help="Strategy to create betas.", show_default=True,
              type=click.Choice(["linear", "cosine"], case_sensitive=False), show_choices=True)
@click.option("--num-diffusion-timesteps", default=1000, help="Number of diffusion timesteps.", show_default=True)
@click.option("--ssn/--no-ssn", default=True, type=bool, help="use scale shift norm or not.")
@click.option("--num-attention-heads", default=4, type=int, help="Number of multihead attention heads", show_default=True)
@click.option("--attention-resolutions", default="16,8", type=str, callback=str_as_integer_list,
              help="Resolutions which attention is applied. Comma separated string should be passed. If None, use default for dataset.")
@click.option("--base-channels", default=None, type=int, help="Base channel size. If None, use default for dataset.")
@click.option("--channel-mult", default=None, type=str, callback=str_as_integer_list,
              help="Channel multipliers for each block. Comma separated string should be passed. If None, use default for dataset.")
@click.option("--num-res-blocks", default=None, type=int, help="# of residual blocks. If None, use default for dataset.")
@click.option("--dropout", default=0.0, type=float, help="Dropout prob.", show_default=True)
@click.option("--model-var-type", type=click.Choice(ModelVarType.get_supported_keys(), case_sensitive=False),
              default="fixed_small", help="A type of the model variance.", show_default=True, show_choices=True)
# data related configs
@click.option("--dataset", default="custum", help="Dataset name to train model on.",
              type=click.Choice(["celebahq", "cifar10", "imagenet", "custom"], case_sensitive=False), show_choices=True)
@click.option("--image-size", default=None,
              type=int, help="Image size. Should be a integer and used for both height and width.")
@click.option("--data-dir", default="./data", help="The path for data directory.", show_default=True)
@click.option("--dataset-root-dir", default=None,
              help="The path for dataset root directory.", show_default=True)
@click.option("--dataset-on-memory/--no-dataset-on-memory", default=False,
              help="If True, the data once loaded will be kept on Host memory.", show_default=True)
@click.option("--fix-aspect-ratio/--no-fix-aspect-ratio", default=True,
              help="Whether keep aspect ratio or not for loaded training images.", show_default=True)
# configs for dumpling
@click.option("--output-dir", default="./logdir", help="output dir", show_default=True)
@click.option("--save-interval", default=int(1e4), help="Number of iters between saves.", show_default=True)
@click.option("--gen-interval", default=int(1e4), help="Number of iters between each generation.", show_default=True)
@click.option("--show-interval", default=10, help="Number of iters between showing current logging values.", show_default=True)
@click.option("--dump-grad-norm/--no-dump-grap-norm",
              default=False, help="Show sum of gradient norm of all params for each iteration.")
def main(**kwargs):
    # set training args
    args = AttrDict(kwargs)
    refine_args_by_dataset(args)

    args.output_dir = get_output_dir_name(args.output_dir, args.dataset)

    comm = init_nnabla(ext_name="cudnn", device_id=args.device_id,
                       type_config="float", random_pseed=True)

    data_iterator = get_dataset(args, comm)

    model = Model(beta_strategy=args.beta_strategy,
                  num_diffusion_timesteps=args.num_diffusion_timesteps,
                  model_var_type=ModelVarType.get_vartype_from_key(
                      args.model_var_type),
                  attention_num_heads=args.num_attention_heads,
                  attention_resolutions=args.attention_resolutions,
                  scale_shift_norm=args.ssn,
                  base_channels=args.base_channels,
                  channel_mult=args.channel_mult,
                  num_res_blocks=args.num_res_blocks)

    # build graph
    x = nn.Variable(args.image_shape)  # assume data_iterator returns [0, 255]
    x_rescaled = x / 127.5 - 1  # rescale to [-1, 1]
    loss_dict, t = model.build_train_graph(x_rescaled,
                                           dropout=args.dropout,
                                           loss_scaling=None if args.loss_scaling == 1.0 else args.loss_scaling)
    assert loss_dict.batched_loss.shape == (args.batch_size, )
    assert t.shape == (args.batch_size, )
    assert t.persistent == True

    # optimizer
    solver = S.Adam()
    solver.set_parameters(nn.get_parameters())

    # 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 args.resume:
        parent = os.path.dirname(os.path.abspath(args.output_dir))
        all_logs = sorted(fnmatch.filter(
            os.listdir(parent), "*{}*".format(args.dataset)))
        if len(all_logs):
            latest_dir = os.path.join(parent, all_logs[-1])
            checkpoints = sorted(fnmatch.filter(
                os.listdir(latest_dir), "checkpoint_*.json"))
            if len(checkpoints):
                latest_cp = os.path.join(latest_dir, checkpoints[-1])
                start_iter = load_checkpoint(latest_cp, solvers)

                for sname, slv in solvers.items():
                    slv.zero_grad()
    comm.barrier()

    # Reporter
    reporter = KVReporter(comm, save_path=args.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(model.model_var_type):
            reporter.set_key(f"vlb_q{i}")

    image_dir = os.path.join(args.output_dir, "image")
    if comm.rank == 0:
        os.makedirs(image_dir, exist_ok=True)

    if args.progress:
        from tqdm import trange
        piter = trange(start_iter + 1, args.n_iters + 1,
                       disable=comm.rank > 0, ncols=0)
    else:
        piter = range(start_iter + 1, args.n_iters + 1)

    # dump config
    if comm.rank == 0:
        args.dump()
        write_yaml(os.path.join(args.output_dir, "config.yaml"), args)

    comm.barrier()

    for i in piter:
        # update solver's lr
        # cur_lr = get_warmup_lr(lr, args.n_warmup, i)
        solver.set_learning_rate(args.lr)

        # evaluate graph
        dummy_solver_ema.zero_grad()  # just in case
        solver.zero_grad()
        for accum_iter in range(args.accum):  # accumelate
            data, label = data_iterator.next()
            x.d = data.copy()

            loss_dict.loss.forward(clear_no_need_grad=True)

            all_reduce_cb = None
            if accum_iter == args.accum - 1:
                all_reduce_cb = comm.get_all_reduce_callback(
                    params=solver.get_parameters().values())

            loss_dict.loss.backward(
                clear_buffer=True, communicator_callbacks=all_reduce_cb)

            # logging
            # loss
            reporter.kv_mean("loss", loss_dict.loss)

            if is_learn_sigma(model.model_var_type):
                reporter.kv_mean("vlb", loss_dict.vlb)

            # loss for each quantile
            for j in range(args.batch_size):
                ti = t.d[j]
                q_level = int(ti) * 4 // args.num_diffusion_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(model.model_var_type):
                    reporter.kv_mean(f"vlb_q{q_level}", loss_dict.vlb.d[j])

        # update
        if args.grad_clip > 0:
            solver.clip_grad_by_norm(args.grad_clip)
        solver.update()

        # update ema params
        ema_op.forward(clear_no_need_grad=True)

        # grad norm
        if args.dump_grad_norm:
            gnorm = sum_grad_norm(solver.get_parameters().values())
            reporter.kv_mean("grad", gnorm)

        # samples
        reporter.kv("samples", i * args.batch_size * comm.n_procs * args.accum)

        # iteration (only for no-progress)
        if not args.progress:
            reporter.kv("iteration", i)

        if i % args.show_interval == 0:
            if args.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=True)

            reporter.flush_monitor(i)

        if i > 0 and i % args.save_interval == 0:
            if comm.rank == 0:
                save_checkpoint(args.output_dir, i, solvers, n_keeps=3)

            comm.barrier()

        if i > 0 and i % args.gen_interval == 0:
            # sampling
            sample_out, _, _ = model.sample(
                shape=(16, ) + x.shape[1:], use_ema=True, progress=False)
            assert sample_out.shape == (16, ) + args.image_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)


if __name__ == "__main__":
    main()


# diffusion-models/unet.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 click

import nnabla as nn
import nnabla.functions as F

from layers import *


class ResidualBlock(object):
    def __init__(self,
                 out_channels=None,
                 scale_shift_norm=False,
                 dropout=0,
                 conv_shortcut=False):
        self.out_channels = out_channels
        self.scale_shift_norm = scale_shift_norm
        self.dropout = dropout
        self.conv_shortcut = conv_shortcut

    def in_layers(self, x):
        h = normalize(x, name='norm_in')
        h = nonlinearity(h)
        h = conv(h, self.out_channels, name='conv_in')

        return h

    def emb_layers(self, emb):
        out_channels = self.out_channels
        if self.scale_shift_norm:
            out_channels *= 2

        return nin(emb, out_channels, name="emb_proj")

    def out_layers(self, h, emb):
        if self.scale_shift_norm:
            scale, shift = chunk(emb, num_chunk=2, axis=1)
            h = normalize(h, name="norm_out") * (scale + 1) + shift
        else:
            h += emb
            h = normalize(h, name="norm_out")

        h = nonlinearity(h)

        if self.dropout > 0:
            h = F.dropout(h, p=self.dropout)

        h = conv(h, self.out_channels, name="conv_out", zeroing_w=True)

        return h

    def shortcut(self, x):
        if self.out_channels == x.shape[1]:
            return x
        elif self.conv_shortcut:
            return conv(x, self.out_channels, name="conv_shortcut")
        else:
            return nin(x, self.out_channels, name="conv_shortcut")

    def __call__(self, x, temb, name):
        C = x.shape[1]
        if self.out_channels is None:
            self.out_channels = C

        with nn.parameter_scope(name):
            # first block
            h = self.in_layers(x)

            # embed
            emb = self.emb_layers(temb)

            # second block
            h = self.out_layers(h, emb)

            # add residual
            out = F.add2(h, self.shortcut(x), inplace=True)

        return out


def attn_block(x, name, num_heads=4, fix_parameters=False):
    """Multihead attention block"""
    B, C, H, W = x.shape

    with nn.parameter_scope(name):
        # Get query, key, value
        h = normalize(x, name="norm")
        # nin(3 * C) -> split is faster?
        q = nin(h, C, name="q")
        k = nin(h, C, name="k")
        v = nin(h, C, name="v")

        # Attention
        w = F.batch_matmul(F.reshape(q, (B * num_heads, -1, H * W)),
                           F.reshape(k, (B * num_heads, -1, H * W)), transpose_a=True)
        w = F.mul_scalar(w, int(C) ** (-0.5), inplace=True)

        assert w.shape == (B * num_heads, H * W, H * W)
        w = F.softmax(w, axis=-1)

        h = F.reshape(v, (B * num_heads, -1, H * W))
        h = F.batch_matmul(h, w)
        h = F.reshape(h, (B, C, H, W))

        # output projection
        h = nin(h, C, name='proj_out', zeroing_w=True)

    assert h.shape == x.shape
    return F.add2(h, x, inplace=True)


class UNet(object):
    def __init__(self,
                 num_classes,
                 model_channels,
                 output_channels,
                 num_res_blocks,
                 attention_resolutions,
                 attention_num_heads,
                 channel_mult=(1, 2, 4, 8),
                 dropout=0.,
                 scale_shift_norm=False,
                 conv_resample=True):
        self.num_classes = num_classes
        self.model_channels = model_channels
        self.output_channels = output_channels
        self.num_res_blocks = num_res_blocks
        self.attention_resolutions = attention_resolutions
        self.attention_num_heads = attention_num_heads
        self.channel_mult = channel_mult
        self.dropout = dropout
        self.scale_shift_norm = scale_shift_norm
        self.conv_resample = conv_resample

    def timestep_embedding(self, t):
        with nn.parameter_scope('timestep_embedding'):
            # sinusoidal embedding
            emb = sinusoidal_embedding(t, self.model_channels)

            # reshape to use conv rather than affine
            emb = F.reshape(emb, emb.shape + (1, 1))

            # linear transforms
            emb = nin(emb, self.model_channels * 4, name='dense0')
            emb = nonlinearity(emb)

            emb = nin(emb, self.model_channels * 4, name='dense1')

        return emb

    def resblock_with_attention(self, h, emb, out_channels, name):
        with nn.parameter_scope(name):
            block = ResidualBlock(out_channels,
                                  scale_shift_norm=self.scale_shift_norm,
                                  dropout=self.dropout)

            h = block(h, emb, f"res_block")

            res = h.shape[-1]
            if self.attention_resolutions is not None and res in self.attention_resolutions:
                h = attn_block(h, f"attention",
                               num_heads=self.attention_num_heads)

        return h

    def downsample_blocks(self, h, emb, out_channels, name):
        hs = []
        with nn.parameter_scope(name):
            for i in range(self.num_res_blocks):
                h = self.resblock_with_attention(
                    h, emb, out_channels, name=f"resblock_{i}")
                hs.append(h)

        return hs

    def upsample_blocks(self, h, emb, hs_down, out_channels, name):
        hs_up = []
        with nn.parameter_scope(name):
            for i in range(self.num_res_blocks + 1):
                # concat skip
                h = F.concatenate(h, hs_down.pop(), axis=1)
                h = self.resblock_with_attention(
                    h, emb, out_channels, name=f"resblock_{i}")
                hs_up.append(h)

        return hs_up

    def middle_block(self, h, emb):
        ch = h.shape[1]
        block = ResidualBlock(
            ch, scale_shift_norm=self.scale_shift_norm, dropout=self.dropout)

        h = block(h, emb, name="resblock_0")

        res = h.shape[-1]
        if self.attention_resolutions is not None and res in self.attention_resolutions:
            h = attn_block(h, f"attention", num_heads=self.attention_num_heads)

        h = block(h, emb, name="resblock_1")

        return h

    def output_block(self, h):
        h = normalize(h, "last_norm")
        h = nonlinearity(h)
        h = conv(h, self.output_channels, name="last_conv", zeroing_w=True)

        return h

    def get_intermediates(self, x, t, name=None):
        ret = dict()

        with nn.auto_forward(True):
            ch = self.model_channels
            with nn.parameter_scope('UNet' if name is None else name):
                h = conv(x, ch, name="first_conv")

                emb = self.timestep_embedding(t)
                ret["emb"] = emb
                emb = nonlinearity(emb)

                hs = [h]
                # downsample block
                with nn.parameter_scope("downsample_block"):
                    for level, mult in enumerate(self.channel_mult):
                        # apply resblock and attention for this resolution
                        outs = self.downsample_blocks(h, emb, ch * mult,
                                                      name=f"block_{level}")
                        hs += outs
                        h = outs[-1]

                        # downsample to lower resolution except last
                        if level < len(self.channel_mult) - 1:
                            h = downsample(h, name=f"downsample_{level}",
                                           with_conv=self.conv_resample)
                            hs.append(h)

                ret["down"] = hs.copy()

                # middle block
                with nn.parameter_scope("middle_block"):
                    h = self.middle_block(h, emb)

                ret["middle"] = h

                # upsample block
                hs_up = []
                with nn.parameter_scope("upsample_block"):
                    for level, mult in enumerate(reversed(self.channel_mult)):
                        # apply resblock and attention for this resolution
                        outs = self.upsample_blocks(h, emb, hs, ch * mult,
                                                    name=f"output_{level}")
                        h = outs[-1]

                        # downsample to lower resolution except last
                        if level < len(self.channel_mult) - 1:
                            h = upsample(h, name=f"upsample_{level}",
                                         with_conv=self.conv_resample)
                            outs.pop()
                            outs.append(h)

                        hs_up += outs
                assert len(hs) == 0

                ret["up"] = hs_up.copy()

                # output block
                with nn.parameter_scope("output_block"):
                    out = self.output_block(h)

                ret["out"] = out

                assert out.shape == x.shape[:1] + \
                    (self.output_channels, ) + x.shape[2:]

                return ret

    def __call__(self, x, t, name=None):
        ch = self.model_channels
        with nn.parameter_scope('UNet' if name is None else name):
            h = conv(x, ch, name="first_conv")
            emb = self.timestep_embedding(t)
            emb = nonlinearity(emb)

            hs = [h]
            # downsample block
            with nn.parameter_scope("downsample_block"):
                for level, mult in enumerate(self.channel_mult):
                    # apply resblock and attention for this resolution
                    outs = self.downsample_blocks(h, emb, ch * mult,
                                                  name=f"block_{level}")
                    hs += outs
                    h = outs[-1]

                    # downsample to lower resolution except last
                    if level < len(self.channel_mult) - 1:
                        h = downsample(h, name=f"downsample_{level}",
                                       with_conv=self.conv_resample)
                        hs.append(h)

            # middle block
            with nn.parameter_scope("middle_block"):
                h = self.middle_block(h, emb)

            # upsample block
            with nn.parameter_scope("upsample_block"):
                for level, mult in enumerate(reversed(self.channel_mult)):
                    # apply resblock and attention for this resolution
                    outs = self.upsample_blocks(h, emb, hs, ch * mult,
                                                name=f"output_{level}")
                    h = outs[-1]

                    # downsample to lower resolution except last
                    if level < len(self.channel_mult) - 1:
                        h = upsample(h, name=f"upsample_{level}",
                                     with_conv=self.conv_resample)

            assert len(hs) == 0

            # output block
            with nn.parameter_scope("output_block"):
                out = self.output_block(h)

            assert out.shape == x.shape[:1] + \
                (self.output_channels, ) + x.shape[2:]

            return out

# Functions below are for dubugging UNet class.


def test_simple_loop():
    nn.clear_parameters()

    x = nn.Variable.from_numpy_array(np.random.randn(10, 3, 128, 128))
    t = nn.Variable.from_numpy_array(np.random.randint(0, 100, (10, )))

    unet = UNet(num_classes=1, model_channels=128, output_channels=3,
                num_res_blocks=2,
                attention_resolutions=(16, 8),
                attention_num_heads=4,
                channel_mult=(1, 1, 2, 2, 4, 4))
    y = unet(x, t)

    loss = F.mean(F.squared_error(y, x))

    import nnabla.solvers as S
    solver = S.Sgd()
    solver.set_parameters(nn.get_parameters())

    from tqdm import trange
    tr = trange(100)
    for i in tr:
        loss.forward(clear_no_need_grad=True)
        solver.zero_grad()
        loss.backward(clear_buffer=True)
        solver.update()

        tr.set_description(f"diff: {loss.d.copy():.5f}")


def test_intermediate():
    import os

    os.environ["NNABLA_CUDNN_DETERMINISTIC"] = '1'

    nn.clear_parameters()

    x = nn.Variable.from_numpy_array(np.full((1, 3, 256, 256), 0.1))
    t = nn.Variable.from_numpy_array([803])

    unet = UNet(num_classes=1, model_channels=128, output_channels=3,
                num_res_blocks=3,
                attention_resolutions=(16, 8),
                attention_num_heads=4,
                channel_mult=(1, 1, 2, 2, 4, 4))
    res = unet.get_intermediates(x, t)

    print("[emb]")
    dump(res["emb"])
    print("")

    print("[down]")
    dump(res["down"])
    print("")

    print("[middle]")
    dump(res["middle"])
    print("")

    print("[up]")
    dump(res["up"])
    print("")

    print("[out]")
    dump(res["out"])
    print("")


def dump(var):
    if isinstance(var, list):
        for x in var:
            dump(x)

        return

    arr = var.d
    mean = arr.mean()
    std = arr.std()
    abs_sum = np.abs(arr).sum()
    print("mean: {:-6.1g} std: {:-6.1g} abs_sum: {:-6.1g} size: {}".format(mean,
          std, abs_sum, arr.size))


@click.command()
@click.option("--loop/--no-loop", default=True)
@click.option("--intermediate/--no-intermediate", default=False)
def test(loop, intermediate):
    # This function is for a unit test of UNet.
    from nnabla.ext_utils import get_extension_context
    ctx = get_extension_context("cudnn")
    nn.set_default_context(ctx)

    from nnabla.logger import logger

    if loop:
        logger.info("Test Unet by simple training loop.")
        test_simple_loop()

    if intermediate:
        logger.info("Test intermediate values of Unet.")
        test_intermediate()


if __name__ == "__main__":
    test()


# diffusion-models/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 nnabla as nn
import nnabla.functions as F
import numpy as np
from nnabla.parameter import get_parameter_or_create


def get_warmup_lr(max_lr, warmup_epoch, step):
    # linear warmup
    if warmup_epoch == 0 or step >= warmup_epoch:
        return max_lr

    return max_lr * step / warmup_epoch


def sum_grad_norm(params):
    norm = nn.NdArray()
    norm.zero()

    for p in params:
        assert isinstance(p, nn.Variable) and not p.grad.clear_called
        norm += F.sum(p.grad ** 2)

    return np.sqrt(norm.data)


def create_ema_op(params, ema_decay=0.9999):
    """
    Define exponential moving average update for trainable params.
    """
    def ema_update(p_ema, p_train):
        return F.assign(p_ema, ema_decay * p_ema + (1. - ema_decay) * p_train)

    ops = []
    with nn.parameter_scope("ema"):
        for name, p_train in params.items():
            p_ema = get_parameter_or_create(
                name, shape=p_train.shape, need_grad=False)
            p_ema.data.copy_from(
                p_train.data, use_current_context=False)  # initialize
            ops.append(ema_update(p_ema, p_train))

        ema_params = nn.get_parameters(grad_only=False)

    return F.sink(*ops), ema_params


# 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/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/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/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_and_get_test_results(model, x_train, y_train, x_test, y_test, **kwargs):
    """
    Train the model on given training data and evaluate its performance on test set

    Args:
        model: various ML models
        x_train: Training data
        y_train: Training targets
        x_test: Testing data
        y_test: Testing targets
        kwargs: dictionary that maps each keyword to the value that we pass alongside it

    Returns:
        Accuracy and predictions made by the specified model
    """
    # Instantiate the classification model
    if "sample_weight" in kwargs:
        model.fit(x_train, y_train.ravel(),
                  sample_weight=kwargs["sample_weight"])
    else:
        model.fit(x_train, y_train)

    predicted = model.predict(x_test)

    # compute the accuracy of the model
    accuracy = model.score(x_test, y_test.ravel())

    return accuracy, predicted


def visualize_model_comparison(ml_models, base_fairness, mitigated_fairness, base_accuracy,
                               mitigated_accuracy):
    """
    Graphical Comparison of fairness & performance

    Args:
        ml_models: list the name of the ML models
        base_fairness: list the fairness values of the ML models, before bias mitigation
        mitigated_fairness: list the fairness values of the ML models, after bias mitigation
        base_accuracy: list the accuracy values of the ML models, before bias mitigation
        mitigated_accuracy: list the fairness values of the ML models, after bias mitigation

    """

    def autolabel(rects, axes):
        for rect in rects:
            h = rect.get_height()
            if h < 0:
                axes.text(rect.get_x() + rect.get_width() / 2., h - 0.04, '%.2f' % float(h),
                          ha='center', va='bottom')
            else:
                axes.text(rect.get_x() + rect.get_width() / 2., 1.02 * h, '%.2f' % float(h),
                          ha='center', va='bottom')

    ind = np.arange(len(ml_models))  # the x locations for the groups
    width = 0.27  # the width of the bars
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))

    fairness_base_bar = ax1.bar(ind, base_fairness, width, color='r',
                                label='Before bias mitigation')
    fairness_mitigated_bar = ax1.bar(ind + width, mitigated_fairness, width, color='g',
                                     label='After bias mitigation')
    ax1.axhline(y=0.0, color='r', linestyle='-')
    ax1.set_xlabel("ML models")
    ax1.set_ylabel('Fairness(SPD)')
    ax1.set_ylim(-0.5, 0.5)
    ax1.set_xticks(ind + width)
    ax1.set_xticklabels(ml_models)
    ax1.title.set_text('(a) Fairness of ML model')
    ax1.legend()
    autolabel(fairness_base_bar, axes=ax1)
    autolabel(fairness_mitigated_bar, axes=ax1)

    base_accuracy_bar = ax2.bar(ind, base_accuracy, width, color='r',
                                label='Before bias mitigation')
    mitigated_accuracy_bar = ax2.bar(ind + width, mitigated_accuracy, width, color='g',
                                     label='After bias mitigation')
    ax2.set_xlabel("ML models")
    ax2.set_ylabel('Accuracy')
    ax2.set_ylim(0.0, 1.0)
    ax2.set_xticks(ind + width)
    ax2.set_xticklabels(ml_models)
    ax2.title.set_text('(b) Performance of ML model')
    ax2.legend()
    autolabel(base_accuracy_bar, axes=ax2)
    autolabel(mitigated_accuracy_bar, axes=ax2)
    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/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/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/three_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))

    # 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_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/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

from setuptools import setup

setup(
    name='neu',
    version='1.23.0dev',  # 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/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
