# Licensed under a 3-clause BSD style license - see LICENSE.rst
import numpy as np
from gammapy.datasets import SpectrumDatasetOnOff

__all__ = ["StandardOGIPDataset"]

class StandardOGIPDataset(SpectrumDatasetOnOff):
    """Dataset containing spectral data as defined by X-ray OGIP compliant files.

    A few elements are added that are not supported by the current SpectrumDataset in gammapy.

    grouping contains the information on the grouping scheme, namely the group number of each bin.

    Parameters
    ----------
    models : `~gammapy.modeling.models.Models`
        Source sky models.
    counts : `~gammapy.maps.WcsNDMap`
        Counts cube
    counts_off : `~gammapy.maps.WcsNDMap`
        Ring-convolved counts cube
    acceptance : `~gammapy.maps.WcsNDMap`
        Acceptance from the IRFs
    acceptance_off : `~gammapy.maps.WcsNDMap`
        Acceptance off
    exposure : `~gammapy.maps.WcsNDMap`
        Exposure cube
    mask_fit : `~gammapy.maps.WcsNDMap`
        Mask to apply to the likelihood for fitting.
    psf : `~gammapy.irf.PSFKernel`
        PSF kernel. Unused here.
    edisp : `~gammapy.irf.EDispKernel`
        Energy dispersion
    mask_safe : `~gammapy.maps.WcsNDMap`
        Mask defining the safe data range.
    grouping_axis : `~gammapy.maps.MapAxis`
        MapAxis defining the grouping scheme.
    gti : `~gammapy.data.GTI`
        GTI of the observation or union of GTI if it is a stacked observation
    meta_table : `~astropy.table.Table`
        Table listing information on observations used to create the dataset.
        One line per observation for stacked datasets.
    name : str
        Name of the dataset.

    """

    stat_type = "wstat"
    tag = "StandardOGIPDataset"

    def __init__(self, *args, **kwargs):
        self._grouped = None
        axis = kwargs.pop("grouping_axis", None)
        super().__init__(*args, **kwargs)
        self.grouping_axis = axis
        self.apply_grouping(self.grouping_axis)

    @property
    def grouped(self):
        """Return grouped SpectrumDatasetOnOff."""
        return self._grouped

    @property
    def _is_grouped(self):
        return self.grouped is not None

    def apply_grouping(self, axis=None):
        """Apply grouping."""
        if axis is None:
            raise ValueError("A grouping MapAxis must be provided.")
        else:
            dataset = self.to_spectrum_dataset_onoff()
            self._grouped = dataset.resample_energy_axis(
                axis, name=f"group_{self.name}"
            )

    @property
    def models(self):
        """Models (`~gammapy.modeling.models.Models`)."""
        return self.grouped._models

    @models.setter
    def models(self, models):
        """Models setter"""
        if self._is_grouped:
            self.grouped.models = models

    @property
    def mask_fit(self):
        """RegionNDMap providing the fitting energy range."""
        if self._is_grouped:
            return self.grouped.mask_fit

    @mask_fit.setter
    def mask_fit(self, mask_fit):
        """RegionNDMap providing the fitting energy range."""
        if self._is_grouped:
            self.grouped.mask_fit = mask_fit.resample_axis(
                axis=self.grouping_axis, ufunc=np.logical_or
            )

    @property
    def mask(self):
        """Combined fit and safe mask"""
        return self.grouped.mask

    def npred(self):
        """Predicted source and background counts
        Returns
        -------
        npred : `Map`
            Total predicted counts
        """
        return self.grouped.npred()

    def npred_signal(self, model_name=None):
        """ "Model predicted signal counts.
        If a model is passed, predicted counts from that component is returned.
        Else, the total signal counts are returned.
        Parameters
        -------------
        model_name: str
            Name of  SkyModel for which to compute the npred for.
            If none, the sum of all components (minus the background model)
            is returned
        Returns
        ----------
        npred_sig: `gammapy.maps.Map`
            Map of the predicted signal counts
        """
        return self.grouped.npred_signal(model_name=model_name)

    def stat_sum(self):
        """Total statistic given the current model parameters."""
        return self.grouped.stat_sum()

    def plot_fit(
        self,
        ax_spectrum=None,
        ax_residuals=None,
        kwargs_spectrum=None,
        kwargs_residuals=None,
    ):
        ax_spectrum, ax_residuals = self.grouped.plot_fit(
            ax_spectrum, ax_residuals, kwargs_spectrum, kwargs_residuals
        )

        return ax_spectrum, ax_residuals

    def plot_residuals_spectral(self, ax=None, method="diff", region=None, **kwargs):
        """Plot spectral residuals.
        The residuals are extracted from the provided region, and the normalization
        used for its computation can be controlled using the method parameter.
        Parameters
        ----------
        ax : `~matplotlib.axes.Axes`
            Axes to plot on.
        method : {"diff", "diff/sqrt(model)"}
            Normalization used to compute the residuals, see `SpectrumDataset.residuals`.
        region: `~regions.SkyRegion` (required)
            Target sky region.
        **kwargs : dict
            Keyword arguments passed to `~matplotlib.axes.Axes.errorbar`.
        Returns
        -------
        ax : `~matplotlib.axes.Axes`
            Axes object.
        """
        return self.grouped.plot_residuals_spectral(
            ax=ax, method=method, region=region, **kwargs
        )

    @classmethod
    def read(cls, filename):
        """Read from file

        For now, filename is assumed to the name of a PHA file where BKG file, ARF, and RMF names
        must be set in the PHA header and be present in the same folder.

        For formats specs see `OGIPDatasetReader.read`

        Parameters
        ----------
        filename : `~pathlib.Path` or str
            OGIP PHA file to read
        """
        from .io_ogip import StandardOGIPDatasetReader

        reader = StandardOGIPDatasetReader(filename=filename)
        return reader.read()

    def write(self, filename, overwrite=False, format="ogip"):
        raise NotImplementedError("Standard OGIP writing is not supported.")

    def to_spectrum_dataset_onoff(self, name=None):
        """convert to spectrum dataset on off by dropping the grouping axis.
        Parameters
        ----------
        name : str
            Name of the new dataset.

        Returns
        -------
        dataset : `SpectrumDatasetOnOff`
            Spectrum dataset on off.
        """
        kwargs = {"name": name}

        kwargs["acceptance"] = self.acceptance
        kwargs["acceptance_off"] = self.acceptance_off
        kwargs["counts_off"] = self.counts_off
        dataset = self.to_spectrum_dataset()
        return SpectrumDatasetOnOff.from_spectrum_dataset(dataset=dataset, **kwargs)


#!/usr/bin/env python3
import sys
import numpy as np
import os, cv2
import argparse
import rospy

import torch
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from lib.mask_rcnn import MaskRCNNPredictor, MaskAffordancePredictor, MaskRCNNHeads
import lib.mask_rcnn as mask_rcnn
import utils
import argparse
from PIL import Image
import numpy as np
import cv2
import time
from config import IITAFF
import os
from skimage.transform import resize
from affordanceService.client import AffordanceClient

from affordance_analyzer.srv import *
from std_msgs.msg import MultiArrayDimension, String

class AffordanceAnalyzer(object):
    """docstring for AffordanceAnalyzer."""

    def __init__(self):

        self.CONF_THRESHOLD = 0.7
        self.root_path = '/workspace/affordanceNet_sytn/'
        self.name = "affordancenet_synth"
        self.device = None

        self.serviceGet = rospy.Service('/affordance/result', getAffordanceSrv, self.getAffordance)
        self.serviceRun = rospy.Service('/affordance/run', runAffordanceSrv, self.analyzeAffordance)
        self.serviceStart = rospy.Service('/affordance/start', startAffordanceSrv, self.startAffordance)
        self.serviceStop = rospy.Service('/affordance/stop', stopAffordanceSrv, self.stopAffordance)
        self.serviceName = rospy.Service('/affordance/name', getNameSrv, self.getName)

        self.net = None

    def getName(self, msg):

        response = getNameSrvResponse()
        response.name = String(self.name)

        return response

    def run_net(self, x, CONF_THRESHOLD = 0.7):

        #ori_height, ori_width, _ = im.shape

        # Detect all object classes and regress object bounds
        #timer = Timer()
        #timer.tic()
        ts = time.time() * 1000

        predictions = self.net(x)[0]
        boxes, labels, scores, masks = predictions['boxes'], predictions['labels'], predictions['scores'], predictions['masks']
        print("Found: ")

        for label, score in zip(labels, scores):
            print(label, score)
        te = time.time() * 1000
        print("Prediction took: ", te - ts, " ms")
        try:
            idx = scores > CONF_THRESHOLD
            labels = labels.cpu().detach().numpy()

            boxes = boxes[idx].cpu().detach().numpy()
            labels = labels[idx.cpu().detach().numpy()]
            scores = scores[idx].cpu().detach().numpy()
            masks = masks[idx].cpu().detach().numpy()

            masks = masks * 255

        except:
            pass

        #timer.toc()
        #print ('Detection took {:.3f}s for '
        #       '{:d} object proposals').format(timer.total_time, rois_final.shape[0])

        try:
            return boxes, labels, masks, scores
        except:
            return 0

    def sendResults(self, bbox, objects, masks, scores):
        intToLabel = {0: 'class', 1: 'height', 2: 'width'}
        msg = getAffordanceSrvResponse()

        aff_client = AffordanceClient(connected = False)

        # constructing mask message
        """
        for i in range(3):
            dimMsg = MultiArrayDimension()
            dimMsg.label = intToLabel[i]
            stride = 1
            for j in range(3-i):
                stride = stride * masks.shape[i+j]
            dimMsg.stride = stride
            dimMsg.size = masks.shape[i]
            msg.masks.layout.dim.append(dimMsg)
        masks = masks.flatten().astype(int).tolist()
        #msg.masks.data = masks
        """
        msg.masks = aff_client.packMasks(masks)

        # constructing bounding box message
        msg.bbox.data = bbox.flatten().astype(int).tolist()

        # constructing object detection class message
        msg.object.data = objects.flatten().tolist()

        # constructing the scores message
        msg.confidence.data = scores.flatten().tolist()

        return msg

    def analyzeAffordance(self, msg):

        img = np.frombuffer(msg.img.data, dtype=np.uint8).reshape(msg.img.height, msg.img.width, -1)
        width, height = img.shape[1], img.shape[0]
        ratio = width / height
        img = cv2.resize(img, (int(450 * ratio), 450), interpolation = cv2.INTER_AREA)
        x = [torchvision.transforms.ToTensor()(img).to(self.device)]
        self.CONF_THRESHOLD = msg.confidence_threshold.data

        print("Analyzing affordance with confidence threshold: ", self.CONF_THRESHOLD)
        try:
            bbox, objects, masks, scores = self.run_net(x, CONF_THRESHOLD=self.CONF_THRESHOLD)
            #bbox = bbox[:,1:]
            print(masks.shape, bbox.shape, objects.shape, scores.shape)
            try:
                m = np.zeros((11, height, width))
                for c, m_t in enumerate(masks[0]):
                    m[c] = resize(m_t, (height, width))
            except Exception as e:
                print(e)
                raise
            #m = resize(m, (width, height))
            #m = cv2.resize(m, (width, height), interpolation=cv2.INTER_NEAREST)
            try:
                for c, mask in enumerate(masks):
                    if c > 0:
                        m_aff = np.zeros((11, height, width))
                        for c_a, aff_mask in enumerate(mask):
                            m_aff[c_a] = resize(aff_mask, (height, width))
                        m = np.vstack((m,m_aff))
            except Exception as e:
                print(e)
                raise
            masks = m

            for b_c, box in enumerate(bbox):
                box[0] = box[0] * (width / (450 * ratio))
                box[2] = box[2] * (width / (450 * ratio))
                box[1] = box[1] * (height / 450)
                box[3] = box[3] * (height / 450)
                bbox[b_c] = box

            self.bbox = bbox
            self.objects = objects
            self.masks = masks
            self.scores = scores
        except:
            bbox = np.zeros((1, 4))
            objects = np.zeros((1,1))
            masks = np.zeros((1, 11, 244, 244))
            scores = np.zeros(1)

        return runAffordanceSrvResponse()

    def getAffordance(self, msg):

        return self.sendResults(self.bbox, self.objects, self.masks, self.scores)

    def startAffordance(self, msg):
        #weights_path = os.path.dirname(os.path.realpath(__file__)) + "/14.pth"
        weights_path = os.path.dirname(os.path.realpath(__file__)) + "/weights.pth"

        GPU = msg.GPU.data
        device = torch.device('cpu')

        if GPU:
            device = torch.device('cuda')
        self.device = device
        print("Device is: ", self.device)

        model = utils.get_model_instance_segmentation(23, 11)
        #model = utils.get_model_instance_segmentation(18, 8)
        model.load_state_dict(torch.load(weights_path))
        model.to(device)
        model.eval()

        # load network
        self.net = model

        msg = startAffordanceSrvResponse()
        return msg

    def stopAffordance(self, msg):
        del self.net
        self.net = None

        msg = stopAffordanceSrvResponse()
        return msg

