from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import random
import logging
import sys
import numbers
import math
import sklearn
import datetime
import numpy as np
import cv2

sys.path.append(os.path.join(os.path.dirname(__file__), '../'))
sys.path.append(os.path.join(os.path.dirname(__file__), '../common'))
import face_preprocess
import multiprocessing

import mxnet as mx
from mxnet import ndarray as nd
from mxnet import io
from mxnet import recordio

logger = logging.getLogger()
from utils import init_random
from augments.common import common_aug

class FaceImageIter(io.DataIter):
    def __init__(self, batch_size, data_shape,
                 path_imgrecs = None,
                 shuffle=False, aug_list=None, mean = None,
                 rand_mirror = False, cutout = None, crop = None, mask = None, gridmask = None,
                 downsample_back = 0.0, motion_blur = 0.0,
                 mx_model = None,
                 data_names=['data'], label_name='softmax_label', **kwargs):
        super(FaceImageIter, self).__init__()
        assert path_imgrecs
        self.kwargs = kwargs
        if path_imgrecs:
            self.rec_num = len(path_imgrecs)
            self.seq, self.oseq = [], []
            self.imgrec, self.header0, self.imgidx, self.id2range, self.seq_identity = [], [], [], [], []
            for path_imgrec in path_imgrecs:
                logging.info('loading recordio %s...',
                             path_imgrec)
                path_imgidx = path_imgrec[0:-4]+".idx"
                self.imgrec.append( recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r'))  # pylint: disable=redefined-variable-type
                s = self.imgrec[-1].read_idx(0)
                header, _ = recordio.unpack(s)
                if header.flag>0:
                  print('header0 label', header.label)
                  self.header0.append( (int(header.label[0]), int(header.label[1])))
                  #assert(header.flag==1)
                  self.imgidx.append(range(1, int(header.label[0])))
                  self.id2range.append( {} )
                  self.seq_identity.append(range(int(header.label[0]), int(header.label[1])))
                  for identity in self.seq_identity[-1]:
                    s = self.imgrec[-1].read_idx(identity)
                    header, _ = recordio.unpack(s)
                    a,b = int(header.label[0]), int(header.label[1])
                    self.id2range[-1][identity] = (a,b)
                    count = b-a
                  print('id2range', len(self.id2range[-1]))
                else:
                  self.imgidx.append(list(self.imgrec[-1].keys))

                self.seq.append(list(self.imgidx[-1]))
                self.oseq.append(self.imgidx[-1])
                print(len(self.seq[-1]))

        self.iteration = 0
        self.margin_policy = self.kwargs.get('margin_policy', 'step')
        self.use_bgr = self.kwargs.get('use_bgr', False)

        self.mx_model = mx_model

        self.check_data_shape(data_shape)
        if crop is not None:
            crop_h, crop_w = crop.crop_h, crop.crop_w
            data_shape = (data_shape[0], crop_h, crop_w)
        if 'loss_type' in self.kwargs and self.kwargs['loss_type'] == 6:
          self.provide_data = [(data_names[0], (batch_size,) + data_shape), (data_names[1], (batch_size,))]
        else:
          self.provide_data = [(data_names[0], (batch_size,) + data_shape)]
        self.batch_size = batch_size
        self.data_shape = data_shape
        self.shuffle = shuffle
        self.image_size = '%d,%d'%(data_shape[1],data_shape[2])
        self.augs = common_aug(rand_mirror = rand_mirror, cutout = cutout, crop = crop,
                  mask = mask, gridmask = gridmask, downsample_back = downsample_back, motion_blur = motion_blur, mean = mean)

        print('rand_mirror: {}'.format( rand_mirror))
        self.provide_label = [(label_name, (batch_size, self.rec_num))]
        #self.provide_label = [(label_name, (batch_size, self.rec_num))] if self.rec_num > 1 else [(label_name, (batch_size, ))]
        #print(self.provide_label[0][1])
        self.cur = [0] * len(path_imgrecs)
        self.nbatch = 0
        self.is_init = False

    def reset(self):
        """Resets the iterator to the beginning of the data."""
        print('call reset()')
        self.cur = [0] * self.rec_num
        if self.shuffle:
          for i in range(self.rec_num):
            random.shuffle(self.seq[i])
        for i in range(self.rec_num):
          if self.seq[i] is None and self.imgrec[i] is not None:
              self.imgrec[i].reset()

    def num_samples(self, dataset_idx):
      return len(self.seq[dataset_idx])

    def margin(self, iteration):
        if self.margin_policy == 'fixed':
          m= self.kwargs['margin_m']
        elif self.margin_policy == 'step':
          steps = [100000, 300000, 1000000]
          values = [0, 0.35, 0.5]
          for i, step in enumerate(steps):
            if step > iteration:
              m = values[i]
              break
        elif self.margin_policy == 'linear':
            m = self.kwargs['margin_m'] / self.kwargs['max_steps'] * iteration
        return m

    def next_sample(self, dataset_idx):
        """Helper function for reading in next sample."""
        #set total batch size, for example, 1800, and maximum size for each people, for example 45
        if self.seq[dataset_idx] is not None:
          while True:
            if self.cur[dataset_idx] >= len(self.seq[dataset_idx]):
                raise StopIteration
            idx = self.seq[dataset_idx][self.cur[dataset_idx]]
            self.cur[dataset_idx] += 1
            if self.imgrec[dataset_idx] is not None:
              try:
                s = self.imgrec[dataset_idx].read_idx(idx)
              except:
                continue
              header, img = recordio.unpack(s)
              label = header.label
              if not isinstance(label, numbers.Number):
                label = label[0]
              return label, img, None, None
            else:
              label, fname, bbox, landmark = self.imglist[dataset_idx][idx]
              return label, self.read_image(fname), bbox, landmark
        else:
            s = self.imgrec[dataset_idx].read()
            if s is None:
                raise StopIteration
            header, img = recordio.unpack(s)
            label = header.label
            if not isinstance(label, numbers.Number):
              label = label[0]
            return label, img, None, None

    def next(self):
        if not self.is_init:
          self.reset()
          self.is_init = True
        """Returns the next batch of data."""
        #print('in next', self.cur, self.labelcur)
        self.nbatch+=1
        batch_size = self.batch_size
        c, h, w = self.data_shape
        batch_data = nd.empty((batch_size, c, h, w))
        batch_margin = nd.empty((batch_size,))
        if self.provide_label is not None:
          batch_label = nd.empty(self.provide_label[0][1])
        i = 0
        try:
            dataset_idx = 0
            while i < batch_size:
                batch_margin[i] = self.margin(self.iteration)
                _label, s, bbox, landmark = self.next_sample(dataset_idx)
                if(len(self.seq) > 1):
                  label = np.ones([self.rec_num,]) * (-1)
                  label[dataset_idx] = _label
                else:
                  label = _label
                dataset_idx = (dataset_idx + 1) % self.rec_num
                _data = self.imdecode(s)

                _data = self.augs.apply(_data)

                data = [_data]
                try:
                    self.check_valid_image(data)
                except RuntimeError as e:
                    logging.debug('Invalid image, skipping:  %s', str(e))
                    continue

                for datum in data:
                    assert i < batch_size, 'Batch size must be multiples of augmenter output length'
                    #print(datum.shape)
                    batch_data[i][:] = self.postprocess_data(datum)
                    batch_label[i][:] = label
                    i += 1
        except StopIteration:
            if i<batch_size:
                raise StopIteration

        #db = mx.io.DataBatch(data=(batch_data,))
        #self.mx_model.forward(db, is_train=False)
        #net_out = self.mx_model.get_outputs()

        #print(net_out)
        self.iteration += 1
        return io.DataBatch([batch_data], [batch_label], batch_size - i)

    def check_data_shape(self, data_shape):
        """Checks if the input data shape is valid"""
        if not len(data_shape) == 3:
            raise ValueError('data_shape should have length 3, with dimensions CxHxW')
        if not data_shape[0] == 3:
            raise ValueError('This iterator expects inputs to have 3 channels.')

    def check_valid_image(self, data):
        """Checks if the input data is valid"""
        if len(data[0].shape) == 0:
            raise RuntimeError('Data shape is wrong')

    def imdecode(self, s):
        """Decodes a string or byte string to an NDArray.
        See mx.img.imdecode for more details."""
        img = mx.image.imdecode(s) #mx.ndarray
        return img

    def read_image(self, fname):
        """Reads an input image `fname` and returns the decoded raw bytes.

        Example usage:
        ----------
        >>> dataIter.read_image('Face.jpg') # returns decoded raw bytes.
        """
        with open(os.path.join(self.path_root, fname), 'rb') as fin:
            img = fin.read()
        return img

    def augmentation_transform(self, data):
        """Transforms input data with specified augmentation."""
        for aug in self.auglist:
            data = [ret for src in data for ret in aug(src)]

        return data

    def postprocess_data(self, datum):
        """Final postprocessing step before image is loaded into the batch."""
        if self.use_bgr:
          return nd.transpose(datum[:, :, ::-1], axes=(2, 0, 1))
        else:
          return nd.transpose(datum, axes=(2, 0, 1))

class FaceImageIterList(io.DataIter):
  def __init__(self, iter_list):
    assert len(iter_list)>0
    self.provide_data = iter_list[0].provide_data
    self.provide_label = iter_list[0].provide_label
    self.iter_list = iter_list
    self.cur_iter = None

  def reset(self):
    self.cur_iter.reset()

  def next(self):
    self.cur_iter = random.choice(self.iter_list)
    while True:
      try:
        ret = self.cur_iter.next()
      except StopIteration:
        self.cur_iter.reset()
        continue
      return ret

if __name__ == '__main__':
  import config
  train_dataiter = FaceImageIter(
      batch_size           = 32,
      data_shape           = (3, 112, 112),
      path_imgrecs         = ['../datasets/faces_emore/train.rec'],
      shuffle              = True,
      rand_mirror          = True,
      mean                 = None,
      cutout               = config.cutout,
      crop                 = config.crop,
      mask                 = config.mask,
      gridmask             = config.gridmask,
      data_names           = ['data'],
      downsample_back      = config.config.downsample_back,
      motion_blur          = config.config.motion_blur,
      use_bgr              = config.config.use_bgr
  )

  batch = train_dataiter.next()
  data = batch.data[0].asnumpy()
  for i in range(32):
    img_data = data[i, ...].transpose([1, 2, 0])[:, :, ::-1]
    cv2.imwrite('temp/%d.png' % i, img_data)


# Copyright (c) Microsoft. All rights reserved.

# Licensed under the MIT license. See LICENSE.md file in the project root
# for full license information.
# ==============================================================================
"""Deep Q-learning and its variants."""

import math

import cntk as C
import numpy as np

import ast

from .agent import AgentBaseClass
from .shared.cntk_utils import huber_loss
from .shared.models import Models
from .shared.qlearning_parameters import QLearningParameters
from .shared.replay_memory import ReplayMemory

class QLearning(AgentBaseClass):
    """
    Q-learning agent.

    Including:
    - DQN https://storage.googleapis.com/deepmind-media/dqn/DQNNaturePaper.pdf
    - Prioritized Experience Replay https://arxiv.org/pdf/1511.05952.pdf
    - Dueling Network https://arxiv.org/pdf/1511.06581.pdf
    - Double Q Learning https://arxiv.org/pdf/1509.06461.pdf
    """

    def __init__(self, config_filename, o_space, a_space):
        """Constructor for Q learning algorithm.

        Widely known as DQN. Use either predefined neural network structure
        (see models.py) or customized network (see customized_models.py).

        Args:
            config_filename: configure file specifying training details.
            o_space: observation space, gym.spaces.tuple_space.Tuple is not
                supported.
            a_space: action space, limits to gym.spaces.discrete.Discrete.
        """
        super(QLearning, self).__init__(o_space, a_space)

        self._parameters = QLearningParameters(config_filename)

        # Create preprocessor.
        if self._parameters.preprocessing:
            try:
                preproc = self._import_method(self._parameters.preprocessing)
                self._preprocessor = preproc(
                    self._shape_of_inputs,
                    *ast.literal_eval(self._parameters.preprocessing_args))
            except ValueError:
                raise ValueError(
                    'Unknown preprocessing method: "{0}"'
                    '\n'.format(self._parameters.preprocessing))

        # Set up the Q-function.
        shape_of_inputs = self._shape_of_inputs \
            if self._preprocessor is None \
            else self._preprocessor.output_shape()
        if self._parameters.q_representation == 'dqn':
            model = Models.feedforward_network(
                shape_of_inputs,
                self._num_actions,
                self._parameters.hidden_layers,
                huber_loss if self._parameters.use_error_clipping else None)
        elif self._parameters.q_representation == 'dueling-dqn':
            model = Models.dueling_network(
                shape_of_inputs,
                self._num_actions,
                self._parameters.hidden_layers,
                huber_loss if self._parameters.use_error_clipping else None)
        else:
            try:
                model_definition_function = self._import_method(
                    self._parameters.q_representation)
                model = model_definition_function(
                    shape_of_inputs,
                    self._num_actions,
                    huber_loss if self._parameters.use_error_clipping else None)
            except ValueError:
                raise ValueError(
                    'Unknown representation for Q-learning: "{0}"'
                    '\n'.format(self._parameters.q_representation))

        self._q = model['f']
        self._input_variables = model['inputs']
        self._output_variables = model['outputs']
        if self._parameters.use_prioritized_replay:
            self._weight_variables = \
                C.ops.input_variable(shape=(1,), dtype=np.float32)
            self._loss = model['loss'] * self._weight_variables
        else:
            self._loss = model['loss']

        minibatch_size = int(self._parameters.minibatch_size)
        # If gradient_clipping_threshold_per_sample is inf, gradient clipping
        # will not be performed. Set gradient_clipping_with_truncation to False
        # to clip the norm.
        # TODO: allow user to specify learner through config file.
        opt = C.learners.adam(
            self._q.parameters,
            C.learners.learning_rate_schedule(
                self._parameters.initial_eta, C.learners.UnitType.sample),
            use_mean_gradient=True,
            momentum=C.learners.momentum_schedule(self._parameters.momentum),
            variance_momentum=C.learners.momentum_schedule(0.999),
            gradient_clipping_threshold_per_sample=
                self._parameters.gradient_clipping_threshold,
            gradient_clipping_with_truncation=False)
        self._trainer = C.train.trainer.Trainer(
            self._q, (self._loss, None), opt)

        # Initialize target Q.
        self._target_q = self._q.clone('clone')

        # Initialize replay memory.
        self._replay_memory = ReplayMemory(
            self._parameters.replay_memory_capacity,
            self._parameters.use_prioritized_replay)

        print('Parameterized Q-learning agent using neural networks '
              '"{0}" with {1} actions.\n'
              ''.format(self._parameters.q_representation,
                        self._num_actions))

        self.episode_count = 0
        self.step_count = 0

    def start(self, state):
        """
        Start a new episode.

        Args:
            state (object): observation provided by the environment.

        Returns:
            action (int): action choosen by agent.
            debug_info (dict): auxiliary diagnostic information.
        """
        if self._preprocessor is not None:
            self._preprocessor.reset()

        self._adjust_exploration_rate()
        self._last_state = self._preprocess_state(state)
        self._last_action, action_behavior = \
            self._choose_action(self._last_state)
        self.episode_count += 1
        return self._last_action, {
            'action_behavior': action_behavior,
            'epsilon': self._epsilon}

    def step(self, reward, next_state):
        """
        Observe one transition and choose an action.

        Args:
            reward (float) : amount of reward returned after previous action.
            next_state (object): observation provided by the environment.

        Returns:
            action (int): action choosen by agent.
            debug_info (dict): auxiliary diagnostic information.
        """
        next_encoded_state = self._preprocess_state(next_state)
        priority = self._compute_priority(
            self._last_state, self._last_action, reward, next_encoded_state)
        self._replay_memory.store(
            self._last_state,
            self._last_action,
            reward,
            next_encoded_state,
            priority)
        self.step_count += 1

        # Update Q every self._parameters.q_update_frequency
        self._update_q_periodically()

        self._adjust_exploration_rate()
        self._last_state = next_encoded_state
        self._last_action, action_behavior = self._choose_action(
            self._last_state)
        return self._last_action, {
            'action_behavior': action_behavior,
            'epsilon': self._epsilon}

    def end(self, reward, next_state):
        """
        Last observed reward/state of the episode (which then terminates).

        Args:
            reward (float) : amount of reward returned after previous action.
            next_state (object): observation provided by the environment.
        """
        priority = self._compute_priority(
            self._last_state, self._last_action, reward, None)
        self._replay_memory.store(
            self._last_state,
            self._last_action,
            reward,
            None,
            priority)
        self.step_count += 1

        # Update Q every self._parameters.q_update_frequency
        self._update_q_periodically()

    def set_as_best_model(self):
        """Copy current model to best model."""
        self._best_model = self._q.clone('clone')

    def enter_evaluation(self):
        """Setup before evaluation."""
        self._epsilon = 0

    def _adjust_learning_rate(self):
        if self._parameters.initial_eta != self._parameters.eta_minimum:
            eta = self._parameters.eta_minimum + max(
                0,
                (self._parameters.initial_eta - self._parameters.eta_minimum) *
                (1 - float(self.step_count)/self._parameters.eta_decay_step_count))

            self._trainer.parameter_learners[0].reset_learning_rate(
                C.learners.learning_rate_schedule(
                    eta, C.learners.UnitType.sample))

    def _adjust_exploration_rate(self):
        self._epsilon = self._parameters.epsilon_minimum + max(
            0,
            (self._parameters.initial_epsilon - self._parameters.epsilon_minimum) *
            (1 - float(self.step_count)/self._parameters.epsilon_decay_step_count))

    def _choose_action(self, state):
        """
        Epsilon greedy policy.

        Args:
            state (object): observation seen by agent, which can be different
                from what is provided by the environment. The difference comes
                from preprcessing.

        Returns:
            action (int): action choosen by agent.
            debug_info (str): auxiliary diagnostic information.
        """
        if self.step_count < self._parameters.replay_start_size or \
                np.random.uniform(0, 1) < self._epsilon:
            return np.random.randint(self._num_actions), 'RANDOM'
        else:
            return np.argmax(self._evaluate_q(self._q, state)), 'GREEDY'

    def save(self, filename):
        """Save model to file."""
        self._best_model.save(filename)

    def save_parameter_settings(self, filename):
        """Save parameter settings to file."""
        self._parameters.save(filename)

    def _evaluate_q(self, model, state, action=None):
        """
        Evaluate Q[state, action].

        If action is None, return values for all actions.
        Args:
            state (object): observation seen by agent, which can be different
                from what is provided by the environment. The difference comes
                from preprcessing.
            action (int): action choosen by agent.
        """
        q = np.squeeze(model.eval({model.arguments[0]: [state]}))
        if action is None:
            return q
        else:
            return q[action]

    def _update_q_periodically(self):
        if self.step_count < self._parameters.replay_start_size or \
                self.step_count % self._parameters.q_update_frequency != 0:
            return

        self._adjust_learning_rate()
        for i in range(self._parameters.replays_per_update):
            self._replay_and_update()

        # Clone target network periodically.
        if self.step_count % \
                self._parameters.target_q_update_frequency == 0:
            self._target_q = self._q.clone('clone')

    def _replay_and_update(self):
        """Perform one minibatch update of Q."""
        input_values = []
        output_values = []
        if self._parameters.use_prioritized_replay:
            # importance sampling weights.
            weight_values = []

        minibatch = self._replay_memory.sample_minibatch(
            self._parameters.minibatch_size)
        for index_transition_pair in minibatch:
            input_value = index_transition_pair[1].state

            # output_value is the same for all actions except last_action.
            output_value = self._evaluate_q(
                self._q, index_transition_pair[1].state)
            td_err = self._compute_td_err(
                index_transition_pair[1].state,
                index_transition_pair[1].action,
                index_transition_pair[1].reward,
                index_transition_pair[1].next_state)
            output_value[index_transition_pair[1].action] += td_err

            input_values.append(input_value)
            output_values.append(output_value)

            if self._parameters.use_prioritized_replay:
                weight_values.append(math.pow(
                    index_transition_pair[1].priority,
                    -self._parameters.priority_beta))

        if self._parameters.use_prioritized_replay:
            w_sum = sum(weight_values)
            weight_values = [[w / w_sum] for w in weight_values]
            self._trainer.train_minibatch(
                {
                    self._input_variables: np.array(input_values).astype(
                        np.float32),
                    self._output_variables: np.array(output_values).astype(
                        np.float32),
                    self._weight_variables: np.array(weight_values).astype(
                        np.float32)
                })

            # Update replay priority.
            position_priority_map = {}
            for index_transition_pair in minibatch:
                position_priority_map[index_transition_pair[0]] = \
                    self._compute_priority(
                        index_transition_pair[1].state,
                        index_transition_pair[1].action,
                        index_transition_pair[1].reward,
                        index_transition_pair[1].next_state)

            self._replay_memory.update_priority(position_priority_map)
        else:
            self._trainer.train_minibatch(
                {
                    self._input_variables: np.array(input_values).astype(
                        np.float32),
                    self._output_variables: np.array(output_values).astype(
                        np.float32)
                })

    def _compute_td_err(self, state, action, reward, next_state):
        td_err = reward
        if next_state is not None:
            if self._parameters.double_q_learning:
                td_err += self._parameters.gamma * \
                    self._evaluate_q(
                        self._target_q,
                        next_state,
                        np.argmax(self._evaluate_q(self._q, next_state)))
            else:
                td_err += self._parameters.gamma * np.max(
                    self._evaluate_q(self._target_q, next_state))
        td_err -= self._evaluate_q(self._q, state, action)
        return td_err

    def _compute_priority(self, state, action, reward, next_state):
        priority = None
        if self._parameters.use_prioritized_replay:
            priority = math.pow(
                math.fabs(self._compute_td_err(
                    state, action, reward, next_state))
                + self._parameters.priority_epsilon,
                self._parameters.priority_alpha)
        return priority


import os
from typing import Union

import numpy as np

import redback.get_data.directory
from redback.transient.transient import OpticalTransient

dirname = os.path.dirname(__file__)

class Kilonova(OpticalTransient):

    DATA_MODES = ['flux_density', 'magnitude', 'luminosity']

    def __init__(
            self, name: str, data_mode: str = 'magnitude', time: np.ndarray = None, time_err: np.ndarray = None,
            time_mjd: np.ndarray = None, time_mjd_err: np.ndarray = None, time_rest_frame: np.ndarray = None,
            time_rest_frame_err: np.ndarray = None, Lum50: np.ndarray = None, Lum50_err: np.ndarray = None,
            flux_density: np.ndarray = None, flux_density_err: np.ndarray = None, magnitude: np.ndarray = None,
            magnitude_err: np.ndarray = None, redshift: float = np.nan, photon_index: float = np.nan,
            bands: np.ndarray = None, system: np.ndarray = None, active_bands: Union[np.ndarray, str] = 'all',
            use_phase_model: bool = False, **kwargs: None) -> None:
        """This is a general constructor for the Kilonova class. Note that you only need to give data corresponding to
        the data mode you are using. For luminosity data provide times in the rest frame, if using a phase model
        provide time in MJD, else use the default time (observer frame).

        :param name: Name of the transient.
        :type name: str
        :param data_mode: Data mode. Must be one from `Afterglow.DATA_MODES`.
        :type data_mode: str, optional
        :param time: Times in the observer frame.
        :type time: np.ndarray, optional
        :param time_err: Time errors in the observer frame.
        :type time_err: np.ndarray, optional
        :param time_mjd: Times in MJD. Used if using phase model.
        :type time_mjd: np.ndarray, optional
        :param time_mjd_err: Time errors in MJD. Used if using phase model.
        :type time_mjd_err: np.ndarray, optional
        :param time_rest_frame: Times in the rest frame. Used for luminosity data.
        :type time_rest_frame: np.ndarray, optional
        :param time_rest_frame_err: Time errors in the rest frame. Used for luminosity data.
        :type time_rest_frame_err: np.ndarray, optional
        :param Lum50: Luminosity values.
        :type Lum50: np.ndarray, optional
        :param Lum50_err: Luminosity error values.
        :type Lum50_err: np.ndarray, optional
        :param flux: Flux values.
        :type flux: np.ndarray, optional
        :type flux_err: np.ndarray, optional
        :param flux_err: Flux error values.
        :param flux_density:Flux density values.
        :type flux_density: np.ndarray, optional
        :param flux_density_err: Flux density error values.
        :type flux_density_err: np.ndarray, optional
        :param magnitude: Magnitude values for photometry data.
        :type magnitude: np.ndarray, optional
        :param magnitude_err: Magnitude error values for photometry data.
        :type magnitude_err: np.ndarray, optional
        :param redshift: Redshift value. Will be read from the metadata table if not given.
        :type redshift: float
        :param photon_index: Photon index value. Will be read from the metadata table if not given.
        :type photon_index: float
        :param use_phase_model: Whether we are using a phase model.
        :type use_phase_model: bool
        :param frequency: Array of band frequencies in photometry data.
        :type frequency: np.ndarray, optional
        :param system: System values.
        :type system: np.ndarray, optional
        :param bands: Band values.
        :type bands: np.ndarray, optional
        :param active_bands: List or array of active bands to be used in the analysis. Use all available bands if 'all' is given.
        :type active_bands: Union[list, np.ndarray]
        :param kwargs: Additional callables:
                        bands_to_frequency: Conversion function to convert a list of bands to frequencies. Use
                        redback.utils.bands_to_frequency if not given.
        :type kwargs: None
        """

        super().__init__(time=time, time_err=time_err, time_rest_frame=time_rest_frame, time_mjd=time_mjd,
                         time_mjd_err=time_mjd_err, time_rest_frame_err=time_rest_frame_err, Lum50=Lum50,
                         Lum50_err=Lum50_err, flux_density=flux_density, flux_density_err=flux_density_err,
                         magnitude=magnitude, magnitude_err=magnitude_err, data_mode=data_mode, name=name, bands=bands,
                         system=system, active_bands=active_bands, use_phase_model=use_phase_model, redshift=redshift,
                         photon_index=photon_index, **kwargs)
        self.directory_structure = redback.get_data.directory.open_access_directory_structure(
            transient=name, transient_type="kilonova")
        self._set_data()


import numpy as np
import cv2
import paddlehub as hub
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '0'
# 创建VideoCapture，传入0即打开系统默认摄像头
vc = cv2.VideoCapture(0)

# 创建module
module = hub.Module(name='ace2p')

n = 0
count = 0
sum_count = 40
colors = np.array([
    # [139, 0, 0],
    # [255, 255, 0],  # yellow
    [255, 215, 0],  # gold
    # [255, 105, 180],  # hotpink
    # [255, 0, 255],  # fuchsia
    [224, 255, 255],  # lightcyan
    [139, 69, 19],  # saddlebrown
    [128, 128, 0],  # olive
    [124, 252, 0],  # lawngreen
    # [0, 0, 255],  # blue
    [0, 128, 0],  # green
])
len_colors = colors.shape[0]

while True:
    # 读取一帧，read()方法是其他两个类方法的结合，具体文档
    # ret为bool类型，指示是否成功读取这一帧
    ret, frame = vc.read()
    frame = cv2.flip(frame, 1)
    # 分割
    img = module.segmentation(images=[frame], use_gpu=True)[0]['data']
    img[img != 2] = 1
    img[img == 2] = 0

    img = img[:, :, np.newaxis]
    img = np.repeat(img, repeats=3, axis=2)

    frame = np.multiply(frame, img)

    img = np.abs(img - 1)
    img = np.multiply(img, colors[n])

    frame += img.astype('uint8')
    # 若没有按下q键，则每1毫秒显示一帧
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break
    cv2.imshow('frame', frame)
    count += 1
    if count == sum_count:
        count = 0
        n += 1
        if n == 5:
            sum_count = 60
        else:
            sum_count = 40
        if n == len_colors:
            n = 0
# 所有操作结束后不要忘记释放
vc.release()
cv2.destroyAllWindows()


"""
This file is used to test the model using the AI Challenger dataset.

@ Author: Yu Sun. vxallset@outlook.com

@ Date created: Jun 04, 2019

@ Last modified: Jun 27, 2019

"""
import numpy as np
import tensorflow as tf
from hrnet import *
import dataset
from heatmap import *
import time
import os

from functools import reduce
from operator import mul

def get_num_params():
    num_params = 0
    for variable in tf.trainable_variables():
        shape = variable.get_shape()
        num_params += reduce(mul, [dim.value for dim in shape], 1)
    return num_params

def main(device_option='/gpu:0'):
    batch_size = 1
    num_epochs = 10
    image_numbers = 378352
    #image_numbers = 4500

    root_path = os.getcwd()[:-3]

    datasetname = os.path.join(root_path, 'dataset/train.tfrecords')
    model_folder = os.path.join(root_path, 'models/')
    modelfile = os.path.join(root_path, 'models/epoch2.ckpt-567528')

    global_step = tf.Variable(0, trainable=False)

    image_name, image, keypoints_ver, keypoints_hor, keypoints_class = dataset.input_batch(
        datasetname=datasetname, batch_size=batch_size, num_epochs=num_epochs)

    input_images = tf.placeholder(tf.float32, [None, 256, 192, 3])
    ground_truth = tf.placeholder(tf.float32, [None, 256, 192, 14])

    input_images = tf.cast(input_images / 255.0, tf.float32, name='change_type')
    net_output = HRNet(input=input_images)
    loss = compute_loss(net_output=net_output, ground_truth=ground_truth)

    saver = tf.train.Saver()
    # os.environ['CUDA_VISIBLE_DEVICES'] = ''

    with tf.Session() as sess:
        with tf.device(device_option):
            sess.run(tf.global_variables_initializer())
            saver.restore(sess=sess, save_path=modelfile)
            #print(get_num_params())

            writer = tf.summary.FileWriter('../log/', sess.graph)
            start_time = time.time()
            try:
                distances = 0
                classes = 0
                for step in range(int(image_numbers / batch_size)):
                    _img, _kp_ver, _kp_hor, _kp_class = sess.run(
                        [image, keypoints_ver, keypoints_hor, keypoints_class])
                    _gt = calculate_groundtruth_heatmap(_kp_ver, _kp_hor, _kp_class)

                    tloss, tnet_output = sess.run([loss, net_output],
                                                  feed_dict={input_images: _img, ground_truth: _gt})

                    prediction = decode_output(tnet_output, threshold=0.001)
                    gt_all = np.zeros((batch_size, 14*3))
                    gt_all[:, ::3] = _kp_ver
                    gt_all[:, 1::3] = _kp_hor
                    gt_all[:, 2::3] = _kp_class
                    distance = calculate_distance(prediction, gt_all)

                    if step == 0:
                        distances = distance
                        classes = _kp_class
                    elif step == 1000:
                        np.save('distances.npy', distances)
                        np.save('classes.npy', classes)
                        break
                    else:
                        distances = np.append(distances, distance, axis=0)
                        classes = np.append(classes, _kp_class, axis=0)

                    timgs = decode_pose(_img, tnet_output, threshold=0.001)
                    for i in range(batch_size):
                        io.imsave('../test_img/step{}_i_{}.jpg'.format(step, i), timgs[i])
                    print('Step = {:>6}/{:>6}, loss = {:.6f}, time = {}'
                          .format(step, int(image_numbers / batch_size), tloss,
                                  time.time() - start_time))
                    print('---------------------------------------------------------------------------------')

            except tf.errors.OutOfRangeError:
                print('End testing...')
            finally:
                total_time = time.time() - start_time
                print('Running time: {} s'.format(total_time))
                print('Done!')

if __name__ == '__main__':
    main(device_option='/gpu:0')


from __future__ import print_function, absolute_import

import os
import numpy as np
import json
import random
import math

import torch
import torch.utils.data as data
import random

from utils.imutils2 import *
from utils.transforms import *
import torchvision.transforms as transforms

import scipy.io as sio
import scipy.misc

# get the video frames
# two patches in the future frame, one is center, the other is one of the 8 patches around

class DavisSet(data.Dataset):
    def __init__(self, params, is_train=True):

        self.filelist = params['filelist']
        self.batchSize = params['batchSize']
        self.imgSize = params['imgSize']
        self.cropSize = params['cropSize']
        self.cropSize2 = params['cropSize2']
        self.videoLen = params['videoLen']
        self.predFrames = params['predFrames'] # 4
        self.sideEdge = params['sideEdge'] # 64

        # prediction distance, how many frames far away
        self.offset = params['offset']
        # gridSize = 3
        # self.gridSize = params['gridSize']

        self.is_train = is_train

        f = open(self.filelist, 'r')
        self.jpgfiles = []
        self.lblfiles = []

        for line in f:
            rows = line.split()
            jpgfile = rows[0]
            lblfile = rows[1]

            self.jpgfiles.append(jpgfile)
            self.lblfiles.append(lblfile)

        f.close()

    def cropimg(self, img, offset_x, offset_y, cropsize):

        img = im_to_numpy(img)
        cropim = np.zeros([cropsize, cropsize, 3])
        cropim[:, :, :] = img[offset_y: offset_y + cropsize, offset_x: offset_x + cropsize, :]
        cropim = im_to_torch(cropim)

        return cropim

    def __getitem__(self, index):

        folder_path = self.jpgfiles[index]
        label_path = self.lblfiles[index]

        imgs = []
        lbls = []
        patches = []
        target_imgs = []

        frame_num = len(os.listdir(folder_path)) + self.videoLen

        mean=[0.485, 0.456, 0.406]
        std=[0.229, 0.224, 0.225]

        for i in range(frame_num):
            if i < self.videoLen:
                img_path = folder_path + "{:05d}.jpg".format(0)
                lbl_path = label_path + "{:05d}.png".format(0)
            else:
                img_path = folder_path + "{:05d}.jpg".format(i - self.videoLen)
                lbl_path = label_path + "{:05d}.png".format(i - self.videoLen)

            img = load_image(img_path)  # CxHxW
            ht, wd = img.size(1), img.size(2)
            newh, neww = ht, wd

            if ht <= wd:
                ratio  = 1.0 #float(wd) / float(ht)
                # width, height
                img = resize(img, int(self.imgSize * ratio), self.imgSize)
                newh = self.imgSize
                neww = int(self.imgSize * ratio)
            else:
                ratio  = 1.0 #float(ht) / float(wd)
                # width, height
                img = resize(img, self.imgSize, int(self.imgSize * ratio))
                newh = int(self.imgSize * ratio)
                neww = self.imgSize

            if i == 0:
                imgs = torch.Tensor(frame_num, 3, newh, neww)

            img = color_normalize(img, mean, std)
            imgs[i] = img
            lblimg  = scipy.misc.imread(lbl_path)
            lblimg  = scipy.misc.imresize( lblimg, (newh, neww), 'nearest' )

            lbls.append(lblimg.copy())

        gridx = 0
        gridy = 0

        for i in range(frame_num):

            img = imgs[i]
            ht, wd = img.size(1), img.size(2)
            newh, neww = ht, wd

            sideEdge = self.sideEdge

            gridy = int(newh / sideEdge)
            gridx = int(neww / sideEdge)

            # img = im_to_numpy(img)
            # target_imgs.append(img)

            for yid in range(gridy):
                for xid in range(gridx):

                    patch_img = img[:, yid * sideEdge: yid * sideEdge + sideEdge, xid * sideEdge: xid * sideEdge + sideEdge].clone()

                    patches.append(patch_img)

        countPatches = frame_num * gridy * gridx
        patchTensor = torch.Tensor(countPatches, 3, self.cropSize2, self.cropSize2)

        for i in range(countPatches):
            patchTensor[i, :, :, :] = patches[i]

        # for i in range(len(imgs)):
        #     imgs[i] = color_normalize(imgs[i], mean, std)

        patchTensor = patchTensor.view(frame_num, gridy * gridx, 3, self.cropSize2, self.cropSize2)

        # Meta info
        meta = {'folder_path': folder_path, 'gridx': gridx, 'gridy': gridy}

        lbls_tensor = torch.Tensor(len(lbls), newh, neww, 3)
        for i in range(len(lbls)):
            lbls_tensor[i] = torch.from_numpy(lbls[i])

        return imgs, patchTensor, lbls_tensor, meta

    def __len__(self):
        return len(self.jpgfiles)


from tensorflow.keras.preprocessing.image import img_to_array
import imutils
import cv2
from tensorflow.keras.models import load_model
import numpy as np
import face_recognition
import imutils
import pickle
import time
import os

def imotion():

    detection_model_path = 'haarcascade_frontalface_default.xml'
    emotion_model_path = '_mini_XCEPTION.106-0.65.hdf5'

    face_detection = cv2.CascadeClassifier(detection_model_path)
    data = pickle.loads(open('face_enc', "rb").read())
    emotion_classifier = load_model(emotion_model_path, compile=False)
    EMOTIONS = ["angry" ,"disgust","scared", "happy", "sad", "surprised",
                "neutral"]

    print("Streaming started")
    cv2.namedWindow('your_face')
    camera = cv2.VideoCapture(0)
    while True:

        frame = camera.read()[1]
        #reading the frame
        frame = imutils.resize(frame,width=400)
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = face_detection.detectMultiScale(gray,scaleFactor=1.1,minNeighbors=5,minSize=(30,30),flags=cv2.CASCADE_SCALE_IMAGE)

        canvas = np.zeros((250, 300, 3), dtype="uint8")
        frameClone = frame.copy()
        if len(faces) > 0:
            faces = sorted(faces, reverse=True,
            key=lambda x: (x[2] - x[0]) * (x[3] - x[1]))[0]
            (fX, fY, fW, fH) = faces
            # Extract the ROI of the face from the grayscale image, resize it to a fixed 48x48 pixels, and then prepare
            # the ROI for classification via the CNN
            roi = gray[fY:fY + fH, fX:fX + fW]
            roi = cv2.resize(roi, (48, 48))
            roi = roi.astype("float") / 255.0
            roi = img_to_array(roi)
            roi = np.expand_dims(roi, axis=0)

            preds = emotion_classifier.predict(roi)[0]
            emotion_probability = np.max(preds)
            label = EMOTIONS[preds.argmax()]

        for (i, (emotion, prob)) in enumerate(zip(EMOTIONS, preds)):

            # construct the label text
            text = "{}: {:.2f}%".format(emotion, prob * 100)
            w = int(prob * 300)
            cv2.rectangle(canvas, (7, (i * 35) + 5),(w, (i * 35) + 35), (0, 0, 255), -1)
            cv2.putText(canvas, text, (10, (i * 35) + 23),cv2.FONT_HERSHEY_SIMPLEX, 0.45,(255, 255, 255), 2)
            cv2.putText(frameClone, label, (fX, fY - 10),cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 0, 255), 2)
            cv2.rectangle(frameClone, (fX, fY), (fX + fW, fY + fH),(0, 0, 255), 2)

        cv2.imshow('your_face', frameClone)
        cv2.imshow("Probabilities", canvas)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    camera.release()
    cv2.destroyAllWindows()

def recognition():
    #find path of xml file containing haarcascade file
    cascPathface = os.path.dirname(
        cv2.__file__) + "/data/haarcascade_frontalface_alt2.xml"
    # load the harcaascade in the cascade classifier
    faceCascade = cv2.CascadeClassifier(cascPathface)
    # load the known faces and embeddings saved in last file
    data = pickle.loads(open('face_enc', "rb").read())

    print("Streaming started")
    video_capture = cv2.VideoCapture(0)
    # loop over frames from the video file stream
    while True:
        # grab the frame from the threaded video stream
        ret, frame = video_capture.read()
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = faceCascade.detectMultiScale(gray,
                                         scaleFactor=1.1,
                                         minNeighbors=5,
                                         minSize=(60, 60),
                                         flags=cv2.CASCADE_SCALE_IMAGE)

        # convert the input frame from BGR to RGB
        rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        # the facial embeddings for face in input
        encodings = face_recognition.face_encodings(rgb)
        names = []
        # loop over the facial embeddings incase
        # we have multiple embeddings for multiple fcaes
        for encoding in encodings:
            #Compare encodings with encodings in data["encodings"]
            #Matches contain array with boolean values and True for the embeddings it matches closely
            #and False for rest
            matches = face_recognition.compare_faces(data["encodings"],
                                                     encoding)
            #set name =inknown if no encoding matches
            name = "Unknown"
            # check to see if we have found a match
            if True in matches:
                #Find positions at which we get True and store them
                matchedIdxs = [i for (i, b) in enumerate(matches) if b]
                counts = {}
                # loop over the matched indexes and maintain a count for
                # each recognized face face
                for i in matchedIdxs:
                    #Check the names at respective indexes we stored in matchedIdxs
                    name = data["names"][i]
                    #increase count for the name we got
                    counts[name] = counts.get(name, 0) + 1
                    #set name which has highest count
                    name = max(counts, key=counts.get)

            # update the list of names
            names.append(name)
            # loop over the recognized faces
            for ((x, y, w, h), name) in zip(faces, names):
                # rescale the face coordinates
                # draw the predicted face name on the image
                cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
                cv2.putText(frame, name, (x, y), cv2.FONT_HERSHEY_SIMPLEX,
                            0.75, (0, 255, 0), 2)

        cv2.imshow("Frame", frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    video_capture.release()
    cv2.destroyAllWindows()

imotion()


import numpy as np
import os
import pandas as pd
from network.utils import bbox_iou
import pickle
from tqdm import tqdm
import argparse
import multiprocessing
from configs.default import __C, cfg_from_file
from dataset.Parsers.structures import *

class GTSingleParser:
    def __init__(self, folder,
                 min_visibility,
                 forward_frames,
                 frame_stride,
                 tube_thre,
                 loose,
                 height_clamp):
        # 1. get the gt path and image folder
        split_path = folder.split('/')
        if folder[0] == '/':
            jta_root = '/' + os.path.join(*split_path[:-3])
        else:
            jta_root = os.path.join(*split_path[:-3])
        type = split_path[-2]
        video_name = split_path[-1]
        gt_file_path = os.path.join(jta_root, 'gt_' + str(loose) + '_' + str(min_visibility) + '_' + str(height_clamp), type, video_name, 'gt.txt')
        # gt_file_path = os.path.join(folder, 'gt/gt.txt')

        self.folder = folder
        self.forward_frames = forward_frames
        self.tube_thre = tube_thre
        self.min_visibility = min_visibility
        self.frame_stride = frame_stride

        self.tube_res_path = os.path.join(jta_root,
                                          'tubes_' + str(self.forward_frames) + '_' + str(
                                              self.frame_stride) + '_' + str(self.min_visibility),
                                          type,
                                          video_name)

        try:
            os.makedirs(self.tube_res_path)
        except:
            pass

        # 2. read the gt data
        gt_file = pd.read_csv(gt_file_path, header=None)
        gt_file = gt_file[gt_file[6] == 1]  # human class
        gt_file = gt_file[gt_file[8] > min_visibility]
        gt_group = gt_file.groupby(0)
        gt_group_keys = gt_group.indices.keys()
        self.max_frame_index = max(gt_group_keys)
        # 3. update tracks
        self.tracks = Tracks()
        self.recorder = {}
        for key in gt_group_keys:
            det = gt_group.get_group(key).values
            ids = np.array(det[:, 1]).astype(int)
            det = np.array(det[:, 2:6])
            det[:, 2:4] += det[:, :2]

            self.recorder[key - 1] = list()
            # 3.1 update tracks
            for id, d in zip(ids, det):
                node = Node(d, key - 1)
                track_index, node_index = self.tracks.add_node(node, id)
                self.recorder[key - 1].append((track_index, node_index))

    def bbox2tube(self, track, mid_id, direction, pos_in_video, thre):
        def get_true_z(mid_node, end_node):
            return end_node.frame_id - mid_node.frame_id

        def get_inter_box(start_box, end_box, inter_id, end_id):
            return start_box * (end_id - inter_id) / end_id + end_box * inter_id / end_id

        mid_node = track.get_node_by_index(mid_id)
        mid_box = mid_node.box
        inter_boxes = []

        z = 1 if direction == 'front' else -1
        if mid_id + z >= len(track.nodes) or mid_id + z < 0:
            return np.array([0, 0, 0, 0, 0])

        true_z = get_true_z(mid_node, track.get_node_by_index(mid_id + z))

        max_len = (self.forward_frames * 2 - 1) * self.frame_stride + 1

        while -1 * pos_in_video <= true_z < max_len - pos_in_video:
            iou_total = 0
            end_node = track.get_node_by_index(mid_id + z)
            end_box = end_node.box
            for i, gt_box in enumerate(inter_boxes):
                iou = sum(bbox_iou(gt_box[None], get_inter_box(mid_box, end_box, i + 1, len(inter_boxes) + 1)[None]))
                iou_total += iou
            iou_total += 1
            iou_total /= (len(inter_boxes) + 1)

            if iou_total < thre:
                break

            inter_boxes.append(end_box)
            if z % self.frame_stride == 0:
                res_z = true_z

            z += 1 if direction == 'front' else -1
            if mid_id + z >= len(track.nodes) or mid_id + z < 0:
                break
            true_z = get_true_z(mid_node, track.get_node_by_index(mid_id + z))

        if not inter_boxes or len(inter_boxes) < self.frame_stride:
            return np.array([0, 0, 0, 0, 0])
        else:
            ret_ind = (len(inter_boxes) // self.frame_stride) * self.frame_stride - 1
            return np.concatenate((np.array([abs(res_z)]), inter_boxes[ret_ind] - mid_box))

    def get_item(self, frame_index):
        start_frame = frame_index
        max_len = (self.forward_frames * 2 - 1) * self.frame_stride + 1
        if self.max_frame_index - start_frame < max_len:
            return 0

        tubes = []
        for i in range(self.forward_frames * 2):
            frame_index = start_frame + i * self.frame_stride
            if frame_index not in self.recorder:
                continue

            det_ids = self.recorder[frame_index]

            # 1. get tubes
            for track_index, node_index in det_ids:
                t = self.tracks.get_track_by_index(track_index)
                n = t.get_node_by_index(node_index)
                mid_box = np.concatenate((n.box, np.array([frame_index - start_frame])))
                # backward
                back_box = self.bbox2tube(track=t, mid_id=node_index, direction='back',
                                          pos_in_video=i * self.frame_stride, thre=self.tube_thre)
                # forward
                front_box = self.bbox2tube(track=t, mid_id=node_index, direction='front',
                                           pos_in_video=i * self.frame_stride, thre=self.tube_thre)

                # remove the fast turning
                mid_box_w = mid_box[2] - mid_box[0]
                if abs(front_box[0] - back_box[0]) > 2.5 * mid_box_w:
                    print('remove turning')
                    continue

                tube = np.concatenate((mid_box, front_box, back_box))
                tubes.append(tube)

        if len(tubes) == 0:
            return 0
        tubes = np.array(tubes)

        pickle.dump(tubes, open(os.path.join(self.tube_res_path, str(start_frame)), 'wb'))
        return 0

    def __len__(self):
        return self.max_frame_index

class GTParser:
    def __init__(self, jta_root,
                 arg,
                 loose,
                 height_clamp,
                 type='train',
                 ):
        # analsis all the folder in jta_root
        # 1. get all the folders
        self.jta_root = jta_root
        jta_root = os.path.join(jta_root, type)
        all_folders = sorted(
            [os.path.join(jta_root, i) for i in os.listdir(jta_root)
             if os.path.isdir(os.path.join(jta_root, i))]
        )
        # 2. create single parser
        print('Init SingleParser')
        self.parsers = [GTSingleParser(folder, forward_frames=arg.forward_frames,
                                       min_visibility=arg.min_visibility,
                                       frame_stride=arg.frame_stride,
                                       tube_thre=arg.tube_thre,
                                       loose=loose,
                                       height_clamp=height_clamp) for folder in tqdm(all_folders, ncols=20)]

        # 3. get some basic information
        self.lens = [len(p) for p in self.parsers]
        self.len = sum(self.lens)

    def __len__(self):
        # get the length of all the matching frame
        return self.len

    def clear(self):
        print('Clearing')

    def run(self):
        print('Running')
        pool = multiprocessing.Pool(processes=40)
        pool_list = []
        for item in tqdm(range(self.len), ncols=20):
            total_len = 0
            index = 0
            current_item = item
            for l in self.lens:
                total_len += l
                if item < total_len:
                    break
                else:
                    index += 1
                    current_item -= l

            if index >= len(self.parsers):
                return
            pool_list.append(pool.apply_async(self.parsers[index].get_item, (current_item,)))
            # self.parsers[index].get_item(current_item)
        for p in tqdm(pool_list, ncols=20):
            p.get()
        pool.close()
        pool.join()

def get_gt(json_path, frames_path, loose, min_visiblity, height_clamp):
    assert os.path.exists(json_path), 'File does not exist: {}'.format(json_path)
    assert os.path.exists(frames_path), 'Folder does not exist: {}'.format(frames_path)
    split_path = frames_path.split('/')
    if frames_path[0] == '/':
        jta_root = '/' + os.path.join(*split_path[:-3])
    else:
        jta_root = os.path.join(*split_path[:-3])
    type = split_path[-2]
    video_name = split_path[-1]
    gt_path = os.path.join(jta_root, 'gt_' + str(loose) + '_' + str(min_visiblity) + '_' + str(height_clamp), type, video_name)
    try:
        os.makedirs(gt_path)
    except:
        pass
    gt_file = os.path.join(gt_path, 'gt.txt')
    df = pd.read_json(json_path)
    df = df.iloc[:, [0, 1, 3, 4, 8]]  # Frame, ID, x, y, occluded
    df_group = df.groupby([0, 1])  # Group by frame and id

    def get_bbox(g):
        assert len(g.columns) == 5
        if g.iloc[:, 4].sum() >= (1 - min_visiblity) * len(g):  # Completely occluded
            return pd.Series([-1, 0, 0, 0, 0, 0, 0], dtype=np.int)
        x1 = np.maximum(0, g.iloc[:, 2].min())
        y1 = np.maximum(0, g.iloc[:, 3].min())
        x2 = np.minimum(1920, g.iloc[:, 2].max())
        y2 = np.minimum(1080, g.iloc[:, 3].max())
        w = x2 - x1
        h = y2 - y1
        # Loose a little bit
        x1 -= np.round(w * loose)
        y1 -= np.round(h * loose)
        x1 = np.maximum(0.0, x1)
        y1 = np.maximum(0.0, y1)
        w = np.round(w * (1 + loose*2))
        h = np.round(h * (1 + loose*2))
        w = np.minimum(1920 - x1, w)
        h = np.minimum(1080 - y1, h)

        return pd.Series([x1, y1, w, h, 1, 1, 1], dtype=np.int)

    res_df = df_group.apply(get_bbox)
    res_df = res_df[res_df.iloc[:, 0] != -1]

    # get mode and remove the small box
    ns, edges = np.histogram(res_df.iloc[:, 3], bins=50)
    max_n = np.argmax(ns)
    mode = np.mean(edges[[max_n, max_n + 1]])
    res_df = res_df[res_df.iloc[:, 3] > height_clamp * mode]
    res_df = res_df[res_df.iloc[:, 3] > 7]

    res_df.to_csv(gt_file, header=False)

def get_gts(jta_root, frames_dir, loose, min_vis, height_clamp):
    pool = multiprocessing.Pool(processes=20)
    pool_list = []
    anno_path = os.path.join(jta_root, 'annotations')
    for type in os.listdir(anno_path):
        for json_file in os.listdir(os.path.join(anno_path, type)):
            json_path = os.path.join(anno_path, type, json_file)
            frames_path = os.path.join(jta_root, frames_dir, type, os.path.splitext(json_file)[0])
            pool_list.append(pool.apply_async(get_gt, (json_path, frames_path, loose, min_vis, height_clamp, )))

    for p in tqdm(pool_list, ncols=20):
        p.get()
    pool.close()
    pool.join()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--jta_root', type=str, help="data path of jta")
    parser.add_argument('--loose', type=float, default=0.1, help="ratio to loose the bbox generated from keypoint")
    parser.add_argument('--height_clamp', type=float, default=0.6, help="get rid of the bboxes whose height is smaller "
                                                                        "than 0.6 of the mean height")
    arg_input, unparsed = parser.parse_known_args()

    arg = __C
    cfg_from_file('../configs/get_jta_tube.yaml')

    print('Generating GT files')
    get_gts(jta_root=arg_input.jta_root, frames_dir='imgs', loose=arg_input.loose, min_vis=arg.min_visibility,
            height_clamp=arg_input.height_clamp)
    print('Generating Tubes')
    parser = GTParser(jta_root=os.path.join(arg_input.jta_root, 'imgs'), arg=arg, type='train', loose=arg_input.loose,
                      height_clamp=arg_input.height_clamp)
    parser.clear()
    parser.run()


import argparse
import itertools
import json
import logging
import pickle
import sys
from collections import defaultdict
from multiprocessing import Pool
from pathlib import Path
import numpy as np
import pandas as pd
import torch
from tqdm import tqdm
import os

# Add current directory to path
sys.path.insert(0, str(Path(__file__).resolve().parent))
from create_json_for_eval import create_json
from track_concat import track_concat
from associate_track import track_associate

def track_and_save(onnx_path1, onnx_path2, output, threshold_ass):

    onnx_detections1 = np.load(onnx_path1)
    onnx_detections2 = np.load(onnx_path2)
    #image_id, track_id, bb_left, bb_top, bb_width, bb_height, score, category, feature

    det_concat = track_concat(onnx_detections1, onnx_detections2)

    det_associate = track_associate(det_concat, threshold_ass)

    output.parent.mkdir(exist_ok=True, parents=True)
    np.save(output, det_associate[:, :8])

def track_and_save_star(args):
    track_and_save(*args)

def main():
    # Use first line of file docstring as description if it exists.
    parser = argparse.ArgumentParser(
        description=__doc__.split('\n')[0] if __doc__ else '',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--onnx_results1',
                        type=Path,
                        required=True,
                        help='results that use the first det')
    parser.add_argument('--onnx_results2',
                        type=Path,
                        required=True,
                        help='results that use the second det')
    parser.add_argument('--annotations',
                        type=Path,
                        required=True,
                        help='Annotations json')
    parser.add_argument(
        '--output-dir',
        type=Path,
        required=True,
        help='Output directory, where a results.json will be output.')

    parser.add_argument('--threshold-ass', type=float, default=0.7, help='threshold for track association')

    parser.add_argument('--workers', default=8, type=int)

    if 'RUN_MODE_IDE' in os.environ:
        args = parser.parse_args([
            "--onnx_results1", "/home/s1238640/Documents/DataSynced/PhD Project/Data/MRC Harwell/Scratch/TAO_Test/tracks",
            "--onnx_results2", "/home/s1238640/Documents/DataSynced/PhD Project/Data/MRC Harwell/Scratch/TAO_Test/tracks",
            "--annotations", "/home/s1238640/Documents/DataSynced/PhD Project/Data/MRC Harwell/Scratch/TAO_Test/train.json",
            "--output-dir", "/home/s1238640/Documents/DataSynced/PhD Project/Data/MRC Harwell/Scratch/TAO_Test/results",
            "--workers", "8",
        ])
        print(f'Running from PyCharm')
    else:
        args = parser.parse_args()
        print(f'Running from Terminal')

    args.output_dir.mkdir(exist_ok=True, parents=True)

    def get_output_path(video):
        return args.output_dir / (video + '.npy')

    with open(args.annotations, 'r') as f:
        groundtruth = json.load(f)

    videos = [x['name'] for x in groundtruth['videos']]

    tasks = []
    for video in tqdm(videos):
        onnx_path1 = args.onnx_results1 / (video + '.npy')
        onnx_path2 = args.onnx_results2 / (video + '.npy')
        output = get_output_path(video)
        tasks.append((onnx_path1, onnx_path2, output, args.threshold_ass))

    if args.workers > 0:
        pool = Pool(args.workers)
        list(
            tqdm(pool.imap_unordered(track_and_save_star, tasks),
                 total=len(tasks),
                 desc='Tracking'))
    else:
        for task in tqdm(tasks):
            track_and_save(*task)
    print(f'Finished')

    if 'train' in str(args.annotations):
        output = args.output_dir / 'train_results.json'
    elif 'val' in str(args.annotations):
        output = args.output_dir / 'val_results.json'
    elif 'test'in str(args.annotations):
        output = args.output_dir / 'test_results.json'
    else:
        print('wrong annotations')

    create_json(args.output_dir, groundtruth, output)

if __name__ == "__main__":
    main()


import numpy as np
from meta_mb.logger import logger
from meta_mb.algos.base import Algo
import tensorflow as tf
from collections import OrderedDict
from meta_mb.utils import create_feed_dict

class VGModel(Algo):
    """
    Algorithm for PPO MAML

    Args:
        policy (Policy): policy object
        name (str): tf variable scope
        learning_rate (float): learning rate for the meta-objective
        exploration (bool): use exploration / pre-update sampling term / E-MAML term
        inner_type (str): inner optimization objective - either log_likelihood or likelihood_ratio
        inner_lr (float) : gradient step size used for inner step
        meta_batch_size (int): number of meta-learning tasks
        num_inner_grad_steps (int) : number of gradient updates taken per maml iteration
        trainable_inner_step_size (boolean): whether make the inner step size a trainable variable
    """
    def __init__(
            self,
            policy,
            dynamics_model,
            value_function,
            tf_reward,
            discount,
            name="svg_1",
            learning_rate=1e-3,
            buffer_size=100000,
            max_grad_norm=5.,
            batch_size=64,
            num_grad_steps=100,
            kl_penalty=0.01,
            **kwargs
            ):
        super(VGModel, self).__init__(policy)

        self.dynamics_model = dynamics_model
        self.value_function = value_function
        self.tf_reward = tf_reward
        self.discount = discount
        self.learning_rate = learning_rate
        self.buffer_size = buffer_size
        self.max_grad_norm = max_grad_norm
        self.batch_size = batch_size
        self.num_grad_steps = num_grad_steps
        self.kl_penalty = kl_penalty

        self.recurrent = getattr(self.policy, 'recurrent', False)
        self._optimization_keys = ['observations', 'actions', 'next_observations', 'advantages', 'agent_infos']
        self.name = name
        self._train_op = None
        self._loss = None
        self._dataset = None

        self.build_graph()

    def build_graph(self):
        """
        Creates the computation graph

        Notes:
            Pseudocode:
            for task in meta_batch_size:
                make_vars
                init_init_dist_sym
            for step in num_inner_grad_steps:
                for task in meta_batch_size:
                    make_vars
                    update_init_dist_sym
            set objectives for optimizer
        """

        """ Create Variables """

        """ ----- Build graph for the meta-update ----- """
        self.op_phs_dict = OrderedDict()
        obs_ph, action_ph, next_obs_ph, adv_ph, dist_info_old_ph, all_phs_dict = \
            self._make_input_placeholders('train', recurrent=False, next_obs=True)

        self.op_phs_dict.update(all_phs_dict)

        distribution_info_vars = self.policy.distribution_info_sym(obs_ph)
        act_ev = distribution_info_vars['mean']

        model_distribution_info_vars = self.dynamics_model.distribution_info_sym(obs_ph, act_ev)
        obs_ev = model_distribution_info_vars['mean']

        value_fun_dist_info_vars = self.value_function.distribution_info_sym(obs_ev)
        value_ev = value_fun_dist_info_vars['mean']

        surr_obj = tf.reduce_mean(self.tf_reward(obs_ph, act_ev, obs_ev) + value_ev)

        # KL-penalty
        last_policy_info_vars = self.policy.distribution_info_sym(obs_ph, params=self.policy.policy_params_ph)
        surr_obj -= self.kl_penalty * self.policy.distribution.kl_sym(distribution_info_vars, last_policy_info_vars)

        self._train_op = tf.train.AdamOptimizer(self.learning_rate).minimize(-surr_obj,
                                                                             var_list=self.policy.get_params())
        self._loss = surr_obj

    def optimize_policy(self, samples_data, log=True):
        sess = tf.get_default_session()
        input_dict = self._extract_input_dict(samples_data, self._optimization_keys, prefix='train')
        if self._dataset is None:
            self._dataset = input_dict

        else:
            for k, v in input_dict.items():
                n_new_samples = len(v)
                n_max = self.buffer_size - n_new_samples
                self._dataset[k] = np.concatenate([self._dataset[k][-n_max:], v], axis=0)

        num_elements = len(list(self._dataset.values())[0])
        policy_params_dict = self.policy.policy_params_feed_dict
        for _ in range(self.num_grad_steps):
            idxs = np.random.randint(0, num_elements, size=self.batch_size)
            batch_dict = self._get_indices_from_dict(self._dataset, idxs)

            feed_dict = create_feed_dict(placeholder_dict=self.op_phs_dict, value_dict=batch_dict)
            feed_dict.update(policy_params_dict)

            _ = sess.run(self._train_op, feed_dict=feed_dict)

    def _get_indices_from_dict(self, d, idxs):
        _d = OrderedDict()
        for k, v in d.items():
            _d[k] = v[idxs]
        return _d


# Prediction interface for Cog ⚙
# Reference: https://github.com/replicate/cog/blob/main/docs/python.md

import clip
import os
from torch import nn
import numpy as np
import torch
import torch.nn.functional as nnf
import sys
from typing import Tuple, List, Union, Optional
from transformers import (
    GPT2Tokenizer,
    GPT2LMHeadModel,
    AdamW,
    get_linear_schedule_with_warmup,
)
import skimage.io as io
import PIL.Image

import cog

# import torch

N = type(None)
V = np.array
ARRAY = np.ndarray
ARRAYS = Union[Tuple[ARRAY, ...], List[ARRAY]]
VS = Union[Tuple[V, ...], List[V]]
VN = Union[V, N]
VNS = Union[VS, N]
T = torch.Tensor
TS = Union[Tuple[T, ...], List[T]]
TN = Optional[T]
TNS = Union[Tuple[TN, ...], List[TN]]
TSN = Optional[TS]
TA = Union[T, ARRAY]

WEIGHTS_PATHS = {
    "coco": "coco_weights.pt",
    "conceptual-captions": "conceptual_weights.pt",
}

D = torch.device
CPU = torch.device("cpu")

class Predictor(cog.Predictor):
    def setup(self):
        """Load the model into memory to make running multiple predictions efficient"""
        self.device = torch.device("cuda")
        self.clip_model, self.preprocess = clip.load(
            "ViT-B/32", device=self.device, jit=False
        )
        self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

        self.models = {}
        self.prefix_length = 10
        for key, weights_path in WEIGHTS_PATHS.items():
            model = ClipCaptionModel(self.prefix_length)
            model.load_state_dict(torch.load(weights_path, map_location=CPU))
            model = model.eval()
            model = model.to(self.device)
            self.models[key] = model

    @cog.input("image", type=cog.Path, help="Input image")
    @cog.input(
        "model",
        type=str,
        options=WEIGHTS_PATHS.keys(),
        default="coco",
        help="Model to use",
    )
    @cog.input(
        "use_beam_search",
        type=bool,
        default=False,
        help="Whether to apply beam search to generate the output text",
    )
    def predict(self, image, model, use_beam_search):
        """Run a single prediction on the model"""
        image = io.imread(image)
        model = self.models[model]
        pil_image = PIL.Image.fromarray(image)
        image = self.preprocess(pil_image).unsqueeze(0).to(self.device)
        with torch.no_grad():
            prefix = self.clip_model.encode_image(image).to(
                self.device, dtype=torch.float32
            )
            prefix_embed = model.clip_project(prefix).reshape(1, self.prefix_length, -1)
        if use_beam_search:
            return generate_beam(model, self.tokenizer, embed=prefix_embed)[0]
        else:
            return generate2(model, self.tokenizer, embed=prefix_embed)

class MLP(nn.Module):
    def forward(self, x: T) -> T:
        return self.model(x)

    def __init__(self, sizes: Tuple[int, ...], bias=True, act=nn.Tanh):
        super(MLP, self).__init__()
        layers = []
        for i in range(len(sizes) - 1):
            layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=bias))
            if i < len(sizes) - 2:
                layers.append(act())
        self.model = nn.Sequential(*layers)