if __name__ == '__main__':

    rospy.init_node('affordance_analyzer')
    affordanceServer = AffordanceAnalyzer()
    print("Affordance analyzer is ready.")

    rospy.spin()


#!/usr/bin/env python3
import argparse
import logging
from distutils.util import strtobool

import kaldiio
import numpy

from paddlespeech.s2t.transform.cmvn import CMVN
from paddlespeech.s2t.utils.cli_readers import file_reader_helper
from paddlespeech.s2t.utils.cli_utils import get_commandline_args
from paddlespeech.s2t.utils.cli_utils import is_scipy_wav_style
from paddlespeech.s2t.utils.cli_writers import file_writer_helper

def get_parser():
    parser = argparse.ArgumentParser(
        description="apply mean-variance normalization to files",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter, )

    parser.add_argument(
        "--verbose", "-V", default=0, type=int, help="Verbose option")
    parser.add_argument(
        "--in-filetype",
        type=str,
        default="mat",
        choices=["mat", "hdf5", "sound.hdf5", "sound"],
        help="Specify the file format for the rspecifier. "
        '"mat" is the matrix format in kaldi', )
    parser.add_argument(
        "--stats-filetype",
        type=str,
        default="mat",
        choices=["mat", "hdf5", "npy"],
        help="Specify the file format for the rspecifier. "
        '"mat" is the matrix format in kaldi', )
    parser.add_argument(
        "--out-filetype",
        type=str,
        default="mat",
        choices=["mat", "hdf5"],
        help="Specify the file format for the wspecifier. "
        '"mat" is the matrix format in kaldi', )

    parser.add_argument(
        "--norm-means",
        type=strtobool,
        default=True,
        help="Do variance normalization or not.", )
    parser.add_argument(
        "--norm-vars",
        type=strtobool,
        default=False,
        help="Do variance normalization or not.", )
    parser.add_argument(
        "--reverse",
        type=strtobool,
        default=False,
        help="Do reverse mode or not")
    parser.add_argument(
        "--spk2utt",
        type=str,
        help="A text file of speaker to utterance-list map. "
        "(Don't give rspecifier format, such as "
        '"ark:spk2utt")', )
    parser.add_argument(
        "--utt2spk",
        type=str,
        help="A text file of utterance to speaker map. "
        "(Don't give rspecifier format, such as "
        '"ark:utt2spk")', )
    parser.add_argument(
        "--write-num-frames",
        type=str,
        help="Specify wspecifer for utt2num_frames")
    parser.add_argument(
        "--compress",
        type=strtobool,
        default=False,
        help="Save in compressed format")
    parser.add_argument(
        "--compression-method",
        type=int,
        default=2,
        help="Specify the method(if mat) or "
        "gzip-level(if hdf5)", )
    parser.add_argument(
        "stats_rspecifier_or_rxfilename",
        help="Input stats. e.g. ark:stats.ark or stats.mat", )
    parser.add_argument(
        "rspecifier", type=str, help="Read specifier id. e.g. ark:some.ark")
    parser.add_argument(
        "wspecifier", type=str, help="Write specifier id. e.g. ark:some.ark")
    return parser

def main():
    args = get_parser().parse_args()

    # logging info
    logfmt = "%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s"
    if args.verbose > 0:
        logging.basicConfig(level=logging.INFO, format=logfmt)
    else:
        logging.basicConfig(level=logging.WARN, format=logfmt)
    logging.info(get_commandline_args())

    if ":" in args.stats_rspecifier_or_rxfilename:
        is_rspcifier = True
        if args.stats_filetype == "npy":
            stats_filetype = "hdf5"
        else:
            stats_filetype = args.stats_filetype

        stats_dict = dict(
            file_reader_helper(args.stats_rspecifier_or_rxfilename,
                               stats_filetype))
    else:
        is_rspcifier = False
        if args.stats_filetype == "mat":
            stats = kaldiio.load_mat(args.stats_rspecifier_or_rxfilename)
        else:
            stats = numpy.load(args.stats_rspecifier_or_rxfilename)
        stats_dict = {None: stats}

    cmvn = CMVN(
        stats=stats_dict,
        norm_means=args.norm_means,
        norm_vars=args.norm_vars,
        utt2spk=args.utt2spk,
        spk2utt=args.spk2utt,
        reverse=args.reverse, )

    with file_writer_helper(
            args.wspecifier,
            filetype=args.out_filetype,
            write_num_frames=args.write_num_frames,
            compress=args.compress,
            compression_method=args.compression_method, ) as writer:
        for utt, mat in file_reader_helper(args.rspecifier, args.in_filetype):
            if is_scipy_wav_style(mat):
                # If data is sound file, then got as Tuple[int, ndarray]
                rate, mat = mat
            mat = cmvn(mat, utt if is_rspcifier else None)
            writer[utt] = mat

if __name__ == "__main__":
    main()


import os, sys
import cv2
import torch
import torchvision
import torch.nn.functional as F
import torch.nn as nn
import numpy as np
from glob import glob
import time
import datetime
import imageio

sys.path.append('./models/')
from FLAME import FLAME, FLAMETex
from renderer import Renderer
import util
torch.backends.cudnn.benchmark = True

class PhotometricFitting(object):
    def __init__(self, config, device='cuda'):
        self.batch_size = config.batch_size
        self.image_size = config.image_size
        self.config = config
        self.device = device
        #
        self.flame = FLAME(self.config).to(self.device)
        self.flametex = FLAMETex(self.config).to(self.device)

        self._setup_renderer()

    def _setup_renderer(self):
        mesh_file = './data/head_template_mesh.obj'
        self.render = Renderer(self.image_size, obj_filename=mesh_file).to(self.device)

    def optimize(self, images, landmarks, image_masks, savefolder=None):
        bz = images.shape[0]
        shape = nn.Parameter(torch.zeros(bz, self.config.shape_params).float().to(self.device))
        tex = nn.Parameter(torch.zeros(bz, self.config.tex_params).float().to(self.device))
        exp = nn.Parameter(torch.zeros(bz, self.config.expression_params).float().to(self.device))
        pose = nn.Parameter(torch.zeros(bz, self.config.pose_params).float().to(self.device))
        cam = torch.zeros(bz, self.config.camera_params); cam[:, 0] = 5.
        cam = nn.Parameter(cam.float().to(self.device))
        lights = nn.Parameter(torch.zeros(bz, 9, 3).float().to(self.device))
        e_opt = torch.optim.Adam(
            [shape, exp, pose, cam, tex, lights],
            lr=self.config.e_lr,
            weight_decay=self.config.e_wd
        )
        e_opt_rigid = torch.optim.Adam(
            [pose, cam],
            lr=self.config.e_lr,
            weight_decay=self.config.e_wd
        )

        gt_landmark = landmarks

        # rigid fitting of pose and camera with 51 static face landmarks,
        # this is due to the non-differentiable attribute of contour landmarks trajectory
        for k in range(200):
            losses = {}
            vertices, landmarks2d, landmarks3d = self.flame(shape_params=shape, expression_params=exp, pose_params=pose)
            trans_vertices = util.batch_orth_proj(vertices, cam);
            trans_vertices[..., 1:] = - trans_vertices[..., 1:]
            landmarks2d = util.batch_orth_proj(landmarks2d, cam);
            landmarks2d[..., 1:] = - landmarks2d[..., 1:]
            landmarks3d = util.batch_orth_proj(landmarks3d, cam);
            landmarks3d[..., 1:] = - landmarks3d[..., 1:]

            losses['landmark'] = util.l2_distance(landmarks2d[:, 17:, :2], gt_landmark[:, 17:, :2]) * config.w_lmks

            all_loss = 0.
            for key in losses.keys():
                all_loss = all_loss + losses[key]
            losses['all_loss'] = all_loss
            e_opt_rigid.zero_grad()
            all_loss.backward()
            e_opt_rigid.step()

            loss_info = '----iter: {}, time: {}\n'.format(k, datetime.datetime.now().strftime('%Y-%m-%d-%H:%M:%S'))
            for key in losses.keys():
                loss_info = loss_info + '{}: {}, '.format(key, float(losses[key]))
            if k % 10 == 0:
                print(loss_info)

            if k % 10 == 0:
                grids = {}
                visind = range(bz)  # [0]
                grids['images'] = torchvision.utils.make_grid(images[visind]).detach().cpu()
                grids['landmarks_gt'] = torchvision.utils.make_grid(
                    util.tensor_vis_landmarks(images[visind], landmarks[visind]))
                grids['landmarks2d'] = torchvision.utils.make_grid(
                    util.tensor_vis_landmarks(images[visind], landmarks2d[visind]))
                grids['landmarks3d'] = torchvision.utils.make_grid(
                    util.tensor_vis_landmarks(images[visind], landmarks3d[visind]))

                grid = torch.cat(list(grids.values()), 1)
                grid_image = (grid.numpy().transpose(1, 2, 0).copy() * 255)[:, :, [2, 1, 0]]
                grid_image = np.minimum(np.maximum(grid_image, 0), 255).astype(np.uint8)
                cv2.imwrite('{}/{}.jpg'.format(savefolder, k), grid_image)

        # non-rigid fitting of all the parameters with 68 face landmarks, photometric loss and regularization terms.
        for k in range(200, 1000):
            losses = {}
            vertices, landmarks2d, landmarks3d = self.flame(shape_params=shape, expression_params=exp, pose_params=pose)
            trans_vertices = util.batch_orth_proj(vertices, cam);
            trans_vertices[..., 1:] = - trans_vertices[..., 1:]
            landmarks2d = util.batch_orth_proj(landmarks2d, cam);
            landmarks2d[..., 1:] = - landmarks2d[..., 1:]
            landmarks3d = util.batch_orth_proj(landmarks3d, cam);
            landmarks3d[..., 1:] = - landmarks3d[..., 1:]

            losses['landmark'] = util.l2_distance(landmarks2d[:, :, :2], gt_landmark[:, :, :2]) * config.w_lmks
            losses['shape_reg'] = (torch.sum(shape ** 2) / 2) * config.w_shape_reg  # *1e-4
            losses['expression_reg'] = (torch.sum(exp ** 2) / 2) * config.w_expr_reg  # *1e-4
            losses['pose_reg'] = (torch.sum(pose ** 2) / 2) * config.w_pose_reg

            ## render
            albedos = self.flametex(tex) / 255.
            ops = self.render(vertices, trans_vertices, albedos, lights)
            predicted_images = ops['images']
            losses['photometric_texture'] = (image_masks * (ops['images'] - images).abs()).mean() * config.w_pho

            all_loss = 0.
            for key in losses.keys():
                all_loss = all_loss + losses[key]
            losses['all_loss'] = all_loss
            e_opt.zero_grad()
            all_loss.backward()
            e_opt.step()

            loss_info = '----iter: {}, time: {}\n'.format(k, datetime.datetime.now().strftime('%Y-%m-%d-%H:%M:%S'))
            for key in losses.keys():
                loss_info = loss_info + '{}: {}, '.format(key, float(losses[key]))

            if k % 10 == 0:
                print(loss_info)

            # visualize
            if k % 10 == 0:
                grids = {}
                visind = range(bz)  # [0]
                grids['images'] = torchvision.utils.make_grid(images[visind]).detach().cpu()
                grids['landmarks_gt'] = torchvision.utils.make_grid(
                    util.tensor_vis_landmarks(images[visind], landmarks[visind]))
                grids['landmarks2d'] = torchvision.utils.make_grid(
                    util.tensor_vis_landmarks(images[visind], landmarks2d[visind]))
                grids['landmarks3d'] = torchvision.utils.make_grid(
                    util.tensor_vis_landmarks(images[visind], landmarks3d[visind]))
                grids['albedoimage'] = torchvision.utils.make_grid(
                    (ops['albedo_images'])[visind].detach().cpu())
                grids['render'] = torchvision.utils.make_grid(predicted_images[visind].detach().float().cpu())
                shape_images = self.render.render_shape(vertices, trans_vertices, images)
                grids['shape'] = torchvision.utils.make_grid(
                    F.interpolate(shape_images[visind], [224, 224])).detach().float().cpu()

                # grids['tex'] = torchvision.utils.make_grid(F.interpolate(albedos[visind], [224, 224])).detach().cpu()
                grid = torch.cat(list(grids.values()), 1)
                grid_image = (grid.numpy().transpose(1, 2, 0).copy() * 255)[:, :, [2, 1, 0]]
                grid_image = np.minimum(np.maximum(grid_image, 0), 255).astype(np.uint8)

                cv2.imwrite('{}/{}.jpg'.format(savefolder, k), grid_image)

        single_params = {
            'shape': shape.detach().cpu().numpy(),
            'exp': exp.detach().cpu().numpy(),
            'pose': pose.detach().cpu().numpy(),
            'cam': cam.detach().cpu().numpy(),
            'verts': trans_vertices.detach().cpu().numpy(),
            'albedos':albedos.detach().cpu().numpy(),
            'tex': tex.detach().cpu().numpy(),
            'lit': lights.detach().cpu().numpy()
        }
        return single_params

    def run(self, imagepath, landmarkpath):
        # The implementation is potentially able to optimize with images(batch_size>1),
        # here we show the example with a single image fitting
        images = []
        landmarks = []
        image_masks = []

        image_name = os.path.basename(imagepath)[:-4]
        savefile = os.path.sep.join([self.config.savefolder, image_name + '.npy'])

        # photometric optimization is sensitive to the hair or glass occlusions,
        # therefore we use a face segmentation network to mask the skin region out.
        image_mask_folder = './FFHQ_seg/'
        image_mask_path = os.path.sep.join([image_mask_folder, image_name + '.npy'])

        image = cv2.resize(cv2.imread(imagepath), (config.cropped_size, config.cropped_size)).astype(np.float32) / 255.
        image = image[:, :, [2, 1, 0]].transpose(2, 0, 1)
        images.append(torch.from_numpy(image[None, :, :, :]).to(self.device))

        image_mask = np.load(image_mask_path, allow_pickle=True)
        image_mask = image_mask[..., None].astype('float32')
        image_mask = image_mask.transpose(2, 0, 1)
        image_mask_bn = np.zeros_like(image_mask)
        image_mask_bn[np.where(image_mask != 0)] = 1.
        image_masks.append(torch.from_numpy(image_mask_bn[None, :, :, :]).to(self.device))

        landmark = np.load(landmarkpath).astype(np.float32)
        landmark[:, 0] = landmark[:, 0] / float(image.shape[2]) * 2 - 1
        landmark[:, 1] = landmark[:, 1] / float(image.shape[1]) * 2 - 1
        landmarks.append(torch.from_numpy(landmark)[None, :, :].float().to(self.device))

        images = torch.cat(images, dim=0)
        images = F.interpolate(images, [self.image_size, self.image_size])
        image_masks = torch.cat(image_masks, dim=0)
        image_masks = F.interpolate(image_masks, [self.image_size, self.image_size])

        landmarks = torch.cat(landmarks, dim=0)
        savefolder = os.path.sep.join([self.config.savefolder, image_name])

        util.check_mkdir(savefolder)
        # optimize
        single_params = self.optimize(images, landmarks, image_masks, savefolder)
        self.render.save_obj(filename=savefile[:-4]+'.obj',
                             vertices=torch.from_numpy(single_params['verts'][0]).to(self.device),
                             textures=torch.from_numpy(single_params['albedos'][0]).to(self.device)
                             )
        np.save(savefile, single_params)

if __name__ == '__main__':
    image_name = str(sys.argv[1])
    device_name = str(sys.argv[2])
    config = {
        # FLAME
        'flame_model_path': './data/generic_model.pkl',  # acquire it from FLAME project page
        'flame_lmk_embedding_path': './data/landmark_embedding.npy',
        'tex_space_path': './data/FLAME_texture.npz',  # acquire it from FLAME project page
        'camera_params': 3,
        'shape_params': 100,
        'expression_params': 50,
        'pose_params': 6,
        'tex_params': 50,
        'use_face_contour': True,

        'cropped_size': 256,
        'batch_size': 1,
        'image_size': 224,
        'e_lr': 0.005,
        'e_wd': 0.0001,
        'savefolder': './test_results/',
        # weights of losses and reg terms
        'w_pho': 8,
        'w_lmks': 1,
        'w_shape_reg': 1e-4,
        'w_expr_reg': 1e-4,
        'w_pose_reg': 0,
    }

    config = util.dict2obj(config)
    util.check_mkdir(config.savefolder)

    config.batch_size = 1
    fitting = PhotometricFitting(config, device=device_name)

    input_folder = './FFHQ'

    imagepath = os.path.sep.join([input_folder, image_name + '.png'])
    landmarkpath = os.path.sep.join([input_folder, image_name + '.npy'])
    fitting.run(imagepath, landmarkpath)


import numpy as np
import tensorflow as tf
import gym
import gym.spaces
from gym.wrappers import Monitor

from utils import *

class Gym_atari_env(object):
    def __init__(self,env_name,max_path_length=None,video=False,video_dir=None,image_per_state=4):
        ## arguments
        ## env_name: name of the gym env
        ## max_path_length: max path length before a hard reset, if None, just use env default
        print("INFO: creating new env of %s" % env_name)
        ## create a new env
        self.env = gym.make(env_name)
        # remove the timelimitwrapper
        #self.env = env.env
        self.image_per_state=image_per_state
        ## define max path, either from user input or env default
        self.max_path_length = max_path_length or self.env.spec.max_episode_steps

        ## define parameters directly
        self.discrete =True
        ##all image will be processed to the dimension below
        self.ob_dim = [84,84,self.image_per_state]
        self.ac_dim = 4 if env_name == "Pong-v0" or env_name == "Breakout-v0" else self.env.action_space.n

        ## if video, use monitor
        if video:
            self.env = Monitor(self.env, directory=video_dir, video_callable=lambda x: True, resume=True)
        ## init process image
        #self.process_image

    #@define_scope
    def process_image(self,raw_image):

        ## grep scale
        out = np.mean(raw_image, axis=2).astype(np.uint8)
        ## crop, down sample and normalise
        out = out[34:194,:][::2, ::2] / 255
        ## dim expand
        out = np.hstack([np.zeros((80,2)),out,np.zeros((80,2))])
        out = np.vstack([np.zeros((2,84)),out,np.zeros((2,84))])

        return out

    def step(self,action,obs,sess=None,training=True):
        ## arguments
        ## action: action to take
        ## return ob, rew, done, info

        raw_image,rew,done,info = self.env.step(action)
        next_obs = self.process_image(raw_image)
        next_obs = np.append(obs[:,:,1:], np.expand_dims(next_obs, 2), axis=2)

        return (next_obs,rew,done,info)

    def reset(self,sess=None):
        raw_image = self.env.reset()
        obs = self.process_image(raw_image)
        return np.stack([obs] * self.image_per_state, axis=2)

    def statistic(self):
        return self.discrete, self.max_path_length, self.ob_dim, self.ac_dim


import numpy as np
import heapq
import torch
import pickle

from sklearn import preprocessing
from transformers import BertModel
from gensim.models.fasttext import FastText
from gensim.models.word2vec import Word2Vec
from nltk.corpus import stopwords

from FNTokenizer import FNTokenizer
import Frames
#from glove import Corpus, Glove

config = {"model":BertModel,
          "tokenizer":FNTokenizer}

LIMIT = 1000

def load_fasttext(input_file):
    return FastText.load(input_file)

def load_w2v(input_file):
    return Word2Vec.load(input_file)

def load_bbc(in_file):
    vecs = pickle.load(open(in_file + ".m", "rb"))
    vocab = pickle.load(open(in_file + ".iw", "rb"))

    return Embedding(vecs, vocab)

# Some issues with glove python package, working on...
'''
def load_glove(input_file):
    model = Glove.load(input_file)
    vocab = list(model.dictionary.keys())
    vecs = np.array([model.word_vectors[model.dictionary[v]] for v in vocab])
    return Embedding(vecs, vocab)
'''