class ClipCaptionModel(nn.Module):

    # @functools.lru_cache #FIXME
    def get_dummy_token(self, batch_size: int, device: D) -> T:
        return torch.zeros(
            batch_size, self.prefix_length, dtype=torch.int64, device=device
        )

    def forward(
        self, tokens: T, prefix: T, mask: Optional[T] = None, labels: Optional[T] = None
    ):
        embedding_text = self.gpt.transformer.wte(tokens)
        prefix_projections = self.clip_project(prefix).view(
            -1, self.prefix_length, self.gpt_embedding_size
        )
        # print(embedding_text.size()) #torch.Size([5, 67, 768])
        # print(prefix_projections.size()) #torch.Size([5, 1, 768])
        embedding_cat = torch.cat((prefix_projections, embedding_text), dim=1)
        if labels is not None:
            dummy_token = self.get_dummy_token(tokens.shape[0], tokens.device)
            labels = torch.cat((dummy_token, tokens), dim=1)
        out = self.gpt(inputs_embeds=embedding_cat, labels=labels, attention_mask=mask)
        return out

    def __init__(self, prefix_length: int, prefix_size: int = 512):
        super(ClipCaptionModel, self).__init__()
        self.prefix_length = prefix_length
        self.gpt = GPT2LMHeadModel.from_pretrained("gpt2")
        self.gpt_embedding_size = self.gpt.transformer.wte.weight.shape[1]
        if prefix_length > 10:  # not enough memory
            self.clip_project = nn.Linear(
                prefix_size, self.gpt_embedding_size * prefix_length
            )
        else:
            self.clip_project = MLP(
                (
                    prefix_size,
                    (self.gpt_embedding_size * prefix_length) // 2,
                    self.gpt_embedding_size * prefix_length,
                )
            )