class Embedding(object):
    """
    Base class for all embeddings.
    Based heavily on the Embedding class from https://github.com/williamleif/histwords
    """

    @staticmethod
    def from_bert_model_and_text(vec_model, texts, output_file=None, agg=np.mean):
        """
        Given a (transformer) embedding model and a text, return an embedding space at the word level.
        Takes the embedding for each word found in the text, and combines them using the agg function.
        """
        vocab, vecs = None, None
        model = BertModel.from_pretrained(vec_model)

        # Need special tokenizer to keep frames as their own tokens
        tokenizer = FNTokenizer.from_pretrained("bert-base-cased")
        tokenizer.add_tokens([f.upper() for f in Frames.frames])

        embeds = {}

        for c, line in enumerate(texts):
            input_ids = torch.tensor(tokenizer.encode(line)).unsqueeze(0)
            outputs = model(input_ids)
            last_hidden_states = outputs[0]
            tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
            for i, token in enumerate(tokens):
                if token not in stopwords.words('english'):
                    embed = last_hidden_states[0][i].detach().numpy()
                    if token not in embeds:
                        embeds[token] = []
                    if len(embeds[token]) < LIMIT:
                        embeds[token].append(embed)

            if (c+1) % 1000 == 0 or c == len(texts)-1:
                print (c, len(texts), len(embeds), np.mean([len(embeds[v]) for v in embeds]))
                types = list(embeds.keys())
                res = {}
                for key in embeds:
                    res[key] = agg(embeds[key], axis=0)

                vocab = sorted(list(res.keys()))
                vecs = np.array([res[v] for v in vocab])
                if (c+1) % 10000 == 0 and output_file:
                    e = Embedding(vecs, vocab)
                    e.save("embeddings/" + output_file + str((c+1)/10000))
        return Embedding(vecs, vocab)

    @staticmethod
    def new_init(input_file, normalize=True, **kwargs):
        """
        Alternate method for initialization which uses gensim formating.

        :param input_file:
        :param normalize:
        :param kwargs:
        :return:
        """
        vecs, vocab = [], []
        for line in open(input_file):
            vocab.append(line.split()[0])
            vecs.append(line.split()[1:])

        vecs = np.array(vecs, dtype="float")
        vocab = set(vocab)
        return Embedding(vecs, vocab, normalize)

    def __init__(self, vecs, vocab, normalize=True, **kwargs):
        """
        Normal initialization, which requires all the parameters

        :param vecs: matrix of vectors
        :param vocab: vocabulary
        :param normalize:
        :param kwargs:
        """
        self.m = vecs
        self.dim = self.m.shape[1]
        self.iw = vocab
        self.wi = {w:i for i,w in enumerate(self.iw)}
        if normalize:
            self.normalize()

    def __getitem__(self, key):
        if self.oov(key):
            raise KeyError
        else:
            return self.represent(key)

    def __iter__(self):
        return self.iw.__iter__()

    def __contains__(self, key):
        return not self.oov(key)

    def normalize(self):
        preprocessing.normalize(self.m, copy=False)

    def oov(self, w):
        return not (w in self.wi)

    def represent(self, w):
        if w in self.wi:
            return self.m[self.wi[w], :]
        else:
            print ("OOV: ", w)
            return np.zeros(self.dim)

    def similarity(self, w1, w2):
        """
        Assumes the vectors have been normalized.
        """
        sim = self.represent(w1).dot(self.represent(w2))
        return sim

    def closest(self, w, n=10):
        """
        Assumes the vectors have been normalized.
        """
        scores = self.m.dot(self.represent(w))
        return heapq.nlargest(n, zip(scores, self.iw))

    def save(self, out_file):
        pickle.dump(self.m, open(out_file + ".m", "wb"))
        pickle.dump(self.iw, open(out_file + ".iw", "wb"))


import os
from werkzeug.utils import secure_filename
import imutils
import cv2
import base64
import os
import cv2
import numpy as np
import tensorflow as tf
import sys

from utils import label_map_util
from utils import visualization_utils as vis_util

stopper = False

global MODEL_NAME
MODEL_NAME = 'trained-inference-graphs' #'inference_graph'

global CWD_PATH
CWD_PATH = os.getcwd()

global PATH_TO_CKPT
PATH_TO_CKPT = os.path.join(
    CWD_PATH, 'services', 'frozen_inference_graph.pb')

global PATH_TO_LABELS
PATH_TO_LABELS = os.path.join(CWD_PATH, 'services', 'label_map.pbtxt')
global NUM_CLASSES
NUM_CLASSES = 15

global label_map
label_map = label_map_util.load_labelmap(PATH_TO_LABELS)

global categories
categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True)

global category_index
category_index = label_map_util.create_category_index(categories)

global detection_graph
detection_graph = tf.Graph()
with detection_graph.as_default():
    od_graph_def = tf.GraphDef()
    with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
        serialized_graph = fid.read()
        od_graph_def.ParseFromString(serialized_graph)
        tf.import_graph_def(od_graph_def, name='')

global sess
sess = tf.Session(graph=detection_graph)

global image_tensor
image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')

global detection_boxes
detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')

global detection_scores
detection_scores = detection_graph.get_tensor_by_name('detection_scores:0')

global detection_classes
detection_classes = detection_graph.get_tensor_by_name('detection_classes:0')

global num_detections
num_detections = detection_graph.get_tensor_by_name('num_detections:0')

def process_image(reqData, upload_folder):
    try:

        f = reqData.files['file']
        file_save_path = os.path.join(
            upload_folder, "images", secure_filename(f.filename))
        f.save(file_save_path)

        image = cv2.imread(os.path.join(file_save_path))

        image_expanded = np.expand_dims(image, axis=0)

        # Perform the actual detection by running the model with the image as input
        (boxes, scores, classes, num) = sess.run(
            [detection_boxes, detection_scores, detection_classes, num_detections],
            feed_dict={image_tensor: image_expanded})

        # Draw the results of the detection (aka 'visulaize the results')

        vis_util.visualize_boxes_and_labels_on_image_array(
            image,
            np.squeeze(boxes),
            np.squeeze(classes).astype(np.int32),
            np.squeeze(scores),
            category_index,
            use_normalized_coordinates=True,
            line_thickness=8,
            min_score_thresh=0.60)

        _, img_encoded = cv2.imencode(".jpg", image)
        imag64 = base64.b64encode(img_encoded)
        # print(imag64)
        return imag64
    except Exception as e:
        print(e)

def process_live():

    cap = cv2.VideoCapture(0)
    ret = cap.set(3,1280)
    ret = cap.set(4,720)
    global stopper

    while True:
        ret, frame = cap.read()
        frame_expanded = np.expand_dims(frame, axis=0)

        # Perform the actual detection by running the model with the image as input
        (boxes, scores, classes, num) = sess.run(
            [detection_boxes, detection_scores, detection_classes, num_detections],
            feed_dict={image_tensor: frame_expanded})

        # Draw the results of the detection (aka 'visulaize the results')
        vis_util.visualize_boxes_and_labels_on_image_array(
            frame,
            np.squeeze(boxes),
            np.squeeze(classes).astype(np.int32),
            np.squeeze(scores),
            category_index,
            use_normalized_coordinates=True,
            line_thickness=8,
            min_score_thresh=0.60)
        (flag, encodeImg) = cv2.imencode(".jpg", frame)

        yield(b'--frame\r\n' b'Content-Type: image/jpeg\r\n\r\n' +
              bytearray(encodeImg) + b'\r\n')
        if stopper == True:
            break

def stop_stream():
    global stopper
    stopper = True


#!/var/scratch/mao540/miniconda3/envs/maip-venv/bin/python

import argparse
import os
import torch
import torch.optim as optim
import torch.backends.cudnn as cudnn
import torch.utils.data as data
import torchvision
import torchvision.transforms as transforms
from torchvision import datasets
from torch.utils.data import DataLoader

from numpy import log
# from models import RealNVP, RealNVPLoss
from model import Glow
from tqdm import tqdm
# from train_r import calc_z_shapes
from shell_util import AverageMeter, bits_per_dim
from load_data import sample_from_directory

def main(args):

    device = torch.device("cuda" if torch.cuda.is_available() and len(args.gpu_ids) > 0 else "cpu")
    print("training on: %s" % device)
    start_epoch = 0

    if args.net == 'glow':
        model = Glow(3, args.n_flow, args.n_block, affine=args.affine, conv_lu=not args.no_lu)
        net = model.to(device)
    else:
        net = RealNVP( **filter_args(args.__dict__) )
        net = net.to(device)

    if str(device).startswith('cuda'):
        net = torch.nn.DataParallel(net, args.gpu_ids)
        cudnn.benchmark = args.benchmark

    if args.resume: # or not args.resume:
        # Load checkpoint.
        # import pdb; pdb.set_trace()
        args.model_dir = find_last_model_relpath(args.root_dir) # /model_{str(i + 1).zfill(6)}.pt'
        print(f'Resuming from checkpoint at {args.model_dir}')
        checkpoint = torch.load(args.model_dir+'/model.pth.tar')
        net.load_state_dict(checkpoint['net'])
        global best_loss
        best_loss = checkpoint['test_loss']
        # we start epoch after the saved one (avoids overwrites).
        start_epoch = checkpoint['epoch'] + 1
    else:
        os.makedirs(args.root_dir, exist_ok=True)
        os.makedirs(args.sample_dir, exist_ok=True)

    if args.net == 'glow':
        optimizer = optim.Adam(net.parameters(), lr=args.lr)
        if args.resume:
            optimizer.load_state_dict(torch.load(f'{model_dir}/optim.pt'))
    else:
        loss_fn = RealNVPLoss()
        param_groups = util.get_param_groups(net, args.weight_decay, norm_suffix='weight_g')
        optimizer = optim.Adam(param_groups, lr=args.lr, eps=1e-7)

    train(args, net, device, optimizer, start_epoch)

def calc_loss(log_p, logdet, image_size, n_bins):
    # TODO: move to utils or Network module
    # log_p = calc_log_p([z_list])
    n_pixel = image_size * image_size * 3

    loss = -log(n_bins) * n_pixel
    loss = loss + logdet + log_p

    return (
        (-loss / (log(2) * n_pixel)).mean(),
        (log_p / (log(2) * n_pixel)).mean(),
        (logdet / (log(2) * n_pixel)).mean(),
    )

def train(args, net, device, optimizer, start_epoch):

    if args.dataset == 'celeba':
        dataset = iter(sample_celeba(args.batch, args.img_size))
    elif args.dataset == 'ffhq':
        dataset = iter(sample_from_directory('data/FFHQ/images1024x1024', args.batch, args.img_size))

    n_bins = 2. ** args.n_bits

    z_sample = find_or_make_z(args.root_dir + '/z_samples.pkl',
                              3, args.img_size, args.n_flow, args.n_block,
                              args.num_sample, args.temp, device)

    loss_meter = AverageMeter()
    bpd_meter = AverageMeter()

    p_imgs = 0
    net.train()
    with tqdm(range(start_epoch, args.iter)) as pbar:
        for i in pbar:
            x, _ = next(dataset)
            x = x.to(device)

            if i == 0:
                with torch.no_grad():
                    # log_p, logdet, _ = net.module(x + torch.rand_like(x) / n_bins)
                    log_p, logdet, _ = net(x + torch.rand_like(x) / n_bins)
                    continue
            else:
                log_p, logdet, _ = net(x + torch.rand_like(x) / n_bins)

            logdet = logdet.mean()

            loss, log_p, log_det = calc_loss(log_p, logdet, args.img_size, n_bins)
            net.zero_grad()
            loss.backward()
            # warmup_lr = args.lr * min(1, i * args.batch / (50000 * 10))
            warmup_lr = args.lr
            optimizer.param_groups[0]['lr'] = warmup_lr
            optimizer.step()
            loss_meter.update(loss.item(), x.size(0))
            bpd_meter.update(bits_per_dim(x, loss_meter.avg))

            pbar.set_description(
                    f'Loss: {loss.item():.5f}; logP: {log_p.item():.5f}; logdet: {log_det.item():.5f}; lr: {warmup_lr:.7f}; imgs: {p_imgs}'
            )
            p_imgs += x.size(0)
            if i % 1000 == 0:
                net.eval()
                # SAMPLE
                with torch.no_grad():
                    torchvision.utils.save_image(net(z_sample, reverse=True).cpu().data,
                                                 f'{args.sample_dir}/{str(i).zfill(6)}.png',
                                                 normalize=True,
                                                 nrow= int(args.num_sample ** 0.5),
                                                 range=(-0.5, 0.5))
                with open(f'{args.root_dir}/log', 'a') as l:
                    report = f'{loss.item():.5f},{log_p.item():.5f},{log_det.item():.5f},{warmup_lr:.7f},{p_imgs}\n'
                    # print("Writing to disk: " + report + ">> {}/log".format(args.root_dir))
                    l.write(report)

                model_dir = args.root_dir
                if i % 10000 == 0:
                    # TEST
                    model_dir += f'/epoch_{str(i).zfill(6)}'
                    os.makedirs(model_dir, exist_ok=True)
                torch.save({'net': net.state_dict(),
                            'test_loss': loss_meter.avg,
                                        'epoch': i
                                        },
                                        f'{model_dir}/model.pth.tar'
                )
                torch.save(
                    optimizer.state_dict(), f'{model_dir}/optim.pt'
                )
            net.train()

def sample_celeba(batch, image_size, test=False):
    if not test:
        split = 'train'
        shuffle = True
        transform = transforms.Compose([
            transforms.CenterCrop(160),
            transforms.Resize(size=image_size),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (1, 1, 1)),
        ])
    else:
        split = 'test'
        transform = transforms.Compose([
        transforms.CenterCrop(160),
        transforms.Resize(size=image_size),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (1, 1, 1)),
    ])
    print(f'shuffle set to {shuffle} for split: {split}')
    target_type = ['attr', 'bbox', 'landmarks']
    dataset = datasets.CelebA(root='data', split=split, target_type=target_type[0], download=True, transform=transform)
    loader = DataLoader(dataset, batch_size=batch, shuffle=shuffle, num_workers=8)
    loader = iter(loader)

    while True:
        try:
            yield next(loader)
        except StopIteration:
            loader = DataLoader(dataset, batch_size=batch, shuffle=True, num_workers=8)
            loader = iter(loader)
            yield next(loader)

def find_or_make_z(path, C, img_size, n_flow, n_block, num_sample, t, device):

    if os.path.isfile(path):
        z_sample = torch.load(path)

    else:
        z_sample = []
        z_shapes = calc_z_shapes(3, img_size, n_flow, n_block)
        for z in z_shapes:
            z_new = torch.randn(num_sample, *z) * t
            z_sample.append(z_new.to(device))

        torch.save(z_sample, path)
    return z_sample

def calc_z_shapes(n_channel, input_size, n_flow, n_block):
    z_shapes = []

    for i in range(n_block - 1):
        input_size //= 2
        n_channel *= 2
        z_shapes.append((n_channel, input_size, input_size))

    input_size //= 2
    z_shapes.append((n_channel * 4, input_size, input_size))
    return z_shapes

def save_imgrid(tensor, name):
    grid = torchvision.utils.make_grid(tensor, nrow=int(tensor.shape[0] ** 0.5), padding=1, pad_value=255)
    torchvision.utils.save_image(grid, name)
    return

def sample(net, num_samples, in_channels, device, resize_hw=None):
    """Sample from RealNVP model.

    Args:
        net (torch.nn.DataParallel): The RealNVP model wrapped in DataParallel.
        batch_size (int): Number of samples to generate.
        device (torch.device): Device to use.
    """

    if not resize_hw:
        side_size = 28
    else:
        side_size, side_size = resize_hw
    print(f"sampling with z space sized: {side_size}x{side_size}.")
    z = torch.randn((num_samples, in_channels, side_size, side_size), dtype=torch.float32, device=device) #changed 3 -> 1
    x, _ = net(z, reverse=True)
    return x, z

def test(epoch, net, testloader, device, loss_fn, **args):
    global best_loss
    net.eval()
    loss_meter = util.AverageMeter()
    bpd_meter = util.AverageMeter()
    with torch.no_grad():
        with tqdm(total=len(testloader.dataset)+1) as progress_bar:
            for x, _ in testloader:
                x = x.to(device)
                z, sldj = net(x, reverse=False)
                loss = loss_fn(z, sldj)
                loss_meter.update(loss.item(), x.size(0))
                # bits per dimensions
                bpd_meter.update(util.bits_per_dim(x, loss_meter.avg), x.size(0))

                progress_bar.set_postfix(loss=loss_meter.avg,
                                                                 bpd=bpd_meter.avg)
                progress_bar.update(x.size(0))

    # Save checkpoint
    save_dir = args['dir_samples'] + '/epoch_{:03d}'.format(epoch) #  + str(epoch)
    os.makedirs(save_dir, exist_ok=True)

    # if loss_meter.avg < best_loss or epoch % 10 == 0 or
    # 		epoch > 100 or epoch < 20:
    if True:
        print('\nSaving...')
        state = {
            'net': net.state_dict(),
            'test_loss': loss_meter.avg,
            'epoch': epoch,
        }
        torch.save(state, save_dir + '/model.pth.tar')
        best_loss = loss_meter.avg

    # import ipdb; ipdb.set_trace()
    sample_fields = ['num_samples', 'in_channels', 'resize_hw']
    images, latent_z = sample(net, device=device, **filter_args( args, fields=sample_fields ) )

    # plot x and z
    num_samples = args['num_samples']
    images_concat = torchvision.utils.make_grid(images, nrow=int(num_samples ** 0.5), padding=2, pad_value=255)
    z_concat = torchvision.utils.make_grid(latent_z, nrow=int(num_samples ** 0.5), padding=2, pad_value=255)
    torchvision.utils.save_image(images_concat, save_dir+'/x.png')
    torchvision.utils.save_image(z_concat, save_dir+'/z.png')

    # with open(, 'wb') as z_serialize:
    # 	pickle.dump(latent_z, z_serialize)
    torch.save(latent_z, f = save_dir+'/z.pkl')

    # dict keys as returned by "train"
    train_loss = args['train_loss']
    train_bpd = args['train_bpd']
    report = [epoch, loss_meter.avg, bpd_meter.avg] + [train_loss, train_bpd]

    dir_samples = args['dir_samples']
    with open('{}/log'.format(dir_samples), 'a') as l:
        report = ", ".join([str(m) for m in report])
        report += "\n"
        print("\nWriting to disk:\n" + report + "At {}".format(dir_samples))
        l.write(report)

def filter_args(arg_dict, fields=None):
    """only pass to network architecture relevant fields."""
    if not fields:
        # if arg_dict['net_type'] == 'resnet':
        fields = ['net_type', 'num_scales', 'in_channels', 'mid_channels']
        # elif arch['net_type'] == 'densenet':
        # 	arch_fields = ['net_type', 'num_scales', 'in_channels', 'mid_channels', 'depth']
    return {k:arg_dict[k] for k in fields if k in arg_dict}

class GaussianNoise(object):

    def __init__(self, mean=0., std=.1, restrict_range=True):
        self.std = std
        self.mean = mean
        self.restrict_range = restrict_range

    def __call__(self, tensor):
        tensor += torch.randn(tensor.size()) * self.std + self.mean
        if self.restrict_range:
            return tensor.clamp(1e-8, 1)
        else:
            return tensor

    def __repr__(self):
        return self.__class__.__name__ + '(mean={0}, std={1})'.format(self.mean, self.std)

class Normie(object):
    '''class for normies'''
    def __init__(self, min, max):
        self.min = min
        self.max = max

    def __call__(self, tensor):
        tensor -= tensor.min()
        tensor /= tensor.max()
        return tensor

def find_last_model_relpath(fp):
    # import pdb; pdb.set_trace()

    dirs_l = os.listdir(fp)
    samples_l = os.listdir(fp + '/samples')
    dirs_e = [int(d.split('_')[-1]) for d in dirs_l if d.startswith('epoch_')]
    samples_fns = [int(png.split('.')[0]) for png in samples_l if png.endswith('png')]
    dirs_e.sort()
    samples_fns.sort()
    if dirs_e[-1] >= samples_fns[-1]:
        return f'{fp}/epoch_{dirs_e[-1]}'
    else:
        return fp

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Glow model')
    parser.add_argument('--benchmark', action='store_true', help='Turn on CUDNN benchmarking')
    # parser.add_argument('--num_epochs', default=100, type=int, help='Number of epochs to train')

    # 1. Dataset : 'celeba', 'MNIST', 'CIFAR' (not tested)
    dataset_ = 'ffhq'
    # 2. Architecture
    net_ = 'glow'  # 2.
    # 3. Samples dir_
    dir_ = net_ + '_' + dataset_
    # 4. GPUs
    gpus_ = '[0, 1]' # if net_ == 'densenet' and dataset_=='mnist'  else '[0]' # 4.
    # 5. resume training?
    resume_ = True # 5.
    # 6. learning_rate
    if net_ == 'glow':
        learning_rate_ = 2e-4
    else:
        learning_rate_ = 1e-4
    # 6. resize
    if dataset_ == 'mnist':
        in_channels_= 1
    elif dataset_ in ['celeba', 'ffhq']:
        in_channels_= 3
        img_size_ = 256
        batch_size_ = 4
        num_sample_ = 16

    root_dir_ = 'data/' + dir_

    parser.add_argument('--img_size', default=img_size_, type=int, help='Image size') # changed from 1e-3 for MNIST
    parser.add_argument('--lr', default=learning_rate_, type=float, help='Learning rate') # changed from 1e-3 for MNIST
    parser.add_argument('--resume', '-r', action='store_true', default=resume_, help='Resume from checkpoint')
    parser.add_argument('--gpu_ids', default=gpus_, type=eval, help='IDs of GPUs to use')
    parser.add_argument('--net', default=net_, help='CNN architecture (resnet or densenet)')
    parser.add_argument('--sample_dir', default= root_dir_ +'/samples', help="Directory for storing generated samples")

    # dataset
    parser.add_argument('--dataset', '-ds', default=dataset_.lower(), type=str, help="MNIST or CIFAR-10")
    parser.add_argument('--in_channels', default=in_channels_, type=int, help='dimensionality along Channels')

    # architecture
    if net_ == 'glow':

        parser.add_argument(
        '--affine', action='store_true', help='use affine coupling instead of additive'
        )
        parser.add_argument('--no_lu', action='store_true', help="don't use LU decomposed convolution")
        parser.add_argument('--n_bits', default=5, type=int, help='number of bits')
        parser.add_argument('--n_flow', default=32, type=int, help='number of bits')
        parser.add_argument('--n_block', default=4, type=int, help='number of bits')
        parser.add_argument('--temp', default=0.7, type=float, help='temperature of sampling')
        parser.add_argument('--iter', default=999999, type=int, help='maximum iterations')
        # if dataset_ == 'celeba':
        num_scales_ = 4

        # elif dataset_ == 'mnist': # data/dense_test6
        # 	batch_size_ = 1024 if len(gpus_) > 3 else 512
        # 	### mid_channels_ = 120
        # 	num_samples_ = 121
        # 	# num_scales_ = 3

    parser.add_argument('--batch', default=batch_size_, type=int, help='Batch size')
    parser.add_argument('--root_dir', default=root_dir_, help="Directory for storing generated samples")
    parser.add_argument('--num_sample', default=num_sample_, type=int, help='Number of samples at test time')
    parser.add_argument('--num_scales', default=num_scales_, type=int, help='Real NVP multi-scale arch. recursions')

    best_loss = 5e5
    main(parser.parse_args())