class ClipCaptionPrefix(ClipCaptionModel):
    def parameters(self, recurse: bool = True):
        return self.clip_project.parameters()

    def train(self, mode: bool = True):
        super(ClipCaptionPrefix, self).train(mode)
        self.gpt.eval()
        return self

def generate_beam(
    model,
    tokenizer,
    beam_size: int = 5,
    prompt=None,
    embed=None,
    entry_length=67,
    temperature=1.0,
    stop_token: str = ".",
):

    model.eval()
    stop_token_index = tokenizer.encode(stop_token)[0]
    tokens = None
    scores = None
    device = next(model.parameters()).device
    seq_lengths = torch.ones(beam_size, device=device)
    is_stopped = torch.zeros(beam_size, device=device, dtype=torch.bool)
    with torch.no_grad():
        if embed is not None:
            generated = embed
        else:
            if tokens is None:
                tokens = torch.tensor(tokenizer.encode(prompt))
                tokens = tokens.unsqueeze(0).to(device)
                generated = model.gpt.transformer.wte(tokens)
        for i in range(entry_length):
            outputs = model.gpt(inputs_embeds=generated)
            logits = outputs.logits
            logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
            logits = logits.softmax(-1).log()
            if scores is None:
                scores, next_tokens = logits.topk(beam_size, -1)
                generated = generated.expand(beam_size, *generated.shape[1:])
                next_tokens, scores = next_tokens.permute(1, 0), scores.squeeze(0)
                if tokens is None:
                    tokens = next_tokens
                else:
                    tokens = tokens.expand(beam_size, *tokens.shape[1:])
                    tokens = torch.cat((tokens, next_tokens), dim=1)
            else:
                logits[is_stopped] = -float(np.inf)
                logits[is_stopped, 0] = 0
                scores_sum = scores[:, None] + logits
                seq_lengths[~is_stopped] += 1
                scores_sum_average = scores_sum / seq_lengths[:, None]
                scores_sum_average, next_tokens = scores_sum_average.view(-1).topk(
                    beam_size, -1
                )
                next_tokens_source = next_tokens // scores_sum.shape[1]
                seq_lengths = seq_lengths[next_tokens_source]
                next_tokens = next_tokens % scores_sum.shape[1]
                next_tokens = next_tokens.unsqueeze(1)
                tokens = tokens[next_tokens_source]
                tokens = torch.cat((tokens, next_tokens), dim=1)
                generated = generated[next_tokens_source]
                scores = scores_sum_average * seq_lengths
                is_stopped = is_stopped[next_tokens_source]
            next_token_embed = model.gpt.transformer.wte(next_tokens.squeeze()).view(
                generated.shape[0], 1, -1
            )
            generated = torch.cat((generated, next_token_embed), dim=1)
            is_stopped = is_stopped + next_tokens.eq(stop_token_index).squeeze()
            if is_stopped.all():
                break
    scores = scores / seq_lengths
    output_list = tokens.cpu().numpy()
    output_texts = [
        tokenizer.decode(output[: int(length)])
        for output, length in zip(output_list, seq_lengths)
    ]
    order = scores.argsort(descending=True)
    output_texts = [output_texts[i] for i in order]
    return output_texts