from ioUtils import getFile
from fsUtils import isFile
from webUtils import getHTML, isBS4
from strUtils import fixName
from math import ceil, floor

from dbBase import dbBase

class artistAMIDClass:
    def __init__(self, ID=None, err=None):
        self.ID=ID
        self.err=err

    def get(self):
        return self.__dict__

class artistAMURLClass:
    def __init__(self, url=None, err=None):
        self.url = url
        self.err = err

    def get(self):
        return self.__dict__

class artistAMNameClass:
    def __init__(self, name=None, err=None):
        self.name = name
        self.err  = err

    def get(self):
        return self.__dict__

class artistAMMediaClass:
    def __init__(self, err=None):
        self.media = {}
        self.err   = err

    def get(self):
        return self.__dict__

class artistAMMediaDataClass:
    def __init__(self, album=None, url=None, aclass=None, aformat=None, artist=None, code=None, year=None, err=None):
        self.album   = album
        self.url     = url
        self.aclass  = aclass
        self.aformat = aformat
        self.artist  = artist
        self.code    = code
        self.year    = year
        self.err     = err

    def get(self):
        return self.__dict__

class artistAMMediaAlbumClass:
    def __init__(self, url=None, album=None, aformat=None, err=None):
        self.url     = url
        self.album   = album
        self.aformat = aformat
        self.err     = err

    def get(self):
        return self.__dict__

class artistAMMediaCountsClass:
    def __init__(self, err=None):
        self.counts = {}
        self.err    = err

    def get(self):
        return self.__dict__

class artistAMPageClass:
    def __init__(self, ppp = None, tot = None, more=None, redo=None, err=None):
        self.ppp   = ppp
        self.tot   = tot
        if isinstance(ppp, int) and isinstance(tot, int):
            self.pages = int(ceil(tot/ppp))
        else:
            self.pages = None

        self.err   = err

        self.more  = more
        self.redo  = redo

    def get(self):
        return self.__dict__

class artistAMProfileClass:
    def __init__(self, profile=None, aliases=None, members=None, sites=None, groups=None, search=None, variations=None, err=None):
        self.profile    = profile
        self.aliases    = aliases
        self.members    = members
        self.sites      = sites
        self.groups     = groups
        self.variations = variations
        self.search     = search
        self.err        = err

    def get(self):
        return self.__dict__

class artistAMURLInfo:
    def __init__(self, name=None, url=None, ID=None, err=None):
        self.name = name
        self.url  = url
        self.ID   = ID
        self.err  = err

    def get(self):
        return self.__dict__

class artistAMDataClass:
    def __init__(self, artist=None, url=None, ID=None, pages=None, profile=None, media=None, mediaCounts=None, err=None):
        self.artist      = artist
        self.url         = url
        self.ID          = ID
        self.pages       = pages
        self.profile     = profile
        self.media       = media
        self.mediaCounts = mediaCounts
        self.err         = err

    def show(self):
        print("AllMusic Artist Data Class")
        print("-------------------------")
        print("Artist:  {0}".format(self.artist.name))
        print("URL:     {0}".format(self.url.url))
        print("ID:      {0}".format(self.ID.ID))
        print("Profile: {0}".format(self.profile.get()))
        print("Pages:   {0}".format(self.pages.get()))
        print("Media:   {0}".format(self.mediaCounts.get()))
        for mediaType,mediaTypeAlbums in self.media.media.items():
            print("   {0}".format(mediaType))
            for album in mediaTypeAlbums:
                print("      {0}".format(album.album))

    def get(self):
        return self.__dict__

class artistAM(dbBase):
    def __init__(self, debug=False):
        self.debug = debug

    def getData(self, inputdata):
        if isinstance(inputdata, str):
            if isFile(inputdata):
                try:
                    bsdata = getHTML(getFile(inputdata))
                except:
                    try:
                        bsdata = getHTML(getFile(inputdata, version=2))
                    except:
                        raise ValueError("Cannot read artist file: {0}".format(inputdata))
            else:
                try:
                    bsdata = getHTML(inputdata)
                except:
                    raise ValueError("Not sure about string input: {0} . It is not a file".format(inputdata))
        elif isBS4(inputdata):
            bsdata = inputdata
            pass
        else:
            raise ValueError("Not sure about input type: {0}".format(type(inputdata)))

        self.bsdata = bsdata

        return self.parse()

    def getNamesAndURLs(self, content):
        data = []
        if content is not None:
            for ref in content.findAll("a"):
                url    = ref.attrs['href']
                name   = ref.text

                ID = self.getartistAMDiscID(url)
                data.append(artistAMURLInfo(name=name, url=url, ID=ID))
        return data

    #######################################################################################################################################
    ## Artist URL
    #######################################################################################################################################
    def getartistAMURL(self):
        result1 = self.bsdata.find("link", {"rel": "canonical"})
        result2 = self.bsdata.find("link", {"hreflang": "en"})
        if result1 and not result2:
            result = result1
        elif result2 and not result1:
            result = result2
        elif result1 and result2:
            result = result1
        else:
            auc = artistAMURLClass(err=True)
            return auc

        if result:
            url = result.attrs["href"]
            url = url.replace("https://www.allmusic.com", "")
            if url.find("/artist/") == -1:
                url = None
                auc = artistAMURLClass(url=url, err="NoArtist")
            else:
                auc = artistAMURLClass(url=url)
        else:
            auc = artistAMURLClass(err="NoLink")

        return auc

    #######################################################################################################################################
    ## Artist ID
    #######################################################################################################################################
    def getartistAMDiscID(self, suburl):
        ival = "/artist"
        if isinstance(suburl, artistAMURLClass):
            suburl = suburl.url
        if not isinstance(suburl, str):
            aic = artistAMIDClass(err="NotStr")
            return aic

        pos = suburl.find(ival)
        if pos == -1:
            aic = artistAMIDClass(err="NotArtist")
            return aic

        data = suburl[pos+len(ival)+1:]
        pos  = data.rfind("-")
        discIDurl = data[(pos+3):]
        discID = discIDurl.split("/")[0]

        try:
            int(discID)
        except:
            aic = artistAMIDClass(err="NotInt")
            return aic

        aic = artistAMIDClass(ID=discID)
        return aic

    #######################################################################################################################################
    ## Artist Name
    #######################################################################################################################################
    def getartistAMName(self):
        artistBios = self.bsdata.findAll("div", {"class": "artist-bio-container"})
        if len(artistBios) > 0:
            for div in artistBios:
                h1 = div.find("h1", {"class": "artist-name"})
                if h1 is not None:
                    artistName = h1.text.strip()
                    if len(artistName) > 0:
                        artist = fixName(artistName)
                        anc = artistAMNameClass(name=artist, err=None)
                    else:
                        anc = artistAMNameClass(name=artist, err="Fix")
                else:
                    anc = artistAMNameClass(err="NoH1")
        else:
            anc = artistAMNameClass(err=True)
            return anc

        return anc

    #######################################################################################################################################
    ## Artist Media
    #######################################################################################################################################
    def getartistAMMediaAlbum(self, td):
        amac = artistAMMediaAlbumClass()
        for span in td.findAll("span"):
            attrs = span.attrs
            if attrs.get("class"):
                if 'format' in attrs["class"]:
                    albumformat = span.text
                    albumformat = albumformat.replace("(", "")
                    albumformat = albumformat.replace(")", "")
                    amac.format = albumformat
                    continue
            span.replaceWith("")

        ref = td.find("a")
        if ref:
            amac.url   = ref.attrs['href']
            amac.album = ref.text
        else:
            amac.err = "NoText"

        return amac

    def getartistAMMedia(self):
        amc  = artistAMMediaClass()
        name = "Albums"
        amc.media[name] = []

        tables = self.bsdata.findAll("table")
        for table in tables:
            trs = table.findAll("tr")

            header  = trs[0]
            ths     = header.findAll("th")
            headers = [x.text.strip() for x in ths]

            for tr in trs[1:]:
                tds = tr.findAll("td")

                ## Name
                key = "Name"
                if len(headers[1]) == 0:
                    idx = 1
                    mediaType = tds[idx].text.strip()
                    if len(mediaType) == 0:
                        mediaType = name
                else:
                    mediaType = name

                ## Year
                key  = "Year"
                idx  = headers.index(key)
                year = tds[idx].text.strip()

                ## Title
                key   = "Album"
                idx   = headers.index(key)
                ref   = tds[idx].findAll("a")
                try:
                    refdata = ref[0]
                    url     = refdata.attrs['href']
                    album   = refdata.text.strip()

                    data = url.split("/")[-1]
                    pos  = data.rfind("-")
                    discIDurl = data[(pos+3):]
                    discID = discIDurl.split("/")[0]

                    try:
                        int(discID)
                        code = discID
                    except:
                        code = None

                except:
                    print("Could not parse {0}".format(ref))

                amdc = artistAMMediaDataClass(album=album, url=url, aclass=None, aformat=None, artist=None, code=code, year=year)
                if amc.media.get(mediaType) is None:
                    amc.media[mediaType] = []
                amc.media[mediaType].append(amdc)

        return amc

    #######################################################################################################################################
    ## Artist Media Counts
    #######################################################################################################################################
    def getartistAMMediaCounts(self, media):

        amcc = artistAMMediaCountsClass()

        credittype = "Releases"
        if amcc.counts.get(credittype) == None:
            amcc.counts[credittype] = {}
        for creditsubtype in media.media.keys():
            amcc.counts[credittype][creditsubtype] = int(len(media.media[creditsubtype]))

        return amcc

        amcc.err = "No Counts"
        return amcc

        results = self.bsdata.findAll("ul", {"class": "facets_nav"})
        if results is None or len(results) == 0:
            amcc.err = "No Counts"
            return amcc

        for result in results:
            for li in result.findAll("li"):
                ref = li.find("a")
                if ref:
                    attrs = ref.attrs
                    span = ref.find("span", {"class": "facet_count"})
                    count = None
                    if span:
                        count = span.text
                        credittype    = attrs.get("data-credit-type")
                        creditsubtype = attrs.get("data-credit-subtype")
                        if credittype and creditsubtype:
                            if amcc.counts.get(credittype) == None:
                                amcc.counts[credittype] = {}
                            if amcc.counts[credittype].get(creditsubtype) == None:
                                try:
                                    amcc.counts[credittype][creditsubtype] = int(count)
                                except:
                                    amcc.counts[credittype][creditsubtype] = count
                                    amcc.err = "Non Int"

        return amcc

    #######################################################################################################################################
    ## Artist Variations
    #######################################################################################################################################
    def getartistAMProfile(self):
        from json import loads
        result = self.bsdata.find("section", {"class": "basic-info"})
        if result is None:
            try:
                content = self.bsdata.find("meta", {"name": "title"})
                searchTerm = content.attrs["content"]
                searchTerm = searchTerm.replace("Artist Search for ", "")
                searchTerm = searchTerm.replace(" | AllMusic", "")
                searchTerm = searchTerm[1:-1]
            except:
                apc = artistAMProfileClass(err="No Profile")
                return apc

            apc = artistAMProfileClass(search=searchTerm)
            return apc

        data   = {}

        members = result.find("div", {"class": "group-members"})
        if members is not None:
            data["Members"] = [item.text.strip() for item in members.findAll("span")]
        else:
            data["Members"] = []

        genres = result.find("div", {"class": "genre"})
        genre  = self.getNamesAndURLs(genres)
        styles = result.find("div", {"class": "styles"})
        style = self.getNamesAndURLs(styles)
        #data["Profile"] = str({'genre': genre, 'style': style})
        data["Profile"] = {'genre': genre, 'style': style}

        apc = artistAMProfileClass(profile=data.get("Profile"), aliases=data.get("Aliases"),
                                 members=data.get("Members"), groups=data.get("In Groups"),
                                 sites=data.get("Sites"), variations=data.get("Variations"))
        return apc

    #######################################################################################################################################
    ## Artist Pages
    #######################################################################################################################################
    def getartistAMPages(self):
        apc   = artistAMPageClass()
        from numpy import ceil
        bsdata = self.bsdata

        apc   = artistAMPageClass(ppp=1, tot=1, redo=False, more=False)
        return apc

        pageData = bsdata.find("div", {"class": "pagination bottom"})
        if pageData is None:
            err = "pagination bottom"
            apc = artistAMPageClass(err=err)
            return apc
        else:
            x = pageData.find("strong", {"class": "pagination_total"})
            if x is None:
                err = "pagination_total"
                apc = artistAMPageClass(err=err)
                return apc
            else:
                txt = x.text
                txt = txt.strip()
                txt = txt.replace("\n", "")
                retval = [tmp.strip() for tmp in txt.split('of')]

                try:
                    ppp   = int(retval[0].split('–')[-1])
                    tot   = int(retval[1].replace(",", ""))
                except:
                    err   = "int"
                    apc   = artistAMPageClass(err=err)
                    return apc

                if ppp < 500:
                    if tot < 25 or ppp == tot:
                        apc   = artistAMPageClass(ppp=ppp, tot=tot, redo=False, more=False)
                    else:
                        apc   = artistAMPageClass(ppp=ppp, tot=tot, redo=True, more=False)
                else:
                    if tot < 500:
                        apc   = artistAMPageClass(ppp=ppp, tot=tot, redo=False, more=False)
                    else:
                        apc   = artistAMPageClass(ppp=ppp, tot=tot, redo=False, more=True)

                return apc

        return artistAMPageClass()

    def parse(self):
        bsdata = self.bsdata

        artist      = self.getartistAMName()
        url         = self.getartistAMURL()
        ID          = self.getartistAMDiscID(url)
        pages       = self.getartistAMPages()
        profile     = self.getartistAMProfile()
        media       = self.getartistAMMedia()
        mediaCounts = self.getartistAMMediaCounts(media)

        err = [artist.err, url.err, ID.err, pages.err, profile.err, mediaCounts.err, media.err]

        adc = artistAMDataClass(artist=artist, url=url, ID=ID, pages=pages, profile=profile, mediaCounts=mediaCounts, media=media, err=err)

        return adc