def generate2(
    model,
    tokenizer,
    tokens=None,
    prompt=None,
    embed=None,
    entry_count=1,
    entry_length=67,  # maximum number of words
    top_p=0.8,
    temperature=1.0,
    stop_token: str = ".",
):
    model.eval()
    generated_num = 0
    generated_list = []
    stop_token_index = tokenizer.encode(stop_token)[0]
    filter_value = -float("Inf")
    device = next(model.parameters()).device

    with torch.no_grad():

        for entry_idx in range(entry_count):
            if embed is not None:
                generated = embed
            else:
                if tokens is None:
                    tokens = torch.tensor(tokenizer.encode(prompt))
                    tokens = tokens.unsqueeze(0).to(device)

                generated = model.gpt.transformer.wte(tokens)

            for i in range(entry_length):

                outputs = model.gpt(inputs_embeds=generated)
                logits = outputs.logits
                logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
                cumulative_probs = torch.cumsum(
                    nnf.softmax(sorted_logits, dim=-1), dim=-1
                )
                sorted_indices_to_remove = cumulative_probs > top_p
                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[
                    ..., :-1
                ].clone()
                sorted_indices_to_remove[..., 0] = 0

                indices_to_remove = sorted_indices[sorted_indices_to_remove]
                logits[:, indices_to_remove] = filter_value
                next_token = torch.argmax(logits, -1).unsqueeze(0)
                next_token_embed = model.gpt.transformer.wte(next_token)
                if tokens is None:
                    tokens = next_token
                else:
                    tokens = torch.cat((tokens, next_token), dim=1)
                generated = torch.cat((generated, next_token_embed), dim=1)
                if stop_token_index == next_token.item():
                    break

            output_list = list(tokens.squeeze().cpu().numpy())
            output_text = tokenizer.decode(output_list)
            generated_list.append(output_text)

    return generated_list[0]


import numpy as np
import torch
import torch.utils.data as Data
from torch.autograd import Variable
from torchvision import transforms
import cv2
import os
import PIL.Image as Image

from EvalBox.Evaluation.evaluation import Evaluation
from EvalBox.Evaluation.evaluation import MIN_COMPENSATION

class SNS(Evaluation):
    def __init__(self, outputs_origin,outputs_adv, device, model, **kwargs):
        '''
        @description:
            该方法计算的是敏感单元。
            在参数设置时应注意
                1. IS_WHITE为True，白盒攻击
                2. IS_PYTHORCH_WHITE为True
                3. IS_COMPARE_MODEL为False
        @param {
            model:
            device:
            kwargs:
        }
        @return: None
        '''
        super(SNS, self).__init__(outputs_origin,outputs_adv, device)
        self.model = model
        self.device = device
        self._parsing_parameters(**kwargs)

    def _parsing_parameters(self, **kwargs):
        '''
        @description:
        @param {
            batch_size:
        }
        @return:
        '''
        self.batch_size = kwargs.get('batch_size', 64)
        self.feature_dict = dict()
        self.module_name = []
        # self.features_in_hook = []
        self.features_out_hook = []
        self.clean_feature_map = []

    def for_hook(self, module, fea_in, fea_out):
        # self.module_name.append(module.__class__)
        # self.features_in_hook.append(fea_in)
        self.features_out_hook.append(np.squeeze(fea_out.data.cpu().numpy()))

    def output_feature_map(self):
        print("*"*5+"hook record features"+"*"*5)
        for i in range(len(self.module_name)):
            print(self.module_name[i])
            print(self.clean_feature_map[i].shape)
            print(self.features_out_hook[i].shape)
        print("*"*5+"hook record features"+"*"*5)

    def evaluate(self,adv_xs=None, cln_xs=None, cln_ys=None,adv_ys=None,target_preds=None, target_flag=False):
        '''
        @description:
        @param {
            adv_xs: 攻击样本
            cln_xs：原始样本
            cln_ys: 原始类别，非目标攻击下原始样本的类型
            adv_ys: 攻击样本的预测类别
            target_preds： 目标攻击下希望原始样本攻击的目标类别
            target_flag：是否是目标攻击
        }
        @return: sns {}
        '''
        #print("target_flag",target_flag)
        total = len(adv_xs)
        print("total",total)
        device = self.device
        self.model = self.model.eval().to(device)
        assert len(adv_xs) == len(cln_ys), 'examples and labels do not match.'

        adv_label=[]
        number = 0

        # 监听所有中间层特征
        hook_handle_list = []
        for name, module in self.model._modules.items():
            self.module_name.append(name)
            handle = module.register_forward_hook(self.for_hook)
            hook_handle_list.append(handle)

        cln_dataset = Data.TensorDataset(cln_xs, cln_ys)
        adv_dataset = Data.TensorDataset(adv_xs, adv_ys)
        cln_loader = Data.DataLoader(cln_dataset, batch_size=self.batch_size, num_workers=3)
        adv_loader = Data.DataLoader(adv_dataset, batch_size=self.batch_size, num_workers=3)

        # 原始样本获取feature_map
        for x, y in cln_loader:
            x, y = Variable(x.to(device)), Variable(y.to(device))
            with torch.no_grad():
                pred = self.model(x)
        self.clean_feature_map = self.features_out_hook.copy()
        self.features_out_hook.clear()

        for x, y in adv_loader:
            x, y = Variable(x.to(device)), Variable(y.to(device))
            with torch.no_grad():
                pred = self.model(x)

        # 释放hook
        for handle in hook_handle_list:
            handle.remove()
        # 输出feature map =  feature_num * N * neuron_num * output
        # self.output_feature_map()

        average_diff_l1_dict = {}
        average_diff_l2_dict = {}
        average_diff_ratio_l1_dict = {}
        average_diff_ratio_l2_dict = {}
        useless_neuron = {}
        use_gpu = torch.cuda.is_available()
        # key表示特征名
        for key in range(len(self.module_name)):
            # k 表示第k个测试样本
            for k in range(len(self.clean_feature_map[key])):
                # c表示第c个神经元
                for c in range(len(self.clean_feature_map[key][k])):
                    #load feature maps and compute the corresponding criteria
                    x = self.clean_feature_map[key][k][c].flatten()
                    x_adv = self.features_out_hook[key][k][c].flatten()
                    x_l1norm = np.linalg.norm(x, ord=1) / len(x)
                    x_l2norm = np.linalg.norm(x, ord=2) / len(x)
                    diff_feat = x_adv - x
                    diff_feat_l1norm = np.linalg.norm(diff_feat, ord=1) / len(x)
                    diff_feat_l2norm = np.linalg.norm(diff_feat, ord=2) / len(x)
                    if(x_l1norm != 0):
                        diff_ratio_l1 = diff_feat_l1norm / x_l1norm
                        diff_ratio_l2 = diff_feat_l2norm / x_l2norm
                    else:
                        if(key not in useless_neuron):
                            useless_neuron[key] = [c]
                        else:
                            useless_neuron[key] = list(set(useless_neuron[key]))
                            useless_neuron[key].append(c)
                        diff_ratio_l1 = 1.0
                        diff_ratio_l2 = 1.0

                    # update dicts
                    if key not in average_diff_l1_dict.keys():
                        average_diff_l1_dict[key] = np.zeros(len(self.clean_feature_map[key][k]))
                        average_diff_l2_dict[key] = np.zeros(len(self.clean_feature_map[key][k]))
                        average_diff_ratio_l1_dict[key] = np.zeros(len(self.clean_feature_map[key][k]))
                        average_diff_ratio_l2_dict[key] = np.zeros(len(self.clean_feature_map[key][k]))
                    average_diff_l1_dict[key][c] += diff_feat_l1norm
                    average_diff_l2_dict[key][c]  += diff_feat_l2norm
                    average_diff_ratio_l1_dict[key][c]  += diff_ratio_l1
                    average_diff_ratio_l2_dict[key][c]  += diff_ratio_l2

        for key in average_diff_l1_dict:
            average_diff_l1_dict[key] /= total
            average_diff_l2_dict[key] /= total
            average_diff_ratio_l1_dict[key] /= total
            average_diff_ratio_l2_dict[key] /= total

        with open("neuron_sensitivity.log", "w") as fp:
            fp.write("-------------------------------\n每层前top 10%神经元敏感度\n")
            print("每层前top 10% (最多展示10个) 神经元敏感度")
            for key in average_diff_l1_dict:
                knum = min(10, max(1, int(len(average_diff_l1_dict[key]) / 10)))
                average_diff_l1_dict[key].sort()
                topk = average_diff_l1_dict[key][-knum:]
                topk = [float('{:.4f}'.format(i)) for i in topk]
                # print(self.module_name[key], topk)
                fp.write("{} {}\n".format(self.module_name[key], topk))

            print("每层mean神经元敏感度")
            fp.write("每层mean神经元敏感度\n")
            for key in average_diff_l1_dict:
                mSNS = format(average_diff_l1_dict[key].mean(), '.4f')
                # print(self.module_name[key], mSNS)
                fp.write("{} {}\n".format(self.module_name[key], mSNS))

        # 获取除Conv 及 输出层外，敏感度最高层
        layer_name = -1
        for key in average_diff_l1_dict:
            if "conv" in self.module_name[key].lower() or key == len(average_diff_l1_dict) - 1:
                continue
            if layer_name == -1 or average_diff_l1_dict[layer_name].mean() < average_diff_l1_dict[key].mean():
                layer_name = key
        print(self.module_name[layer_name])
        return 1.0