import os
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler
from transformers import AdamW, get_linear_schedule_with_warmup
from sklearn.metrics import f1_score
from .camembert_emotion_classif_model import MyCamemBERTModel
from ..data.open_emotion_classif_dataset import open_emotion_classif_dataset
from ..features.tokenizer_emotion_classif import EmotionClassifTokenizer

num_labels = MyCamemBERTModel.num_labels
emotions_dict = MyCamemBERTModel.emotions_dict
emotions = MyCamemBERTModel.emotions
model_dirname = os.path.dirname(__file__)
model_path = os.path.join(model_dirname, '../../res/models/emotion_classif/camembert_trained_on_GoEmotions/camemBERT_finetuned_GoEmotions_french2_dict.pt')
device = 'cuda' if torch.cuda.is_available() else 'cpu'

def accuracy_and_pred_vects_on_each_label(labels, preds, detection_thresholds):
    """
    function used to compute accuracy between labels and preds, also used AND to store predictions
    labels and preds are expected to be pt tensors
    detection_thresholds is expected to be an array of length 7
    """
    accuracies = []
    pred_vects = []

    ones = torch.ones_like(labels[:,0])
    zeros = torch.zeros_like(labels[:,0])

    for i in range(num_labels):

        preds_col_i = torch.where(preds[:,i] > detection_thresholds[i], ones, zeros)
        result_col_i = ((labels[:,i] == preds_col_i).sum().item() / list(preds_col_i.size())[0])
        accuracies.append(result_col_i)

        pred_vects.append((labels[:,i].cpu().numpy(),preds_col_i.cpu().numpy()))

    return accuracies, pred_vects