"""=== main Title =>

    Author : Saifeddine ALOUI
    Description :
        A code to test FaceAnalyzer: puts a face mask on a realtime video feed. Can be used to build Halloween masks...
        Feel free to use other assets if you want.
        In the new versions, a reference triangulation file is used. Dont forget to regenerate another one if you change the landmarks you use.
        The image file used in this code is under creative commons licence.
<================"""

import numpy as np
from pathlib import Path
import cv2
import matplotlib.pyplot as plt
import time
from PIL import Image
from FaceAnalyzer import FaceAnalyzer, Face
from pathlib import Path
import pickle
# Parameters
use_simplified_face = True  # If true then the simplified faster version of the face landmarks will be used instead of the full more accurate but slower version
seemless_cloning=False #If true, seemless cloning will be used to blend the mask onto the face

# Select landmarks
if use_simplified_face:
    # Simplified facial features (fast)
    lm_indices = list(set(
                        Face.simplified_face_features+
                        Face.mouth_inner_indices+
                        Face.mouth_outer_indices+
                        Face.left_eyelids_indices+
                        Face.right_eyelids_indices+
                        Face.left_eye_contour_indices+
                        Face.right_eye_contour_indices)) # list(range(468)) #
else:
    # Full face (Slower)
    lm_indices = list(range(478))