class EmotionDetector(object):
    """
    Class used to make simple predictions on french tokenized sentences.
    No labels needed here, unless a training is executed.
    """

    model = MyCamemBERTModel()
    model.load_state_dict(torch.load(model_path, map_location='cpu'))

    def __init__(self, detection_thresholds=[0.3 for _ in range(num_labels)]):
        """
        uses detection thresholds for each emotion
        """
        if len(detection_thresholds) > num_labels:
            self.detection_thresholds = detection_thresholds[:num_labels]
        elif len(detection_thresholds) < num_labels:
            self.detection_thresholds = detection_thresholds + [0.3 for _ in range(num_labels-len(detection_thresholds))]
        else:
            self.detection_thresholds = detection_thresholds

    def detect_emotions(self, input_ids):
        """
        makes detections of emotion based on the underlying camembert model, the input_ids tensor (from the tokenizer) fed to the model and the detection thresholds for each emotion.
        Method that really adds the business value to the model.
        """
        self.model.eval()
        sigmoids = nn.Sigmoid()(self.model(input_ids)[0])
        detections = []

        for i in range(sigmoids.size()[0]):
            list_sigmoids_i = sigmoids[i,:].tolist()
            detections_i = {}

            for j in range(sigmoids.size()[1]):
                if list_sigmoids_i[j] > self.detection_thresholds[j]:
                    detections_i[emotions[j]] = list_sigmoids_i[j]

            # rule-based checks to remove potentially incoherent predictions (eg: predicting that a phrase is both marked with sadness and joy)
            if (
                ('anger' in detections_i or 'sadness' in detections_i or 'fear' in detections_i)
                and 'joy' in detections_i
            ):
                print('hello')
                if (('anger' in detections_i and detections_i['anger'] > detections_i['joy'])
                    or ('sadness' in detections_i and detections_i['sadness'] > detections_i['joy'])
                    or ('fear' in detections_i and detections_i['fear'] > detections_i['joy'])
                ):
                    detections_i.pop('joy', None)
                else:
                    detections_i.pop('anger', None)
                    detections_i.pop('fear', None)
                    detections_i.pop('sadness', None)

            detections.append(detections_i)

        return detections

    def train_underlying_model(
        self,
        train_set_filename='train_ekman_french2.tsv',
        dev_set_filename='dev_ekman_french2.tsv',
        max_len=50,
        batch_size=16,
        learning_rate=5e-5,
        warmup_proportion=0.1,
        weight_decay=1e-7,
        adam_epsilon=1e-8,
        nb_epochs=4
    ):
        """
        trains a new instance of the underlying model (with camembert body and top classifier) using provided data of the form of the GoEmotions dataset.
        Warning: because camembert was pre-trained using a French corpus, sentences in the task-specific datasets have to be in French.
        """
        # we instantiate a new model and its tokenizer so that we can train it independently of the one currently used in the static attribute of the class,
        # and replace it after training
        new_camembert_model = MyCamemBERTModel()
        camembert_tokenizer = EmotionClassifTokenizer()
        # we check if a GPU is available for training
        if device == 'cuda':
            new_camembert_model.cuda()

        # we import and wrangle the train dataset
        train_sentences, train_labels = open_emotion_classif_dataset('train', train_set_filename)
        # we tokenize the list of sentences it contains
        train_seqs = camembert_tokenizer.tokenize(train_sentences, max_length=max_len)
        # we build a tensor dataset then a dataloader
        train_set = TensorDataset(train_seqs.input_ids, train_seqs.attention_mask, train_labels)
        train_loader = DataLoader(train_set, sampler=RandomSampler(train_set), batch_size=batch_size)
        # we follow the same process for the dev set if there is one
        dev_loader = None
        if dev_set_filename is not None:
            dev_sentences, dev_labels = open_emotion_classif_dataset('dev', dev_set_filename)
            dev_seqs = camembert_tokenizer.tokenize(dev_sentences, max_length=max_len)
            dev_set = TensorDataset(dev_seqs.input_ids, dev_seqs.attention_mask, dev_labels)
            dev_loader = DataLoader(dev_set, sampler=SequentialSampler(dev_set), batch_size=batch_size)

        # we add a bit of weight decay to regularize only the weights of certain types of layers
        no_decay = ['bias', 'LayerNorm.weight']
        optimizer_grouped_parameters = [      # weight decay applicable only on certains params, not all
            {'params': [p for n, p in new_camembert_model.named_parameters() if not any(nd in n for nd in no_decay)],
            'weight_decay': weight_decay},
            {'params': [p for n, p in new_camembert_model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
        ]
        # we instantiate our optimizer
        optimizer = AdamW(optimizer_grouped_parameters, lr=learning_rate, eps=adam_epsilon)
        # we will use a scheduler
        total_nb_steps = nb_epochs * len(train_loader)    # nb_epochs * nb_batches
        scheduler = get_linear_schedule_with_warmup(
            optimizer,
            num_warmup_steps=int(total_nb_steps * warmup_proportion),
            num_training_steps=total_nb_steps
        )

        # training loop
        training_stats = []

        for _ in range(nb_epochs):

            # TRAINING:
            print("")
            print('======== Epoch {:} / {:} ========'.format(_ + 1, nb_epochs))
            print('Training on '+device+'...')
            new_camembert_model.train()

            total_train_loss = 0
            total_train_accuracy = [0 for i in range(num_labels)]
            total_train_vects = [([],[]) for i in range(num_labels)]

            for step, batch in enumerate(train_loader):

                train_batch_input_ids, train_batch_attention_masks, train_batch_labels = batch[0].to(device), batch[1].to(device), batch[2].to(device)

                optimizer.zero_grad()
                (loss, output) = new_camembert_model(train_batch_input_ids, train_batch_attention_masks, labels=train_batch_labels)
                loss.backward()

                total_train_loss += loss.item()

                preds = torch.sigmoid(output)

                accuracies, pred_vects = accuracy_and_pred_vects_on_each_label(train_batch_labels, preds, self.detection_thresholds)
                total_train_accuracy = [total_train_accuracy[i]+accuracies[i] for i in range(num_labels)]
                total_train_vects = [(total_train_vects[i][0]+[pred_vects[i][0]],total_train_vects[i][1]+[pred_vects[i][1]]) for i in range(num_labels)]

                # Clip the gradient norms to 1.0.
                nn.utils.clip_grad_norm_(new_camembert_model.parameters(), 1.0)

                optimizer.step()
                scheduler.step()

                if step%16 == 0:
                    print("Batch "+str(step+1)+"/"+str(len(train_loader))+" training loss: "+str(loss.item())[:5]+" & accuracies: "+str(accuracies))

            avg_train_loss = total_train_loss / len(train_loader)
            avg_train_accuracy = [total_train_accuracy_i / len(train_loader) for total_train_accuracy_i in total_train_accuracy]
            # f1 scores computation on the training set:
            total_train_f1 = []
            for i in range(num_labels):
                labels_i = np.concatenate(total_train_vects[i][0], axis=0)
                preds_i = np.concatenate(total_train_vects[i][1], axis=0)
                total_train_f1.append(f1_score(labels_i, preds_i))

            print(" Epoch "+str(_+1)+" avg training loss: "+str(avg_train_loss))
            print(" Epoch "+str(_+1)+" avg training accuracy: "+str(avg_train_accuracy))
            print(" Epoch "+str(_+1)+" avg training f1: "+str(total_train_f1))

            # VALIDATION (if a dev set is provided)
            if dev_loader is not None:
                print("")
                print('Validating...')
                avg_dev_loss, avg_dev_accuracy, total_dev_f1 = self.evaluate_model(new_camembert_model, dev_loader)
                print(" Epoch "+str(_+1)+" avg dev loss: "+str(avg_dev_loss))
                print(" Epoch "+str(_+1)+" avg dev accuracy: "+str(avg_dev_accuracy))
                print(" Epoch "+str(_+1)+" avg dev f1: "+str(total_dev_f1))
                training_stats.append(
                    {
                        'epoch': _ + 1,
                        'Training Loss': avg_train_loss,
                        'Training Accur.': avg_train_accuracy,
                        'Training f1': total_train_f1,
                        'Valid. Loss': avg_dev_loss,
                        'Valid. Accur.': avg_dev_accuracy,
                        'Valid. f1': total_dev_f1,
                    }
                )
            else:
                training_stats.append(
                    {
                        'epoch': _ + 1,
                        'Training Loss': avg_train_loss,
                        'Training Accur.': avg_train_accuracy,
                        'Training f1': total_train_f1,
                    }
                )

        # after the training is over, we replace self.model in the class by the newly trained model
        self.model = new_camembert_model

        # we save the weights of the new model
        torch.save(self.model.state_dict(), model_path)

        return training_stats

    def evaluate_other_model(self, my_model, dev_loader):
        """
        evaluates the performance of the underlying camembert model of the class on a torch DataSet dev_set (or test_set) loaded in a torch DataLoader dev_loader (or test_loader)
        returns scores obtained with accuracy and f1_score metrics
        """
        my_model.eval()

        total_dev_accuracy = [0 for i in range(num_labels)]
        total_dev_vects = [([],[]) for i in range(num_labels)]
        total_dev_loss = 0

        for batch in dev_loader:

            dev_batch_input_ids, dev_batch_attention_masks, dev_batch_labels = batch[0].to(device), batch[1].to(device), batch[2].to(device)

            with torch.no_grad():
                (loss, output) = my_model(dev_batch_input_ids, dev_batch_attention_masks, labels=dev_batch_labels)

            total_dev_loss += loss.item()

            preds = torch.sigmoid(output)

            accuracies, pred_vects = accuracy_and_pred_vects_on_each_label(dev_batch_labels, preds, self.detection_thresholds)
            total_dev_accuracy = [total_dev_accuracy[i]+accuracies[i] for i in range(num_labels)]
            total_dev_vects = [(total_dev_vects[i][0]+[pred_vects[i][0]],total_dev_vects[i][1]+[pred_vects[i][1]]) for i in range(num_labels)]

        avg_dev_loss = total_dev_loss / len(dev_loader)
        avg_dev_accuracy = [total_dev_accuracy_i / len(dev_loader) for total_dev_accuracy_i in total_dev_accuracy]
        # f1 scores computation on the dev set:
        total_dev_f1 = []
        for i in range(num_labels):
            labels_i = np.concatenate(total_dev_vects[i][0], axis=0)
            preds_i = np.concatenate(total_dev_vects[i][1], axis=0)
            total_dev_f1.append(f1_score(labels_i, preds_i))

        return avg_dev_loss, avg_dev_accuracy, total_dev_f1


import os
import torch
from torch.autograd import Variable
import torch.nn as nn
import cv2
import sys
sys.path.append("game/")

import random
import numpy as np
from collections import deque

from create_game import GameState
from config import cfg
from net import Net

use_cuda = torch.cuda.is_available()
net = Net()
net.init()
net = net.cuda() if use_cuda else net

criterion = nn.MSELoss().cuda()
optimizer = torch.optim.Adam(net.parameters(), lr=1e-6)

# 開啟遊戲程序
gameState = GameState()

# Replay Memory，類似 list 的儲存容器
D = deque()

# 將遊戲設定為初始狀態，並獲得一個 80x80 的遊戲畫面
do_nothing = np.zeros(cfg._ACTIONS_NUM)
do_nothing[0] = 1
x_t, r_0, terminal = gameState.frame_step(do_nothing)
x_t = cv2.cvtColor(cv2.resize(x_t, (80, 80)), cv2.COLOR_BGR2GRAY)
ret, x_t = cv2.threshold(x_t, 1, 255, cv2.THRESH_BINARY)

# 將初始的遊戲畫面覆蓋成 4 張，作為初始 s_t
s_t = np.stack((x_t, x_t, x_t, x_t), axis=0)

# 設定初始的 epsilon，並準備 trainingepsilon = INITIAL_EPSILON
epsilon = cfg._INITIAL_EPSILON
t = 0

# 設定初始的 epsilon，並準備 trainingepsilon = INITIAL_EPSILON
# 紀錄每輪平均得分的容器
scores = []
all_turn_scores = []

# start game iter
while 'flappy bird' != 'angry bird':
    # 首先，按照 greedy 選擇一個動作
    s = Variable(torch.from_numpy(s_t).type(torch.FloatTensor))
    s = s.cuda() if use_cuda else s
    s = s.view(-1, s.size()[0], s.size()[1], s.size()[2])

    # 取得目前時刻的遊戲畫面，輸入神經網路中
    readout, h_fc1 = net(s)

    # 神經網路產生的輸出為 readout: 選擇每一個行動產生的預期 Q 值
    readout = readout.cpu() if use_cuda else readout
    readout_t = readout.data.numpy()[0]

    # 按照 epsilon greedy 產生小鳥的行動，即以 epsilon 的機率隨機行動
    # 或以 1-epsilon 的機率按照預期輸出最大的 Q 值列出動作
    a_t = np.zeros([cfg._ACTIONS_NUM])
    action_index = 0
    if t % cfg._FRAME_PER_ACTION == 0:
        if random.random() <= epsilon:
            # create random action
            # print("----------------Random Action----------------")
            action_index = random.randrange(cfg._ACTIONS_NUM)
        else:
            action_index = np.argmax(readout_t)

        a_t[action_index] = 1
    else:
        a_t[0] = 1  # do-nothing

    # 降低 epsilon
    if epsilon > cfg._FINAL_EPSILON and t > cfg._OBSERVE:
        epsilon -= (cfg._INITIAL_EPSILON - cfg._FINAL_EPSILON) / cfg._EXPLORE

    # 其次，將選擇好的行動輸入給遊戲引擎，並獲得下一幀狀態
    x_t1_colored, r_t, terminal = gameState.frame_step(a_t)

    # record score in every step
    scores.append(r_t)
    if terminal:
        # 當遊戲結束時，計算本輪的總成績，並將總成績輸入到 all_turn_scores
        all_turn_scores.append(sum(scores))
        scores = []

    # 對遊戲畫面進行處理，進一步變成一張 80x80 的無背景圖
    x_t1 = cv2.cvtColor(cv2.resize(x_t1_colored, (80, 80)), cv2.COLOR_BGR2GRAY)
    ret, x_t1 = cv2.threshold(x_t1, 1, 255, cv2.THRESH_BINARY)
    x_t1 = np.reshape(x_t1, (1, 80, 80))
    s_t1 = np.append(x_t1, s_t[:3, :, :], axis=0)  # (0~2, 80, 80)

    # 產生一個訓練資料對，並將其存入 D
    D.append((s_t, a_t, r_t, s_t1, terminal))
    if len(D) > cfg._REPLAY_MEMORY:
        D.popleft()

    if t > cfg._OBSERVE:
        # 從 D 中隨機取出一筆資料做訓練
        mini_batch = random.sample(D, cfg._BATCH_SIZE)
        optimizer.zero_grad()

        # 將這個 batch 中的變數全部分別存到清單中
        s_j_batch = [d[0] for d in mini_batch]
        a_batch = [d[1] for d in mini_batch]
        r_batch = [d[2] for d in mini_batch]
        s_j1_batch = [d[3] for d in mini_batch]

        # 接下來要經過 s_j1_batch 預測未來的 Q 值
        s = Variable(torch.FloatTensor(np.array(s_j1_batch, dtype=float)))
        s = s.cuda() if use_cuda else s
        readout, h_fc1 = net(s)
        readout = readout.cpu() if use_cuda else readout
        readout_j1_batch = readout.data.numpy()

        # 根據 Q 的預測值、目前回饋 r 與 terminal，更新待訓練的目標函數
        y_batch = []
        for i in range(len(mini_batch)):
            terminal = mini_batch[i][4]

            # 當遊戲結束的時候，則用環境的回饋作為目標
            # 否則用下一狀態的 Q 值加本期的環境回饋
            if terminal:
                y_batch.append(r_batch[i])
            else:
                # readout_j1_batch[i] 儲存了 32 個 [不扇, 扇]
                # np.max 會不分 dimesion 取最大數值
                y_batch.append(r_batch[i] +  cfg._GAMMA * np.max(readout_j1_batch[i]))

        # 開始梯度更新
        y = Variable(torch.FloatTensor(y_batch))
        a = Variable(torch.FloatTensor(a_batch))
        s = Variable(torch.FloatTensor(np.array(s_j_batch, dtype=float)))
        if use_cuda:
            y = y.cuda()
            a = a.cuda()
            s = s.cuda()

        # 計算 s_j_batch 的 Q 值
        readout, h_fc1 = net(s)
        # 被選擇的 action 會為 1，另一動作為 0，因此相乘後會留下被選擇的動作的 Q 值
        readout_action = readout.mul(a).sum(1)
        # readout_action 為依照舊有的 action 計算的預估 Q 值
        # y 則為被 Q(s',a') 影響計算過的目標 Q 值
        loss = criterion(readout_action, y)
        loss.backward()
        optimizer.step()
        if t % 1000 == 0:
            print('loss:', loss)

    # 將狀態更新一次
    s_t = s_t1
    t += 1

    if not os.path.exists('saving_nets'):
        os.mkdir('saving_nets')

    # 每隔 10000 次循環，儲存網路
    if t % 10000 == 0:
        torch.save(net, 'saving_nets/' + cfg._GAME_NAME + '-dqn' + str(t) + '.txt')

    if t <= cfg._OBSERVE:
        state = 'observe'
    elif t > cfg._OBSERVE and t <= cfg._OBSERVE + cfg._EXPLORE:
        state = 'explore'
    else:
        state = 'train'

    if t % 1000 == 0:
        record = "time_step:{};\tstate:{};\tEpsilon:{};\tAction:{};\tReward:{};\tQ_MAX:{:e};\tAll_turn_scors:{};".format(
                t,
                state,
                epsilon,
                action_index,
                r_t,
                np.max(readout_t),
                np.mean(all_turn_scores[-1000:])
                )

        print(record)
        with open('log_file.txt', 'a') as f:
            f.write(record +'\n')