# open reference triangulation file
file = Path(__file__).parent/"reference.pkl"
with open(str(file),"rb") as f:
    triangles = pickle.load(f)

# open an image and recover all faces inside it (here there is a single face)
fa_mask = FaceAnalyzer.from_image(str(Path(__file__).parent/"assets/pennywize.jpg"))
mask_face = fa_mask.faces[0]
mask_face.triangles=triangles

cap = cv2.VideoCapture(0)
# cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
# cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)

# Build a window
cv2.namedWindow('Face Mask', flags=cv2.WINDOW_NORMAL)
cv2.resizeWindow('Face Mask', (640,480))
# Smask_face between viewing the original image with the face triangles on, or applying the mask to the video stream
view_original = True
if view_original:
    cv2.namedWindow('Original', flags=cv2.WINDOW_NORMAL)
    cv2.resizeWindow('Original', (640,480))

# Plot
start_time = time.time_ns()
fa = FaceAnalyzer(max_nb_faces=3)
for face in fa.faces:
    face.triangles = mask_face.triangles

if view_original:
    image = fa_mask.image.copy()
    mask_face.draw_delaunay(image, landmark_indices=lm_indices)
    cv2.imshow('Original', cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

while cap.isOpened():
    success, image = cap.read()
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    files = []
    fa.process(image)
    if fa.found_faces:
        for i, face in enumerate(fa.faces):
            if face.ready:
                image = mask_face.copyToFace(
                    face,
                    fa_mask.image,
                    image,
                    opacity = 1,
                    landmark_indices=lm_indices,
                    min_output_triangle_cross=0.01,
                    min_input_triangle_cross=20,
                    seemless_cloning=seemless_cloning)
    try:
        cv2.imshow('Face Mask', cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    except Exception as ex:
        print(ex)
    # Wait for key stroke
    wk = cv2.waitKey(5)
    if wk & 0xFF == 27: # If escape is pressed then return
        break
    if wk & 0xFF == 115: # If s is pressed then take a snapshot
        sc_dir = Path(__file__).parent/"screenshots"
        if not sc_dir.exists():
            sc_dir.mkdir(exist_ok=True, parents=True)
        i = 1
        file = sc_dir /f"sc_{i}.jpg"
        while file.exists():
            i+=1
            file = sc_dir /f"sc_{i}.jpg"
        cv2.imwrite(str(file),cv2.cvtColor(image,cv2.COLOR_BGR2RGB))
        print("Shot")
cap.release()


"""
 Copyright (C) 2020 Intel Corporation

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

      http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
"""

import numpy as np

from .model import Model
from .utils import Detection, resize_image, load_labels

class SSD(Model):
    def __init__(self, ie, model_path, labels=None, keep_aspect_ratio_resize=False):
        super().__init__(ie, model_path)

        self.keep_aspect_ratio_resize = keep_aspect_ratio_resize
        if isinstance(labels, (list, tuple)):
            self.labels = labels
        else:
            self.labels = load_labels(labels) if labels else None

        self.image_blob_name, self.image_info_blob_name = self._get_inputs()
        self.n, self.c, self.h, self.w = self.net.input_info[self.image_blob_name].input_data.shape

        self.output_parser = self._get_output_parser(self.net, self.image_blob_name)

    def _get_inputs(self):
        image_blob_name = None
        image_info_blob_name = None
        for blob_name, blob in self.net.input_info.items():
            if len(blob.input_data.shape) == 4:
                image_blob_name = blob_name
            elif len(blob.input_data.shape) == 2:
                image_info_blob_name = blob_name
            else:
                raise RuntimeError('Unsupported {}D input layer "{}". Only 2D and 4D input layers are supported'
                                   .format(len(blob.shape), blob_name))
        if image_blob_name is None:
            raise RuntimeError('Failed to identify the input for the image.')
        return image_blob_name, image_info_blob_name

    def _get_output_parser(self, net, image_blob_name, bboxes='bboxes', labels='labels', scores='scores'):
        try:
            parser = SingleOutputParser(net.outputs)
            self.logger.info('Use SingleOutputParser')
            return parser
        except ValueError:
            pass

        try:
            parser = MultipleOutputParser(net.outputs, bboxes, scores, labels)
            self.logger.info('Use MultipleOutputParser')
            return parser
        except ValueError:
            pass

        try:
            parser = BoxesLabelsParser(net.outputs, net.input_info[image_blob_name].input_data.shape[2:][::-1])
            self.logger.info('Use BoxesLabelsParser')
            return parser
        except ValueError:
            pass
        raise RuntimeError('Unsupported model outputs')

    def preprocess(self, inputs):
        image = inputs

        resized_image = resize_image(image, (self.w, self.h), self.keep_aspect_ratio_resize)
        meta = {'original_shape': image.shape,
                'resized_shape': resized_image.shape}

        h, w = resized_image.shape[:2]
        if h != self.h or w != self.w:
            resized_image = np.pad(resized_image, ((0, self.h - h), (0, self.w - w), (0, 0)),
                                   mode='constant', constant_values=0)
        resized_image = resized_image.transpose((2, 0, 1))  # Change data layout from HWC to CHW
        resized_image = resized_image.reshape((self.n, self.c, self.h, self.w))

        dict_inputs = {self.image_blob_name: resized_image}
        if self.image_info_blob_name:
            dict_inputs[self.image_info_blob_name] = [self.h, self.w, 1]
        return dict_inputs, meta

    def postprocess(self, outputs, meta):
        detections = self.output_parser(outputs)
        orginal_image_shape = meta['original_shape']
        resized_image_shape = meta['resized_shape']
        scale_x = self.w / resized_image_shape[1] * orginal_image_shape[1]
        scale_y = self.h / resized_image_shape[0] * orginal_image_shape[0]
        for detection in detections:
            detection.xmin *= scale_x
            detection.xmax *= scale_x
            detection.ymin *= scale_y
            detection.ymax *= scale_y
        return detections

def find_layer_by_name(name, layers):
    suitable_layers = [layer_name for layer_name in layers if name in layer_name]
    if not suitable_layers:
        raise ValueError('Suitable layer for "{}" output is not found'.format(name))

    if len(suitable_layers) > 1:
        raise ValueError('More than 1 layer matched to "{}" output'.format(name))

    return suitable_layers[0]

class SingleOutputParser:
    def __init__(self, all_outputs):
        if len(all_outputs) != 1:
            raise ValueError('Network must have only one output.')
        self.output_name, output_data = next(iter(all_outputs.items()))
        last_dim = np.shape(output_data)[-1]
        if last_dim != 7:
            raise ValueError('The last dimension of the output blob must be equal to 7, '
                             'got {} instead.'.format(last_dim))

    def __call__(self, outputs):
        return [Detection(xmin, ymin, xmax, ymax, score, label)
                for _, label, score, xmin, ymin, xmax, ymax in outputs[self.output_name][0][0]]

class MultipleOutputParser:
    def __init__(self, layers, bboxes_layer='bboxes', scores_layer='scores', labels_layer='labels'):
        self.labels_layer = find_layer_by_name(labels_layer, layers)
        self.scores_layer = find_layer_by_name(scores_layer, layers)
        self.bboxes_layer = find_layer_by_name(bboxes_layer, layers)

    def __call__(self, outputs):
        bboxes = outputs[self.bboxes_layer][0]
        scores = outputs[self.scores_layer][0]
        labels = outputs[self.labels_layer][0]
        return [Detection(*bbox, score, label) for label, score, bbox in zip(labels, scores, bboxes)]

class BoxesLabelsParser:
    def __init__(self, layers, input_size, labels_layer='labels', default_label=0):
        try:
            self.labels_layer = find_layer_by_name(labels_layer, layers)
        except ValueError:
            self.labels_layer = None
            self.default_label = default_label

        self.bboxes_layer = self.find_layer_bboxes_output(layers)
        self.input_size = input_size

    @staticmethod
    def find_layer_bboxes_output(layers):
        filter_outputs = [name for name, data in layers.items() if len(np.shape(data)) == 2 and np.shape(data)[-1] == 5]
        if not filter_outputs:
            raise ValueError('Suitable output with bounding boxes is not found')
        if len(filter_outputs) > 1:
            raise ValueError('More than 1 candidate for output with bounding boxes.')
        return filter_outputs[0]

    def __call__(self, outputs):
        bboxes = outputs[self.bboxes_layer]
        scores = bboxes[:, 4]
        bboxes = bboxes[:, :4]
        bboxes[:, 0::2] /= self.input_size[0]
        bboxes[:, 1::2] /= self.input_size[1]
        if self.labels_layer:
            labels = outputs[self.labels_layer]
        else:
            labels = np.full(len(bboxes), self.default_label, dtype=bboxes.dtype)

        detections = [Detection(*bbox, score, label) for label, score, bbox in zip(labels, scores, bboxes)]
        return detections