'''
KnowledgeBase Class
'''
import faiss
import numpy as np
import pickle
from glob import glob
from jel.common_config import ENTITY_DATA_PATH, ENTITY_VEC_DIR_PATH, RESOURCES_DIRPATH
from tqdm import tqdm
import os

class TitleIndexerWithFaiss:
    def __init__(self, kbemb_dim=300,
                 search_method_for_faiss='indexflatip',
                 how_many_top_hits_preserved=20):
        '''
        TODO: separate Entities to mysql DB
        :param kbemb_dim:
        :param search_method_for_faiss:
        :param how_many_top_hits_preserved:
        '''
        self.kbemb_dim = kbemb_dim
        self.entity_idx2emb, self.entity_id2title = self._entity2vec_loader()
        self.entity_title2id = {}
        for idx, title in self.entity_id2title.items():
            self.entity_title2id.update({''.join(title): idx})
        self.entity_num = len(self.entity_idx2emb)
        self.search_method_for_faiss = search_method_for_faiss
        self._indexed_faiss_loader()
        self.KBmatrix, self.kb_idx2entity_idx = self._KBmatrixloader()
        self._indexed_faiss_KBemb_adder(KBmatrix=self.KBmatrix)

        self.how_many_top_hits_preserved = how_many_top_hits_preserved

    def _entity2vec_loader(self):
        if os.path.exists(RESOURCES_DIRPATH + 'entity_id2vec.pkl') and \
                os.path.exists(RESOURCES_DIRPATH + 'entity_id2name.pkl'):
            with open(RESOURCES_DIRPATH + 'entity_id2vec.pkl', 'rb') as f:
                entity_idx2emb = pickle.load(f)
            with open(RESOURCES_DIRPATH + 'entity_id2name.pkl', 'rb') as g:
                entity_id2name = pickle.load(g)

            return entity_idx2emb, entity_id2name

        pickles = glob(ENTITY_VEC_DIR_PATH+'*.pkl')
        entity_idx2emb, entity_id2name = {}, {}
        for pkl_path in tqdm(pickles):
            with open(pkl_path, 'rb') as f:
                for ent in pickle.load(f):
                    title = ent['entity_name']
                    if 'contextualized_entity' in ent['vec']:
                        vec = ent['vec']['contextualized_entity']
                    else:
                        vec = ent['vec']
                    idx = len(entity_idx2emb)
                    entity_idx2emb.update({idx: vec})
                    entity_id2name.update({idx: title})

        if not os.path.exists(RESOURCES_DIRPATH):
            os.mkdir(RESOURCES_DIRPATH)

        with open(RESOURCES_DIRPATH + 'entity_id2vec.pkl', 'wb') as f:
            pickle.dump(entity_idx2emb, f)
        with open(RESOURCES_DIRPATH + 'entity_id2name.pkl', 'wb') as g:
            pickle.dump(entity_id2name, g)

        return entity_idx2emb, entity_id2name

    def _KBmatrixloader(self):
        KBemb = np.random.randn(self.entity_num, self.kbemb_dim).astype('float32')
        kb_idx2mention_idx = {}
        for idx, (mention_idx, emb) in enumerate(self.entity_idx2emb.items()):
            KBemb[idx] = emb
            kb_idx2mention_idx.update({idx: mention_idx})

        return KBemb, kb_idx2mention_idx

    def _indexed_faiss_loader(self):
        if self.search_method_for_faiss == 'indexflatl2':  # L2
            self.indexed_faiss = faiss.IndexFlatL2(self.kbemb_dim)
        elif self.search_method_for_faiss == 'indexflatip':  #
            self.indexed_faiss = faiss.IndexFlatIP(self.kbemb_dim)
        elif self.search_method_for_faiss == 'cossim':  # innerdot * Beforehand-Normalization must be done.
            self.indexed_faiss = faiss.IndexFlatIP(self.kbemb_dim)

    def _indexed_faiss_KBemb_adder(self, KBmatrix):
        if self.search_method_for_faiss == 'cossim':
            KBemb_normalized_for_cossimonly = np.random.randn(self.entity_num, self.kbemb_dim).astype('float32')
            for idx, emb in enumerate(KBmatrix):
                if np.linalg.norm(emb, ord=2, axis=0) != 0:
                    KBemb_normalized_for_cossimonly[idx] = emb / np.linalg.norm(emb, ord=2, axis=0)
            self.indexed_faiss.add(KBemb_normalized_for_cossimonly)
        else:
            self.indexed_faiss.add(KBmatrix)

    def _indexed_faiss_returner(self):
        return self.indexed_faiss

    def search_with_emb(self, emb):
        scores, faiss_search_candidate_result_kb_idxs = self.indexed_faiss.search(
            np.array([emb]).astype('float32'),
            self.how_many_top_hits_preserved)
        top_titles, scores_from_dot = [], []

        for kb_idx, score in zip(faiss_search_candidate_result_kb_idxs[0], scores[0]):
            entity_idx = self.kb_idx2entity_idx[kb_idx]
            candidate_title = ''.join(self.entity_id2title[entity_idx])
            top_titles.append(candidate_title)
            scores_from_dot.append(score)

        return top_titles, scores_from_dot

    def title2entity_vec(self, title: str):
        if title in self.entity_title2id:
            return np.array(self.entity_idx2emb[self.entity_title2id[title]])
        else:
            return np.random.randn(300,)

if __name__ == '__main__':
    kb = TitleIndexerWithFaiss()


import numpy as np

import time # only for testing
class Front(object) :
    """
    Key component of NeuroMac. A Front is a phenomenological \
    growth cone with a dual character:

    - A Front is a point in space with a radius. When connected \
       to another Front, a frustrum (or cylinder) is created.
    - A Front is a phenomenological growth cone that contains a \
       set of growth rules.

    Each Front its own set of rules expressed in a Python file \
    that listed in the algo_name variable. Initially, a \
    neuronal structure consists of one Front but later, when a structure \
    develops by simultaneous expanding of multiple front, each front \
    can contain different growth-rules.

    .. warning:: This code is not to be changed and is directly called \
       by NeuroMaC
    """
    def __init__(self,entity_name,algo_name,xyz,radius,\
                 path_length,order) :
        """
        Constructor

        Parameters
        ----------
        entity_name : string
        algo_name : string
           File name containing the growth rules must be named :code:`<algo_name>.py`. \
           It's this name that is also specified in the configuration file.
        xyz : np.array
           3D vector
        radius : float
        """
        #: The name given to and entity
        #: This name is to be used to query the substrate
        self.entity_name = entity_name
        #: Growth-rules to use
        self.algo_name = algo_name
        #: 3D location
        self.xyz = xyz
        #: Radius of the front
        #: Consecutive fronts form a frustrum or cylinder
        self.radius = radius
        #: *Not directly set*. Attrribute set by :py:func:`growth_procs.prepare_next_front`
        self.path_length= path_length
        #: *Not directly set*. Attrribute set by :py:func:`growth_procs.prepare_next_front`
        self.order = order
        #:  SWC-type as define in Cannon et al 1998
        self.swc_type = 7 # SWC-type field
        self.soma_pos =None
        self.parent = None

    def __str__(self):
        return "Front " + self.entity_name +": " + str(self.xyz)

    def __eq__(self,other):
        if other == None: return False

        if not self.entity_name == other.entity_name: return False
        if not self.algo_name == other.algo_name: return False
        if not np.allclose(self.xyz,other.xyz): return False
        if not self.radius == other.radius: return False
        if not self.path_length == other.path_length: return False
        if not self.order == other.order: return False
        if not self.swc_type == other.swc_type: return False
        if other.soma_pos == None:
            if not self.soma_pos == None:
                return False
        else:
            #print "self.soma_pos=", self.soma_pos, ", other.soma_pos=",other.soma_pos
            if not np.allclose(self.soma_pos,other.soma_pos): return False
        if not self.parent == other.parent: return False

        # otherwise, these objects are the same!
        return True

    def __ne__(self,other):
        return not self == other

    def __cmp__(self,other):
        if self.__eq__(other): return 0
        else: return -1

    # quick and dirty way to create a unique hash, needed to work with sets
    def __key(self):
        # WARNING!!! tuple(self.xyz) potentially not correct with rounding errors...
        return (self.entity_name, self.algo_name, tuple(self.xyz),\
                self.radius,self.path_length,self.order,\
                self.swc_type)

    def __hash__(self):
        return hash(self.__key())

    def extend_front(self,seed,constellation,virtual_substrate={}) :
        """
        Main function. Outsources the execution to the actual growth-rules \
        specified in the configuration file (:code:`cell_type_x` > \
        :code:`algorithm`).

        See :ref:`implement-rules`

        .. note:: This function is not to be modified and is only used \
        by framework code (Subvolume.py)

        Parameters
        ----------
        seed : int
        constellation : dict of list of np.array
        virtual_substrate : dict of list

        Returns
        -------
        front : :class:`Front`
           List of fronts that are the "extension" of the current front. \
           By convention, an empty list of None represents a terminating \
           front. A list with one or two entries represents an elongating =\
           or a branching front, respectively. One exception can be the \
           soma from which multiple new fronts can sprout.
        """

        # process the virtual substrate information and set variable in "self"
        for key in virtual_substrate.keys():
            # setattr(x, attr, 'magic')
            for x in virtual_substrate[key].keys():
                rectangle = virtual_substrate[key][x]
                b0 = np.array(rectangle[0])
                b1 = np.array(rectangle[1])
                if np.all(self.xyz>b0) and np.all(self.xyz<b1): # works for rectangle only...
                    setattr(self,key,x)
                    print(">>>> FOUND: ", key, " = ", x)

        # outsource the real call to extend a front
        import  sys
        sys.path.insert(0,'D:/gen_vir_experiment_code/neuromac-master/')
        globals()[self.algo_name] = __import__(self.algo_name)
        ret = globals()[self.algo_name].extend_front(self,seed,constellation)

        return ret


"""
Licensed Materials - Property of IBM
Restricted Materials of IBM
20190891
© Copyright IBM Corp. 2021 All Rights Reserved.
"""

import logging
import keras
import time
import json
import numpy as np
import tensorflow as tf

# if tf.__version__[0:4] != "1.15":
#     raise ImportError("KerasFLModel requires TensorFlow v1.15.")

from keras_preprocessing.image.numpy_array_iterator import NumpyArrayIterator
from keras.preprocessing.image import ImageDataGenerator
from tensorflow.python.keras.backend import set_session

from ibmfl.util import config
from ibmfl.util import fl_metrics
from ibmfl.model.fl_model import FLModel
from ibmfl.model.model_update import ModelUpdate
from ibmfl.exceptions import FLException, LocalTrainingException, ModelException

logger = logging.getLogger(__name__)

class KerasFLModel(FLModel):
    """
    Wrapper class for importing keras and tensorflow.keras models.
    """

    def __init__(self, model_name,
                 model_spec,
                 keras_model=None,
                 **kwargs):
        """
        Create a `KerasFLModel` instance from a Keras model.
        If keras_model is provided, it will use it; otherwise it will take
        the model_spec to create the model.
        Assumes the `model` passed as argument is compiled.

        :param model_name: String specifying the type of model e.g., Keras_CNN
        :type model_name: `str`
        :param model_spec: Specification of the keras_model
        :type model_spec: `dict`
        :param keras_model: Compiled keras model.
        :type keras_model: `keras.models.Model`
        :param kwargs: A dictionary contains other parameter settings on \
         to initialize a Keras or TensorFlow.Keras model, \
         e.g., GPU configuration.
        :type kwargs: `dict`
        """
        super().__init__(model_name, model_spec, **kwargs)

        self.graph = tf.get_default_graph()
        self.sess = tf.Session()
        set_session(self.sess)

        if keras_model is None:
            if model_spec is None or (not isinstance(model_spec, dict)):
                raise ValueError('Initializing model requires '
                                 'a model specification or '
                                 'compiled keras model. '
                                 'None was provided')
            # In this case we need to recreate the model from model_spec
            self.model = self.load_model_from_spec(model_spec)
        else:
            if not issubclass(type(keras_model), (keras.models.Model,
                                                  tf.keras.models.Model)):
                raise ValueError('Compiled keras model needs to be provided '
                                 '(keras.models/tensorflow.keras.models). '
                                 'Type provided' + str(type(keras_model)))
            if not self.use_gpu_for_training or self.num_gpus == 1:
                self.model = keras_model
            elif issubclass(type(keras_model), keras.models.Model):
                from keras.utils import multi_gpu_model
                self.model = multi_gpu_model(keras_model, gpus=self.num_gpus)
                self.model.compile(optimizer=keras_model.optimizer,
                                   loss=keras_model.loss,
                                   metrics=keras_model.metrics)
            else:
                strategy = tf.distribute.MirroredStrategy()
                with strategy.scope():
                    self.model = keras_model

        # keras flag
        if issubclass(type(self.model), keras.models.Model):
            self.is_keras = True
            self.model_type = 'Keras-2.2.4'
        else:
            self.is_keras = False
            self.model_type = 'TensorFlow-Keras-1.15'

        # Default values for local training
        self.batch_size = 128
        self.epochs = 1
        self.steps_per_epoch = 100
        self.is_classification = True if not (model_spec and model_spec.get(
            'is_classification')) else model_spec.get('is_classification')

        # For storing model output
        self.train_result = None

    def fit_model(self, train_data, fit_params=None, **kwargs):
        """
        Fits current model with provided training data.

        :param train_data: Training data, a tuple given in the form \
        (x_train, y_train) or a datagenerator of type `keras.utils.Sequence`, \
        `keras.preprocessing.image.ImageDataGenerator`
        :type train_data: `np.ndarray`
        :param fit_params: (optional) Dictionary with hyperparameters \
        that will be used to call Keras fit function.\
        Hyperparameter parameters should match keras expected values \
        e.g., `epochs`, which specifies the number of epochs to be run. \
        If no `epochs` or `batch_size` are provided, a default value \
        will be used (1 and 128, respectively).
        :type fit_params: `dict`
        :return: None
        """

        hyperparams = fit_params.get('hyperparams', {}) or {} if fit_params else {}
        local_hp = hyperparams.get('local', {}) or {}
        training_hp = local_hp.get('training', {}) or {}

        # Initialized with default values if not in training_hp
        batch_size = training_hp.get('batch_size', self.batch_size)
        epochs = training_hp.get('epochs', self.epochs)
        steps_per_epoch = training_hp.get('steps_per_epoch', self.steps_per_epoch)

        logger.info('Training hps for this round => '
            'batch_size: {}, epochs {}, steps_per_epoch {}'
            .format(batch_size, epochs, steps_per_epoch))

        if epochs is None:
            logger.exception('epochs need to be provided')
            raise ModelException("Invalid hyperparams, epoch can't be None")

        try:

            if type(train_data) is tuple and type(train_data[0]) is np.ndarray:
                self.fit(
                    train_data, batch_size=batch_size, epochs=epochs)
            else:
                # batch_size won't be used for data generator
                self.fit_generator(train_data, epochs=epochs,
                                   steps_per_epoch=steps_per_epoch)

        except Exception as e:
            logger.exception(str(e))
            raise LocalTrainingException(
                'Error occurred while performing model.fit')

    def get_train_result(self):
        whitelist = []
        return {k: v for k, v in iter(self.train_result.items()) if k in whitelist}

    def fit(self, train_data, batch_size, epochs):
        """
        Fits current model using model.fit with provided training data.

        :param train_data: Training data, a tuple \
        given in the form (x_train, y_train).
        :type train_data: `np.ndarray`
        :param batch_size: Number of samples per gradient update.
        :type batch_size: `int`
        :param epochs: Number of epochs to train the model.
        :type epochs: `int`
        :return: None
        """
        x = train_data[0]
        y = train_data[1]
        with self.graph.as_default():
            set_session(self.sess)
            history = self.model.fit(x, y, batch_size=batch_size, epochs=epochs)
            self.train_result = history.history

    def fit_generator(self, training_generator, epochs, steps_per_epoch=None):
        """
        Fits current model using model.fit_generator with provided
        training data generator.

        :param training_generator: Training datagenerator of type \
        `keras.utils.Sequence`, or \
        `keras.preprocessing.image.ImageDataGenerator`. \
        :type training_generator: `ImageDataGenerator` or \
        `keras.utils.Sequence`.
        :param epochs: Number of epochs to train the model.
        :type epochs: `int`
        :param steps_per_epoch: Total number of steps (batches of samples) \
                to yield from `generator` before declaring one epoch. \
                Optional for `Sequence` data generator` as a number of steps.
        :type steps_per_epoch: `int`
        :return: None
        """

        if type(training_generator) is NumpyArrayIterator and not steps_per_epoch:
            raise LocalTrainingException(
                "Variable steps_per_epoch cannot be None for generators not \
                    of type keras.utils.Sequence!")

        with self.graph.as_default():
            set_session(self.sess)
            history = self.model.fit_generator(
                training_generator, steps_per_epoch=steps_per_epoch, epochs=epochs)
            self.train_result = history.history

    def update_model(self, model_update):
        """
        Update keras model with provided model_update, where model_update
        should be generated according to `KerasFLModel.get_model_update()`.

        :param model_update: `ModelUpdate` object that contains the weight \
        that will be used to update the model.
        :type model_update: `ModelUpdate`
        :return: None
        """
        if isinstance(model_update, ModelUpdate):
            with self.graph.as_default():
                set_session(self.sess)
                w = model_update.get("weights")
                self.model.set_weights(w)
        else:
            raise LocalTrainingException('Provided model_update should be of '
                                         'type ModelUpdate. '
                                         'Instead they are:' +
                                         str(type(model_update)))

    def get_model_update(self):
        """
        Generates a `ModelUpdate` object that will be sent to other entities.

        :return: ModelUpdate
        :rtype: `ModelUpdate`
        """
        w = self.model.get_weights()
        return ModelUpdate(weights=w)

    def predict(self, x, batch_size=128, **kwargs):
        """
        Perform prediction for a batch of inputs. Note that for classification
        problems, it returns the resulting probabilities.

        :param x: Samples with shape as expected by the model.
        :type x: `np.ndarray`
        :param batch_size: Size of batches.
        :type batch_size: `int`
        :param kwargs: Dictionary of keras-specific arguments.
        :type kwargs: `dict`
        :return: Array of predictions
        :rtype: `np.ndarray`
        """
        with self.graph.as_default():
            set_session(self.sess)
            try:
                if type(x) is np.ndarray:
                    preds = self.model.predict(
                        x, batch_size=batch_size, **kwargs)
                else:
                    steps = self.steps_per_epoch
                    if 'steps' in kwargs:
                        steps = kwargs['steps']

                    if not type(x) is NumpyArrayIterator and not steps:
                        raise LocalTrainingException(
                            "Variable steps cannot be None for generator "
                            "not of type keras.utils.Sequence")
                    preds = self.model.predict_generator(x, **kwargs)
            except Exception as ex:
                logger.exception(str(ex))
                raise LocalTrainingException(
                    "Error occurred during prediction.")
            return preds

    def evaluate(self, test_dataset, **kwargs):
        """
        Evaluates the model given testing data.

        :param test_dataset: Testing data, a tuple given in the form \
        (x_test, y_test) or a datagenerator of type `keras.utils.Sequence`, \
        `keras.preprocessing.image.ImageDataGenerator`
        :type test_dataset: `np.ndarray`
        :param kwargs: Dictionary of metrics available for the model
        :type kwargs: `dict`
        :return: metrics
        :rtype: `dict`
        """

        if type(test_dataset) is tuple:
            x_test = test_dataset[0]
            y_test = test_dataset[1]

            return self.evaluate_model(x_test, y_test, **kwargs)

        else:
            return self.evaluate_generator_model(
                test_dataset, **kwargs)

    def evaluate_model(self, x, y, batch_size=128, **kwargs):
        """
        Evaluates the model given x and y.

        :param x: Samples with shape as expected by the model.
        :type x: `np.ndarray`
        :param y: Corresponding labels to x
        :type y: `np.ndarray`
        :param batch_size: Size of batches.
        :type batch_size: `int`
        :param kwargs: Dictionary of metrics available for the model
        :type kwargs: `dict`
        :return: metrics
        :rtype: `dict`
        """

        with self.graph.as_default():
            set_session(self.sess)
            metrics = self.model.evaluate(
                x, y, batch_size=batch_size, **kwargs)
            names = self.model.metrics_names
            dict_metrics = {}
            additional_metrics = {}
            if type(metrics) == list:
                for metric, name in zip(metrics, names):
                    dict_metrics[name] = metric
            else:
                dict_metrics[names[0]] = metrics

            y_pred = self.predict(x, batch_size)
            if self.is_classification:
                additional_metrics = fl_metrics.get_eval_metrics_for_classificaton(
                    y, y_pred)
            else:
                additional_metrics = fl_metrics.get_eval_metrics_for_regression(
                    y, y_pred)

            logger.info(additional_metrics)
            dict_metrics = {**dict_metrics, **additional_metrics}
            logger.info(dict_metrics)

        return dict_metrics

    def evaluate_generator_model(self, test_generator, **kwargs):
        """
        Evaluates the model based on the provided data generator.

        :param test_generator: Testing datagenerator of type \
        `keras.utils.Sequence`, or \
        `keras.preprocessing.image.ImageDataGenerator`.
        :type test_generator: `ImageDataGenerator` or `keras.utils.Sequence`
        :return: metrics
        :rtype: `dict`
        """

        steps = self.steps_per_epoch
        if steps in kwargs:
            steps = kwargs.get('steps')

        if not type(test_generator) is NumpyArrayIterator and not steps:
            raise LocalTrainingException(
                "Variable steps cannot be None for generator "
                "not of type keras.utils.Sequence")
        with self.graph.as_default():
            metrics = self.model.evaluate_generator(
                test_generator, steps=steps)
            names = self.model.metrics_names
            dict_metrics = {}
            additional_metrics = {}

            if type(metrics) == list:
                for metric, name in zip(metrics, names):
                    dict_metrics[name] = metric
            else:
                dict_metrics[names[0]] = metrics

        return dict_metrics

    def save_model(self, filename=None):
        """
        Save a model to file in the format specific to the backend framework.

        :param filename: Name of the file where to store the model.
        :type filename: `str`
        :return: filename
        :rtype: `str`
        """
        if filename is None:
            file = self.model_name if self.model_name else self.model_type
            filename = '{}_{}.h5'.format(file, time.time())

        full_path = super().get_model_absolute_path(filename)
        self.model.save(full_path)
        logger.info('Model saved in path: %s.', full_path)
        return filename

    def load_model(self, file_name, custom_objects={}):
        """
        Loads a model from disk given the specified file_name

        :param file_name: Name of the file that contains the model to be loaded.
        :type file_name: `str`
        :param custom_objects: A dictionary of customized objects to be loaded.
        :type custom_objects: `dict`
        :return: Keras model loaded to memory
        :rtype: `keras.models.Model`
        """
        # try loading model from keras
        model = self.load_model_via_keras(file_name, custom_objects)
        if not model:
            # try loading model from tf.keras
            model = self.load_model_via_tf_keras(file_name, custom_objects)
            if model is None:
                logger.error('Loading model failed! '
                             'An acceptable compiled model should be of type '
                             '(keras.models/tensorflow.keras.models)!')
                raise FLException(
                    'Unable to load the provided compiled model!')

        return model

    def load_model_via_keras(self, file_name, custom_objects={}):
        """
        Loads a model from disk given the specified file_name via keras.

        :param file_name: Name of the file that contains the model to be loaded.
        :type file_name: `str`
        :param custom_objects: A dictionary of customized objects to be loaded.
        :type custom_objects: `dict`
        :return: Keras model loaded to memory
        :rtype: `keras.models.Model`
        """
        # try loading model from keras
        model = None
        try:
            if not self.use_gpu_for_training or self.num_gpus == 1:
                # CPU training or use only 1 GPU for training
                model = keras.models.load_model(
                    file_name, custom_objects=custom_objects)
                model._make_predict_function()
            else:
                # use multiple GPU for training
                tmp_model = keras.models.load_model(
                    file_name, custom_objects=custom_objects)
                from keras.utils import multi_gpu_model
                model = multi_gpu_model(tmp_model, gpus=self.num_gpus)
                model.compile(optimizer=tmp_model.optimizer,
                              loss=tmp_model.loss,
                              metrics=tmp_model.metrics)
        except Exception as ex:
            logger.error(
                'Loading model via keras.models.load_model failed!')

        return model

    def load_model_via_tf_keras(self, file_name, custom_objects={}):
        """
        Loads a model from disk given the specified file_name via tf.keras.

        :param file_name: Name of the file that contains the model to be loaded.
        :type file_name: `str`
        :param custom_objects: A dictionary of customized objects to be loaded.
        :type custom_objects: `dict`
        :return: tf.keras model loaded to memory
        :rtype: `tf.keras.models.Model`
        """
        # try load from tf.keras
        model = None
        try:
            if self.use_gpu_for_training:
                strategy = tf.distribute.MirroredStrategy()
                with strategy.scope():
                    model = tf.keras.models.load_model(
                        file_name, custom_objects=custom_objects)
                    model._make_predict_function()
            else:
                model = tf.keras.models.load_model(
                    file_name, custom_objects=custom_objects)
                model._make_predict_function()
        except Exception as ex:
            logger.error('Loading model via tf.keras.models.load_model '
                         'failed!')

        return model

    @staticmethod
    def model_from_json_via_keras(json_file_name):
        """
        Loads a model architecture from disk via keras
        given the specified json file name.

        :param json_file_name: Name of the file that contains \
        the model architecture to be loaded.
        :type json_file_name: `str`
        :return: Keras model with only model architecture loaded to memory
        :rtype: `keras.models.Model`
        """
        # try loading model from keras
        model = None
        json_file = open(json_file_name, 'r')
        f = json_file.read()
        json_file.close()
        try:
            model = keras.models.model_from_json(f)
        except Exception as ex:
            logger.error('Loading model via '
                         'keras.models.model_from_json failed!')

        return model

    @staticmethod
    def model_from_json_via_tf_keras(json_file_name):
        """
        Loads a model architecture from disk via tf.keras
        given the specified json file name.

        :param json_file_name: Name of the file that contains \
        the model architecture to be loaded.
        :type json_file_name: `str`
        :return: tf.keras model with only model architecture loaded to memory
        :rtype: `tf.keras.models.Model`
        """
        # try loading model from keras
        model = None
        json_file = open(json_file_name, 'r')
        f = json_file.read()
        json_file.close()
        try:
            model = tf.keras.models.model_from_json(f)
        except Exception as ex:
            logger.error(
                'Loading model via tf.keras.models.model_from_json failed! ')

        return model

    def load_model_from_spec(self, model_spec):
        """
        Loads model from provided model_spec, where model_spec is a `dict`
        that contains two items: model_spec['model_architecture'] has a
        pointer to the file where the keras model architecture in stored
        in json format, and model_spec['model_weights'] contains
        the path where the associated weights are stored as h5.

        :param model_spec: A dictionary of specifications for a Keras model.
        :type model_spec: `dict`
        :return: model
        :rtype: `keras.models.Model`
        """

        if 'model_definition' in model_spec:
            model_file = model_spec['model_definition']
            model_absolute_path = config.get_absolute_path(model_file)
            custom_objects = {}
            if 'custom_objects' in model_spec:

                custom_objects_config = model_spec['custom_objects']
                for custom_object in custom_objects_config:
                    key = custom_object['key']
                    value = custom_object['value']
                    path = custom_object['path']
                    custom_objects[key] = config.get_attr_from_path(
                        path, value)

            model = self.load_model(model_absolute_path,
                                    custom_objects=custom_objects)
        else:
            # Load architecture from json file
            try:
                model = KerasFLModel.model_from_json_via_keras(
                    model_spec['model_architecture'])
                if not model:
                    model = KerasFLModel.model_from_json_via_tf_keras(
                        model_spec['model_architecture'])
                    if model is None:
                        logger.error(
                            'An acceptable compiled model should be of type '
                            '(keras.models/tensorflow.keras.models)!')
                        raise FLException(
                            'Unable to load the provided compiled model!')
            except Exception as ex:
                logger.error(str(ex))
                raise FLException(
                    'Unable to load the provided compiled model!')

            # Load weights from h5 file
            if 'model_weights' in model_spec:
                model.load_weights(model_spec['model_weights'])
            # model.load_weights(weights)

            # Compile model with provided parameters:
            compiled_option = model_spec['compile_model_options']
            try:
                if 'optimizer' in compiled_option:
                    optimizer = compiled_option['optimizer']
                else:
                    logger.warning('No optimizer information was provided '
                                   'in the compile_model_options, '
                                   'set keras optimizer to default: SGD')
                    optimizer = 'sgd'
                if 'loss' in compiled_option:
                    loss = compiled_option['loss']
                else:
                    logger.warning('No loss function was provided '
                                   'in the compile_model_options.'
                                   'set keras loss function to default: None')
                    loss = None
                if 'metrics' in compiled_option:
                    metrics = compiled_option['metrics']
                    metrics = [metrics] if isinstance(
                        metrics, str) else metrics
                else:
                    logger.warning('No metrics information was provided '
                                   'in the compile_model_options,'
                                   'set keras metrics to default: None')
                    metrics = None
                model.compile(optimizer=optimizer,
                              loss=loss,
                              metrics=metrics)
            except Exception as ex:
                logger.exception(str(ex))
                logger.exception('Failed to compiled keras model.')
        return model

    def expand_model_by_layer_name(self, new_dimension, layer_name="dense"):
        """
        Expand the current Keras model with provided dimension of
        the hidden layers or model weights.
        This method by default expands the dense layer of
        the current neural network.
        It can be extends to expand other layers specified by `layer_name`,
        for example, it can be use to increase the number of CNN filters or
        increase the hidden layer size inside LSTM.

        :param new_dimension: New number of dimensions for \
        the fully connected layers
        :type new_dimension: `list`
        :param layer_name: layer's name to be expanded
        :type layer_name: `str`
        :return: None
        """
        if new_dimension is None:
            raise FLException('No information is provided for '
                              'the new expanded model. '
                              'Please provide the new dimension of '
                              'the resulting expanded model.')

        model_config = json.loads(self.model.to_json())
        i = 0

        for layer in model_config['config']['layers']:
            # find the specified layers
            if 'class_name' in layer and \
                    layer['class_name'].strip().lower() == layer_name:
                layer['config']['units'] = new_dimension[i]
                i += 1
        if self.is_keras:
            new_model = keras.models.model_from_json(json.dumps(model_config))
        else:
            new_model = tf.keras.models.model_from_json(
                json.dumps(model_config))

        metrics = self.model.metrics_names
        if 'loss' in metrics:
            metrics.remove('loss')
        if not self.use_gpu_for_training or self.num_gpus == 1:
            new_model.compile(optimizer=self.model.optimizer,
                              loss=self.model.loss,
                              metrics=metrics)
        elif self.is_keras:
            from keras.utils import multi_gpu_model
            new_model = multi_gpu_model(new_model, gpus=self.num_gpus)
            new_model.compile(optimizer=self.model.optimizer,
                              loss=self.model.loss,
                              metrics=metrics)
        else:
            strategy = tf.distribute.MirroredStrategy()
            with strategy.scope():
                new_model.compile(optimizer=self.model.optimizer,
                                  loss=self.model.loss,
                                  metrics=metrics)

        self.model = new_model

    def get_gradient(self, train_data):
        """
        Compute the gradient with the provided dataset at the current local
        model's weights.

        :param train_data: Training data, a tuple \
        given in the form (x_train, y_train).
        :type train_data: `np.ndarray`
        :return: gradients
        :rtype: `list` of `np.ndarray`
        """
        with self.graph.as_default():
            set_session(self.sess)
            # set up symbolic variables
            try:
                grads = self.model.optimizer.get_gradients(
                    self.model.total_loss,
                    self.model.trainable_weights)
            except Exception as ex:
                logger.exception(str(ex))
                raise FLException('Error occurred when defining '
                                  'gradient expression. ')
            symb_inputs = (self.model._feed_inputs +
                           self.model._feed_targets +
                           self.model._feed_sample_weights)

            # define the symbolic function
            if self.is_keras:
                from keras import backend as k
            else:
                from tensorflow.python.keras import backend as k

            f = k.function(symb_inputs, grads)
            try:
                x, y, sample_weight = self.model._standardize_user_data(
                    train_data[0],
                    train_data[1])
            except Exception as ex:
                logger.exception(str(ex))
                raise FLException('Error occurred when feeding data samples '
                                  'to compute current gradient.')

            if sample_weight:
                return f(x + y + sample_weight)
            else:
                return f(x + y)

    def is_fitted(self):
        """
        Return a boolean value indicating if the model is fitted or not.
        In particular, check if the keras model has weights.
        If it has, return True; otherwise return false.

        :return: res
        :rtype: `bool`
        """
        try:
            self.model.get_weights()
        except Exception:
            return False
        return True

    def get_weights(self):
        """
        Returns current model weights.

        :return: A list containing the current model weights.
        :rtype: `list` of `np.ndarray`
        """
        return self.model.get_weights()

    def get_loss(self, dataset):
        """
        Return the resulting loss computed based on the provided dataset.

        :param dataset: Provided dataset, a tuple given in the form \
        (x_test, y_test) or a datagenerator of type `keras.utils.Sequence`, \
        `keras.preprocessing.image.ImageDataGenerator`.
        :type dataset: `np.ndarray`
        :return: The resulting loss.
        :rtype: `float`
        """
        if 'loss' not in self.model.metrics_names:
            self.model.metrics_names.append('loss')

        res = self.evaluate(dataset)

        if 'loss' in res:
            return res['loss']
        else:
            raise FLException("Loss is not listed in the "
                              "model's metrics_names.")


#!/usr/bin/env python
# -*- coding: utf-8 -*-
# pylint: disable=C0411,broad-except,too-many-statements,too-many-branches,logging-fstring-interpolation,import-error,too-many-arguments
# pylint: disable=too-many-locals,missing-function-docstring

import numpy as np
import torch

from homan.lib2d import maskutils
from homan.mocap import process_handmocap_predictions
from homan.prepare.gtmasks import render_gt_masks
from homan.utils.bbox import bbox_wh_to_xy, bbox_xy_to_wh
from homan.viz.vizframeinfo import viz_frame_info

import os
from libyana.verify import checkshape

def process_hand_boxes(image, hand_boxes, hand_preds, mask_extractor,
                       image_size):
    if isinstance(hand_boxes, list):
        hand_boxes = np.stack(hand_boxes)
    hand_annots = mask_extractor.masks_from_bboxes(image,
                                                   hand_boxes,
                                                   class_idx=0,
                                                   pred_classes=None,
                                                   image_size=image_size)
    full_masks = np.stack([annot["full_mask"] for annot in hand_annots])
    hand_parameters = process_handmocap_predictions(
        mocap_predictions=hand_preds,
        bboxes=bbox_wh_to_xy(hand_boxes),
        masks=full_masks,
        image_size=image_size)
    return hand_parameters

def get_frame_infos(images_np,
                    hand_predictor=None,
                    mask_extractor=None,
                    sample_folder=None,
                    hand_bboxes=None,
                    obj_bboxes=None,
                    camintr=None,
                    debug=True,
                    image_size=640):
    """
    Arguments:
        images_np (list[np.ndarray]): List of input images
        hand_bboxes (dict): dictionnary containing {left_hand: [None|frame_nb x 4], right_hand: [None|frame_nb x 4]}
            sequence bounding boxes in xywh format
        obj_bboxes (np.ndarray): (1, frame_nb, 4) xywh object bounding boxes
        camintr (list[np.ndarray]): (frame_nb, 3, 3) intrinsic camera parameters
        image_size (int): image size
    """
    checkshape.check_shape(obj_bboxes, (1, -1, 4), "obj_bboxes")
    checkshape.check_shape(camintr[0], (3, 3), "camintr")

    person_parameters = []
    obj_mask_infos = []
    super2d_imgs = []
    with torch.no_grad():
        for image_idx, image in enumerate(images_np):
            image_hand_boxes = {
                key: boxes[image_idx]
                for key, boxes in hand_bboxes.items() if boxes is not None
            }
            frame_info = get_frame_info(
                image,
                hand_predictor,
                mask_extractor,
                sample_folder=sample_folder,
                hand_bboxes=[image_hand_boxes],
                obj_bboxes=obj_bboxes[:, image_idx],
                camintr=camintr[image_idx],
                # Save visualization of middle frame
                debug=debug and (image_idx == len(images_np) // 2),
                image_size=image_size,
            )
            person_parameters.append(frame_info["person_parameters"])
            obj_mask_infos.append(frame_info["obj_mask_infos"])
            super2d_img = viz_frame_info(frame_info,
                                         sample_folder=sample_folder,
                                         save=False)
            super2d_imgs.append(super2d_img)
        super2d_imgs = np.concatenate(super2d_imgs[::len(super2d_imgs) // 10],
                                      1)
    return person_parameters, obj_mask_infos, super2d_imgs

def get_frame_info(image,
                   hand_predictor=None,
                   mask_extractor=None,
                   sample_folder=None,
                   hand_bboxes=None,
                   obj_bboxes=None,
                   camintr=None,
                   debug=True,
                   image_size=640):
    """
    Regress frame hand pose and hand+object masks

    Arguments:
        image (np.ndarray): hand-object image
        hand_bboxes (list): [{'left_hand': np.array(4,), 'right_hand': np.array(4,)}, ...] in xywh format
        hand_predictor: Hand pose regressor
        mask_extractor: Instance segmentor
    Returns:
        frame_infos (dict): Contains person parameters and mask information
    """
    person_parameters = {}
    left_boxes = [
        boxes['left_hand'].clip(0, None) for boxes in hand_bboxes
        if (("left_hand" in boxes) and (boxes["left_hand"] is not None))
    ]
    if hand_predictor is not None:
        mocap_predictions = hand_predictor.regress(image[..., ::-1],
                                                   hand_bboxes,
                                                   add_margin=False,
                                                   debug=debug,
                                                   K=camintr,
                                                   viz_path=os.path.join(
                                                       sample_folder,
                                                       "hands.png"))
        left_preds = [pred['left_hand'] for pred in mocap_predictions]
    else:
        left_preds = None
    all_parameters = []
    if len(left_boxes) > 0:
        left_parameters = process_hand_boxes(image,
                                             hand_boxes=left_boxes,
                                             hand_preds=left_preds,
                                             mask_extractor=mask_extractor,
                                             image_size=image_size)
        all_parameters.append(left_parameters)
    if hand_predictor is not None:
        right_preds = [pred['right_hand'] for pred in mocap_predictions]
    else:
        right_preds = None

    right_boxes = [
        boxes['right_hand'].clip(0, None) for boxes in hand_bboxes
        if (("right_hand" in boxes) and (boxes["right_hand"] is not None))
    ]
    if len(right_boxes) > 0:
        right_boxes = np.stack(right_boxes)
        right_parameters = process_hand_boxes(image,
                                              hand_boxes=right_boxes,
                                              hand_preds=right_preds,
                                              mask_extractor=mask_extractor,
                                              image_size=image_size)
        all_parameters.append(right_parameters)

    for key in all_parameters[0].keys():
        if isinstance(all_parameters[0][key], str):
            # Process labels separately
            person_parameters[key] = [param[key] for param in all_parameters]
        else:
            person_parameters[key] = torch.cat(
                [param[key] for param in all_parameters])
    # Handling only 1 object
    obj_mask_infos = mask_extractor.masks_from_bboxes(
        image,
        bbox_xy_to_wh(obj_bboxes),
        pred_classes=None,
        image_size=image_size)[0]

    # Masks with -1 for occluded parts, by merging rendered and segmentation masks
    if (len(person_parameters) > 0) and ("rend" in person_parameters):
        hand_occlusions = (
            (person_parameters["rend"].sum([0, 1]).transpose(1, 0) +
             person_parameters["masks"]) > 0)
    else:
        hand_occlusions = person_parameters["masks"] > 0
    target_masks = maskutils.add_occlusions([obj_mask_infos["crop_mask"]],
                                            hand_occlusions,
                                            [obj_mask_infos["square_bbox"]])[0]
    obj_mask_infos["target_crop_mask"] = target_masks
    frame_infos = dict(
        person_parameters=person_parameters,
        obj_mask_infos=obj_mask_infos,
        image=image,
    )

    return frame_infos

def get_gt_infos(images_np,
                 annots,
                 obj_mask_infos,
                 person_parameters,
                 image_size=None,
                 sample_folder="tmp"):
    render_gt_masks(annots,
                    obj_mask_infos,
                    person_parameters,
                    image_size=image_size)
    gt2d_imgs = []
    for image, obj_params, person_params in zip(images_np, obj_mask_infos,
                                                person_parameters):
        gt2d_img = viz_frame_info(
            {
                "image": image,
                "person_parameters": person_params,
                "obj_mask_infos": obj_params,
            },
            sample_folder=sample_folder,
            save=False)
        gt2d_imgs.append(gt2d_img)
    gt2d_imgs = np.concatenate(gt2d_imgs[::len(gt2d_imgs) // 2], 1)
    return gt2d_imgs


# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.11.3
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="view-in-github" colab_type="text"
#

# + [markdown] id="N8swpNugjZ2m"
# # Pandas
#
# [Pandas](https://pandas.pydata.org/) is a widely used Python library for storing and manipulating tabular data, where feature columns may be of different types (e.g., scalar, ordinal, categorical, text). We give some examples of how to use it below.
#
# For very large datasets, you might want to use [modin](https://github.com/modin-project/modin), which provides the same pandas API but scales to multiple cores, by using [dask](https://github.com/dask/dask) or [ray](https://github.com/ray-project/ray) on the backend.

# + [markdown] id="uCLsMbr6jXHp"
# ### Install necessary libraries
#
#
#

# + id="ds8fKSWOjX6J"
# Standard Python libraries
from __future__ import absolute_import, division, print_function, unicode_literals

import os
import time
import numpy as np
import glob
import matplotlib.pyplot as plt
import PIL
import imageio

from IPython.display import display, HTML

import sklearn

import seaborn as sns;
sns.set(style="ticks", color_codes=True)

import pandas as pd
pd.set_option('precision', 2) # 2 decimal places
pd.set_option('display.max_rows', 20)
pd.set_option('display.max_columns', 30)
pd.set_option('display.width', 100) # wide windows

import xarray as xr

# + [markdown] id="fW9crfnZlUmU"
# ### Auto-mpg dataset

# + id="k6SpViZcjhAQ" colab={"base_uri": "https://localhost:8080/"} outputId="b7d6e7ff-391e-453d-d956-0b15892957b2"
url = 'https://archive.ics.uci.edu/ml/machine-learning-databases/auto-mpg/auto-mpg.data'
column_names = ['MPG','Cylinders','Displacement','Horsepower','Weight',
                'Acceleration', 'Year', 'Origin', 'Name']
df = pd.read_csv(url, names=column_names, sep='\s+', na_values="?")

# The last column (name) is a unique id for the car, so we drop it
df = df.drop(columns=['Name'])

df.info()

# + [markdown] id="ryYgM3I6l1VH"
# We notice that there are only 392 horsepower rows, but 398 of the others.
# This is because the HP column has 6 **missing values** (also called NA, or
# not available).
# There are 3 main ways to deal with this:
# - Drop the rows with any missing values using dropna()
# - Drop any columns with any missing values using drop()
# - Replace the missing vales with some other valye (eg the median) using fillna. (This is called missing value imputation.)
# For simplicity, we adopt the first approach.
#

# + id="yq2XTXk4lwFk" colab={"base_uri": "https://localhost:8080/"} outputId="4ebb0303-3bbc-4525-8b8a-1fb8e6a3b4b2"
# Ensure same number of rows for all features.
df = df.dropna()
df.info()

# + id="BdwRK6ovl4Qy" colab={"base_uri": "https://localhost:8080/", "height": 277} outputId="98f0fb44-f4bd-46af-cc51-50d321ec42cc"
# Summary statistics
df.describe(include='all')

# + id="mJES73Pil6da" colab={"base_uri": "https://localhost:8080/"} outputId="b1b43078-10b3-40d7-cf3e-7e69b666dd0d"
# Convert Origin feature from int to categorical factor
df['Origin'] = df.Origin.replace([1,2,3],['USA','Europe','Japan'])
df['Origin'] = df['Origin'].astype('category')

# Let us check the categories (levels)
print(df['Origin'].cat.categories)

# Let us check the datatypes of all the features
print(df.dtypes)

# + id="ElSW55EVl8mE" colab={"base_uri": "https://localhost:8080/", "height": 191} outputId="1d8737cc-7ff7-4c19-f3f3-e768bac32454"
# Let us inspect the data. We see meaningful names for Origin.
df.tail()

# + id="aKkS27k5l-2K" colab={"base_uri": "https://localhost:8080/"} outputId="f8b5dab0-1ef5-4a87-cb82-b3b77081e6d5"
# Create latex table from first 5 rows
tbl = df[-5:].to_latex(index=False, escape=False)
print(tbl)

# + id="rrv5YoXkmFUJ" colab={"base_uri": "https://localhost:8080/", "height": 285} outputId="bc9e42f2-a82d-41b5-88a7-5b046f2569d9"
# Plot mpg distribution for cars from different countries of origin
data = pd.concat( [df['MPG'], df['Origin']], axis=1)
fig, ax = plt.subplots()
ax = sns.boxplot(x='Origin', y='MPG', data=data)
ax.axhline(data.MPG.mean(), color='r', linestyle='dashed', linewidth=2)
#plt.savefig(os.path.join(figdir, 'auto-mpg-origin-boxplot.pdf'))
plt.show()

# + id="Zhd7OaXwmHPJ" colab={"base_uri": "https://localhost:8080/", "height": 285} outputId="98aaf055-7058-4690-b971-e4535e003357"
# Plot mpg distribution for cars from different years
data = pd.concat( [df['MPG'], df['Year']], axis=1)
fig, ax = plt.subplots()
ax = sns.boxplot(x='Year', y='MPG', data=data)
ax.axhline(data.MPG.mean(), color='r', linestyle='dashed', linewidth=2)
#plt.savefig(os.path.join(figdir, 'auto-mpg-year-boxplot.pdf'))
plt.show()

# + [markdown] id="KpKG_NgomPmd"
# ### Iris dataset

# + id="vZwptsuWmJIY" colab={"base_uri": "https://localhost:8080/"} outputId="2686df6a-2587-447c-ade4-5b16ddf41a84"
# Get the iris dataset and look at it
from sklearn.datasets import load_iris
iris = load_iris()
# show attributes of this object
print(dir(iris))

# Extract numpy arrays
X = iris.data
y = iris.target
print(np.shape(X)) # (150, 4)
print(np.c_[X[0:3,:], y[0:3]]) # concatenate columns

# + id="iAdAEXxrmWIA" colab={"base_uri": "https://localhost:8080/"} outputId="520cbfa9-d235-430b-d775-5092e0dd7c2d"
# The data is sorted by class. Let's shuffle the rows.
N = np.shape(X)[0]
rng = np.random.RandomState(42)
perm = rng.permutation(N)
X = X[perm]
y = y[perm]
print(np.c_[X[0:3,:], y[0:3]])

# + id="4eWOuMkJmYCg" colab={"base_uri": "https://localhost:8080/", "height": 363} outputId="301f3c15-dec8-4a2f-836c-83b8399439cd"
# Convert to pandas dataframe
df = pd.DataFrame(data=X, columns=['sl', 'sw', 'pl', 'pw'])
# create column for labels
df['label'] = pd.Series(iris.target_names[y], dtype='category')

# Summary statistics
df.describe(include='all')

# + id="zbHZKIdRmZpA" colab={"base_uri": "https://localhost:8080/", "height": 191} outputId="89063ae9-319b-4356-ec48-0fc38399abb6"
# Peak at the data
df.head()

# + id="OIgZSoOambhw" colab={"base_uri": "https://localhost:8080/"} outputId="200edc2c-9def-4fe7-998d-08df9d71f271"
# Create latex table from first 5 rows
tbl = df[:6].to_latex(index=False, escape=False)
print(tbl)

# + id="1EMJx8FOmdOf" colab={"base_uri": "https://localhost:8080/", "height": 547} outputId="74107c7c-f4a1-420c-88db-5da5ac302b88"
# 2d scatterplot
#https://seaborn.pydata.org/generated/seaborn.pairplot.html
import seaborn as sns;
sns.set(style="ticks", color_codes=True)
# Make a dataframe with nicer labels for printing
#iris_df = sns.load_dataset("iris")
iris_df = df.copy()
iris_df.columns = iris['feature_names'] + ['label']
g = sns.pairplot(iris_df, vars = iris_df.columns[0:3] , hue="label")
#save_fig("iris-scatterplot.pdf")
plt.show()

# + [markdown] id="b18tmIORmjmw"
# ### Boston housing dataset

# + id="1IBp-q42me-o" colab={"base_uri": "https://localhost:8080/", "height": 277} outputId="75464b36-c127-489b-add1-ac1f9d1e633c"
# Load data (creates numpy arrays)
boston = sklearn.datasets.load_boston()
X = boston.data
y = boston.target

# Convert to Pandas format
df = pd.DataFrame(X)
df.columns = boston.feature_names
df['MEDV'] = y.tolist()

df.describe()

# + id="2I1RZQJyml2v" colab={"base_uri": "https://localhost:8080/", "height": 300} outputId="1375eb86-5ba6-4cc4-8854-006a56a2d093"
# plot marginal histograms of each column (13 features, 1 response)
plt.figure()
df.hist()
plt.show()

# + id="TvM6U_lYmzF-" colab={"base_uri": "https://localhost:8080/", "height": 737} outputId="a00e988c-d376-4ce1-e058-14b8c6799105"
# scatter plot of response vs each feature
nrows = 3; ncols = 4;
fig, ax = plt.subplots(nrows=nrows, ncols=ncols, sharey=True, figsize=[15, 10])
plt.tight_layout()
plt.clf()
for i in range(0,12):
    plt.subplot(nrows, ncols, i+1)
    plt.scatter(X[:,i], y)
    plt.xlabel(boston.feature_names[i])
    plt.ylabel("house price")
    plt.grid()
#save_fig("boston-housing-scatter.pdf")
plt.show()

# + [markdown] id="H-KJ-inpAnGi"
# # Xarray
#
# [Xarray](http://xarray.pydata.org/en/stable/quick-overview.html) generalizes pandas to multi-dimensional indexing. Put another way, xarray is a way to create multi-dimensional numpy arrays, where each dimension has a label (instead of having to remember axis ordering), and each value along each dimension can also have a specified set of allowable values (instead of having to be an integer index). This allows for easier slicing and dicing of data. We give some examples below.
#

# + id="JvQPxOjCBMT3"
import xarray as xr

# + [markdown] id="w2m5WTa1pRdh"
# ## DataArray
#
#
# A data-array is for storing a single, multiply-indexed variable. It is a generalization of a Pandas series.

# + [markdown] id="lj_1sc4pBYom"
# We create a 2d DataArray, where the first dimension is labeled 'gender' and has values 'male', 'female' and 'other' for its coordinates; the second dimension is labeled 'age', and has integer coordinates. We also associate some arbitrary attributes to the array.

# + colab={"base_uri": "https://localhost:8080/", "height": 236} id="c6Jug1-BAoop" outputId="4692eaaf-1a4d-4d89-e5db-4e667be83566"
X = np.reshape(np.arange(15), (3,5))
print(X)
attrs = {'authors': ['John', 'Mary'], 'date': '2021-01-29'}
data = xr.DataArray(X,
                    dims=("gender", "age"),
                    coords={"gender": ["male", "female", "other"]},
                    attrs = attrs)
data

# + colab={"base_uri": "https://localhost:8080/", "height": 201} id="MtD8hTRjB-IT" outputId="965a9fcf-1eb3-4ef4-cb98-7da6a899e348"
# select on dimension name and coordinate label
data.sel(gender="female")

# + colab={"base_uri": "https://localhost:8080/"} id="JUyA35LDDy3Q" outputId="8f6d0fe9-172d-4b78-b1ef-f2d4e209cc79"
v = data.sel(gender="female").values
print(v)
assert np.all(v == X[1,:])

# + id="DEflsVfPM0-S" colab={"base_uri": "https://localhost:8080/", "height": 183} outputId="cb3491b0-82d7-4dfe-87bd-1faac809b6d4"
# the dict indexing method is equivalent to  data.sel(gender="other")
data.loc[dict(gender="other")]
data

# + id="RL4muFaPLSaa" colab={"base_uri": "https://localhost:8080/", "height": 183} outputId="d62b07d8-436d-4aa5-ccc3-8a3f41606f1b"
# For assignment, we need to use the dict indexing method
data.loc[dict(gender="other")] = 42
data

# + colab={"base_uri": "https://localhost:8080/", "height": 183} id="sj2-e_0dC89_" outputId="3e6ef850-57a7-4023-b291-36118c266a51"
# select on dimension name and coordinate value
data.sel(age=3)

# + colab={"base_uri": "https://localhost:8080/"} id="rmRwUXMbDi5A" outputId="f32e90c8-6a38-44b3-9dce-70ad57704754"
v = data.sel(age=3).values
print(v)
assert np.all(v == X[:,3])

# + colab={"base_uri": "https://localhost:8080/", "height": 183} id="UB7-hyXsD_Rv" outputId="e99aaae1-141d-497b-d3ba-e014404a99c6"
# select on dimension name and integer index
data.isel(gender=1)

# + colab={"base_uri": "https://localhost:8080/"} id="89Tr7KUJEV8U" outputId="e5429eb6-0a2f-45a2-dd08-982a764d3eeb"
# regular numpy indexing
data[1,:].values

# + [markdown] id="50gRESr-EiCj"
# We can also do [broadcasting](http://xarray.pydata.org/en/stable/computation.html#broadcasting-by-dimension-name) on xarrays.

# + colab={"base_uri": "https://localhost:8080/", "height": 150} id="UIxZiguRFqTt" outputId="6148efd4-02cf-4630-e896-19c8531946be"
a = xr.DataArray([1, 2], [("x", ["a", "b"])])
a

# + colab={"base_uri": "https://localhost:8080/", "height": 150} id="ZXLoJFYiJTGa" outputId="131deed1-4e59-4acc-890c-8305b3055902"
b = xr.DataArray([-1, -2, -3], [("y", [10, 20, 30])])
b

# + colab={"base_uri": "https://localhost:8080/", "height": 188} id="nonx9ZF-JYjx" outputId="008253dd-56ce-4221-df2a-4bb19ab7505b"

c = a*b
print(c.shape)
c

# + colab={"base_uri": "https://localhost:8080/", "height": 150} id="QrXOmKt_Jfx3" outputId="05b8face-a188-454f-a5ca-e565e92cd832"
data2 = xr.DataArray([10,20,30],dims=("gender"), coords={"gender": ["male", "female", "other"]})
data2

# + colab={"base_uri": "https://localhost:8080/", "height": 256} id="C--JRtYkJrp8" outputId="c02dadca-0542-4c6b-a437-2c895b6d138d"
c = data + data2
print(c.shape)
print(c.sel(gender="female"))
c

# + id="dc_r0HKlJswB"

# + [markdown] id="kaVvJ-nDpToC"
# ## DataSet
#
# An xarray DataSet is a collection of related DataArrays.
#

# + id="nI-uACUWpVNx"


from .detector3d_template import Detector3DTemplate
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from ...ops.iou3d_nms.iou3d_nms_utils import boxes_iou3d_gpu
from ..model_utils.meter_utils import AverageMeter

class Self_VoxelSECOND(Detector3DTemplate):
    def __init__(self, model_cfg, num_class, dataset):
        super().__init__(model_cfg=model_cfg, num_class=num_class, dataset=dataset)
        self.module_list = self.build_networks()
        self.forward_ret_dict = {}

    def forward(self, batch_dict):
        if self.training:
            batch_dict = self.vfe(batch_dict)
            batch_dict = self.backbone_3d(batch_dict)
            batch_dict = self.map_to_bev_module(batch_dict)
            batch_dict = self.backbone_2d(batch_dict)
            if not self.model_cfg.SELFKD:
                batch_dict = self.dense_head(batch_dict)
                batch_dict = self.roi_head(batch_dict)
            else:
                batch_dict = self.dense_head(batch_dict)
                batch_dict = self.roi_head(batch_dict)
                self.forward_ret_dict['stage_one_cls'] = batch_dict['stage_one_cls']
                self.forward_ret_dict['stage_one_box'] = batch_dict['stage_one_box']
                self.forward_ret_dict['stage_two_cls'] = batch_dict['batch_cls_preds']
                self.forward_ret_dict['stage_two_box'] = batch_dict['batch_box_preds']
                self.forward_ret_dict['stage_two_label'] = batch_dict['roi_labels']
                self.forward_ret_dict['cur_epoch'] = batch_dict['cur_epoch']

            if self.model_cfg.ATTENTION_DISTILLATION:
                self.forward_ret_dict['attention_group1'] = batch_dict['attention_group1']
                self.forward_ret_dict['attention_group2'] = batch_dict['attention_group2']
                self.forward_ret_dict['attention_label']  = batch_dict['attention_group3']

            loss, tb_dict, disp_dict = self.get_training_loss()

            ret_dict = {
                'loss': loss
            }
            return ret_dict, tb_dict, disp_dict
        else:
            batch_dict = self.vfe(batch_dict)
            batch_dict = self.backbone_3d(batch_dict)
            batch_dict = self.map_to_bev_module(batch_dict)
            batch_dict = self.backbone_2d(batch_dict)
            batch_dict = self.dense_head(batch_dict)
            #batch_dict = self.roi_head(batch_dict)
            pred_dicts, recall_dicts = self.post_processing(batch_dict)
            return pred_dicts, recall_dicts

    def get_training_loss(self):
        disp_dict = {}
        loss = 0
        if not self.model_cfg.SELFKD:
            loss_rpn, tb_dict  = self.dense_head.get_loss()
            loss_rcnn, tb_dict = self.roi_head.get_loss(tb_dict)
            loss = loss + loss_rpn + loss_rcnn
        else:
            loss_rpn, tb_dict   = self.dense_head.get_loss()
            loss_rcnn, tb_dict  = self.roi_head.get_loss(tb_dict)
            loss_kd, tb_kd_dict = self.get_self_distillation_loss()
            loss = loss + loss_rpn + loss_rcnn + loss_kd
            tb_dict.update(tb_kd_dict)

        if self.model_cfg.ATTENTION_DISTILLATION:
            loss_attn, tb_attn_dict  = self.get_atteintion_self_distillation_loss()
            loss += loss_attn
            tb_dict.update(tb_attn_dict)

        return loss, tb_dict, disp_dict

    def get_self_distillation_loss(self):
        # get parameters.....
        epoch         = self.forward_ret_dict['cur_epoch']
        stage_one_cls = F.sigmoid(self.forward_ret_dict['stage_one_cls'])
        stage_one_box = self.forward_ret_dict['stage_one_box']
        stage_two_cls = F.sigmoid(self.forward_ret_dict['stage_two_cls'])
        stage_two_box = self.forward_ret_dict['stage_two_box']
        stage_two_label = self.forward_ret_dict['stage_two_label']
        batch_sz = stage_one_box.shape[0]

        # filter once by confidence
        one_confidence, _ = torch.max(stage_one_cls, dim=-1)
        two_confidence, _ = torch.max(stage_two_cls, dim=-1)
        if epoch < 10:
            THRESH = (epoch + 1) / 20 + 0.05
        else:
            THRESH = 0.7
        one_mask = one_confidence > THRESH
        two_mask = two_confidence > THRESH
        stage_one_cls, stage_one_box = stage_one_cls[one_mask], stage_one_box[one_mask]
        if stage_one_cls.shape == 0:
            tb_dict = {
                'Self-kd-cls-loss': 0
            }
            return 0, tb_dict
        stage_two_cls, stage_two_box, stage_two_label = stage_two_cls[two_mask], stage_two_box[two_mask], stage_two_label[two_mask]
        one_hot_targets = torch.zeros(
            *list(stage_two_cls.squeeze(-1).shape), self.num_class + 1, dtype=stage_two_cls.dtype, device=stage_two_cls.device
        )
        one_hot_targets.scatter_(-1, stage_two_label.unsqueeze(dim=-1).long(), stage_two_cls)
        one_hot_targets = one_hot_targets[:, 1:]
        # soft the one_hot_targets
        zero_mask = one_hot_targets == 0
        one_hot_targets[zero_mask] = 0.01

        # match the one_stage_box and two_stage_box
        num_teacher_box = stage_two_box.shape[0]

        teacher_centers = stage_two_box[:, :3]
        student_centers = stage_one_box[:, :3]

        with torch.no_grad():
            teacher_class = stage_two_label.unsqueeze(-1)
            student_class = torch.max(stage_one_cls, dim=-1, keepdim=True)[1]
            not_same_class = (teacher_class != student_class.T).float() # [Nt, Ns]
            MAX_DISTANCE = 1000000
            dist = teacher_centers[:, None, :] - student_centers[None, :, :] # [Nt, Ns, 3]
            dist = (dist ** 2).sum(-1) # [Nt, Ns]
            dist += not_same_class * MAX_DISTANCE # penalty on different classes
            student_dist_of_teacher, student_index_of_teacher = dist.min(1) # [Nt]
            # different from standard sess, we only consider distance<1m as matching
            MATCHED_DISTANCE = 1
            matched_student_mask = (student_dist_of_teacher < MATCHED_DISTANCE).float().unsqueeze(-1) # [Nt, 1]

        matched_student_cls_preds = stage_one_cls[student_index_of_teacher]
        cls_loss = F.mse_loss(matched_student_cls_preds, one_hot_targets, reduction='none')
        cls_loss = (cls_loss * matched_student_mask).sum() / (num_teacher_box * batch_sz)
        tb_dict = {
            'Self-kd-cls-loss': cls_loss.item()
        }
        return cls_loss, tb_dict

    def get_atteintion_self_distillation_loss(self):
        cost_function = nn.MSELoss()
        attention_label = self.forward_ret_dict['attention_label']
        attnetion_group1= self.forward_ret_dict['attention_group1']
        attention_group2= self.forward_ret_dict['attention_group2']
        loss = self.model_cfg.ATTENTION_WEIGHT * (cost_function(attention_label, attnetion_group1) + cost_function(attention_label, attention_group2))
        tb_dict = {
            'attention_loss': loss.item()
        }
        return loss, tb_dict


#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
Functions for use with the Harvard-NASA Emissions Component (HEMCO)

Use help(<name of function>) to get details on a particular function.

"""
# - Required modules:
# I/O / Low level
import os
import sys
import glob
import pandas as pd
import logging
# Math/Analysis
import numpy as np
# Time
import time
import datetime as datetime

def rm_eruptive_volcancos_from_files(dates=None, sdate=None, edate=None,
                                     folder=None, output_folder=None):
    """
    Remove the eruptive volcanos from the input HEMCO emission files
    """
    # Set dates to update files as NASA ATom4, unless others provided
    if isinstance(dates, type(None)):
        if isinstance(sdate, type(datetime.datetime)):
            sdate = datetime.datetime(2018, 1, 1)
        if isinstance(edate, type(datetime.datetime)):
            edate = datetime.datetime(2018, 6, 1)
        dates = pd.date_range(sdate, edate)
    assert folder != type(None), 'ABORTING: folder for volcano files needed'
    # Write out the new files to the same folder if a new one is not given
    if isinstance(output_folder, type(None)):
        output_folder = folder
    # Helper functions for processing volcano files

    def get_volc_subfolder4dt(folder=None, dt=None):
        """ Get the volcano folder for a specific datetime """
        year = dt.year
        month = dt.month
        folder = '{}/{}/{:0>2}/'.format(folder, year, month)
        return folder

    def get_volc_filename4dt(dt):
        """ Get the filename for a specific datetime """
        dt_str = dt.strftime('%Y%m%d')
        return 'so2_volcanic_emissions_Carns.{}.rc'.format(dt_str)

    def get_volc_file_lines4date(dt, folder=None):
        """ Open a volcano file for given date and return its lines """
        # Get string contain filename string and another for folder
        filename = get_volc_filename4dt(dt=dt)
        folder2use = get_volc_subfolder4dt(dt=dt, folder=folder)
        # Extract lines
        lines = read_lines_from_txt_file(filename=filename,
                                         folder=folder2use)
        return lines

    def rm_eruption_from_volc_file_lines(lines, skiplines=4, dt_str='',
                                         verbose=True, debug=False):
        """
        Remove the eruptive volcano lines from the HEMCO file

        Notes
        ------
         - The columns (and units) in the files are:
            LAT (-90,90), LON (-180,180), SULFUR [kg S/s], ELEVATION [m],
            CLOUD_COLUMN_HEIGHT [m]
        """
        # Local variables
        pstr1 = "NOTE: rm'd emission of {} kg S/s ({}) - ({}N, {}E, {}m, {}m)"
        pstr2 = 'WARNING: line inc. without check: {}'
        NewLines = []
        # Now loop lines and only save those without eruptive volcanoes
        for n_line, line in enumerate(lines):
            include_line = True
            if (n_line+1) > skiplines:
                try:
                    # Get elevation and emission height
                    tmp_line = line.strip().split()
                    LAT = tmp_line[0]
                    LON = tmp_line[1]
                    S = tmp_line[2]  # SULFUR
                    ELEV = tmp_line[3]  # ELEVATION
                    CLOUD = tmp_line[4]  # CLOUD_COLUMN_HEIGHT
                    # If not equal, then EXCLUDE the line
                    if ELEV != CLOUD:
                        if verbose:
                            print(pstr1.format(S, dt_str, LAT, LON, ELEV,
                                               CLOUD))
                        include_line = False
                except IndexError:
                    if debug:
                        print(pstr2.format(line.strip()))
            else:
                if debug:
                    print(pstr2.format(line.strip()))
            if include_line:
                NewLines += [line]
        return NewLines

    # - Now loop dates and re write files
    for dt in dates:
        dt_str = dt.strftime('%Y/%m/%d')
        # Get lines of volcano file for date
        lines = get_volc_file_lines4date(dt=dt, folder=folder)
        # Remove the eruptive volcanoes from the file
        NewLines = rm_eruption_from_volc_file_lines(lines=lines, dt_str=dt_str)
        if len(lines) != len(NewLines):
            print('WARNING: # of lines updated for {}'.format(dt_str))
#        SubFolder='VOLCANO_NO_ERUPT/v2019-08/'
        # Get string contain filename string
        filename = get_volc_filename4dt(dt=dt)
        folder2use = get_volc_subfolder4dt(dt=dt, folder=output_folder)
        # Save out the new file
        write_lines2txt_file(NewLines, folder=folder2use, filename=filename)


"""
Created on Oct 12, 2017

@author: Siyuan Qi

Description of the file.

"""

import os
import time
import json
import pickle
import warnings
import datetime

import numpy as np
import scipy.misc
import cv2
import torch
import torch.autograd
import torchvision.models

# import vsrl_utils as vu
import hico_config
import roi_pooling
import roi_feature_model
import metadata

import matplotlib.pyplot as plt

feature_mode = 'resnet'

part_ids = {'Torso': [1, 2],
            'Right Hand': [3],
            'Left Hand': [4],
            'Left Foot': [5],
            'Right Foot': [6],
            'Upper Leg Right': [7, 9],
            'Upper Leg Left': [8, 10],
            'Lower Leg Right': [11, 13],
            'Lower Leg Left': [12, 14],
            'Upper Arm Left': [15, 17],
            'Upper Arm Right': [16, 18],
            'Lower Arm Left': [19, 21],
            'Lower Arm Right': [20, 22],
            'Head': [23, 24]
            }

part_names = list(part_ids.keys())

def get_model(paths, feature_type):
    if feature_type == 'vgg':
        feature_network = roi_feature_model.Vgg16(num_classes=len(metadata.action_classes))
    elif feature_type == 'resnet':
        feature_network = roi_feature_model.Resnet152(num_classes=len(metadata.action_classes))
    elif feature_type == 'densenet':
        feature_network = roi_feature_model.Densenet(num_classes=len(metadata.action_classes))
    else:
        raise ValueError('feature type not recognized')

    if feature_type.startswith('alexnet') or feature_type.startswith('vgg'):
        feature_network.features = torch.nn.DataParallel(feature_network.features)
        feature_network.cuda()
    else:
        feature_network = torch.nn.DataParallel(feature_network).cuda()

    checkpoint_dir = os.path.join(paths.tmp_root, 'checkpoints', 'hico', 'finetune_{}'.format(feature_type))
    best_model_file = os.path.join(checkpoint_dir, 'model_best.pth')
    checkpoint = torch.load(best_model_file)
    feature_network.load_state_dict(checkpoint['state_dict'])

    transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor(),torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])])

    return feature_network, transform

def combine_box(box1, box2):
    return np.hstack((np.minimum(box1[:2], box2[:2]), np.maximum(box1[2:], box2[2:])))

def get_info(paths):
    hico_path = paths.data_root
    hico_voc_path = os.path.join(hico_path, 'Deformable-ConvNets/data/hico/VOC2007')
    feature_path = os.path.join(hico_path, 'processed', 'features_roi_vgg')

    # image_list_file = os.path.join(hico_voc_path, 'ImageSets/Main/test.txt')
    # det_res_path = os.path.join(hico_path, 'Deformable-ConvNets/output/rfcn_dcn/hico/hico_detect/2007_test',
    #                             'hico_detect_test_detections.pkl')

    image_list_file = os.path.join(hico_voc_path, 'ImageSets/Main/trainvaltest.txt')
    det_res_path = os.path.join(hico_path, 'Deformable-ConvNets/output/rfcn_dcn/hico/hico_detect/2007_trainvaltest',
                                'hico_detect_trainvaltest_detections.pkl')

    classes = ['__background__',  # always index 0
               'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat',
               'traffic_light', 'fire_hydrant', 'stop_sign', 'parking_meter', 'bench', 'bird', 'cat', 'dog', 'horse',
               'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie',
               'suitcase', 'frisbee', 'skis', 'snowboard', 'sports_ball', 'kite', 'baseball_bat', 'baseball_glove',
               'skateboard', 'surfboard', 'tennis_racket', 'bottle', 'wine_glass', 'cup', 'fork', 'knife', 'spoon',
               'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot_dog', 'pizza', 'donut',
               'cake', 'chair', 'couch', 'potted_plant', 'bed', 'dining_table', 'toilet', 'tv', 'laptop', 'mouse',
               'remote', 'keyboard', 'cell_phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book',
               'clock', 'vase', 'scissors', 'teddy_bear', 'hair_drier', 'toothbrush']

    return hico_path, hico_voc_path, det_res_path, feature_path, classes, image_list_file

def extract_features(paths, imageset):
    input_h, input_w = 244, 244
    feature_size = (7, 7)
    adaptive_max_pool = roi_pooling.AdaptiveMaxPool2d(*feature_size)

    hico_path, hico_voc_path, det_res_path, feature_path, classes, image_list_file = get_info(paths)
    if not os.path.exists(feature_path):
        os.makedirs(feature_path)
    vgg16, transform = get_model(paths, feature_mode)

    det_res = pickle.load(open('/home/tengyu/Documents/mmdetection/outputs/hico-det.%s.pkl'%imageset, 'rb'))

    print('Extracting %s'%imageset)

    total = len(list(det_res.keys()))
    t0 = time.time()

    for i_image, img_fn in enumerate(det_res.keys()):

        if not os.path.exists(os.path.join('/home/tengyu/Documents/densepose/DensePoseData/infer_out/hico-det/%s/%s.pkl'%(imageset, img_fn))):
            continue

        if os.path.exists(os.path.join(feature_path, '{}_edge_features'.format(img_fn))):
            continue

        # Read image
        image_path = os.path.join(hico_path, 'images/', '{}2015'.format(imageset), img_fn)
        assert os.path.exists(image_path)
        original_img = cv2.imread(image_path)

        obj_boxes_all = np.empty((0,4))
        obj_classes_all = list()
        part_boxes_all = np.empty((0,4))
        part_classes_all = list()
        part_human_id = list()
        edge_boxes_all = np.empty((0,4))
        edge_human_id = list()

        # plt.imshow(original_img)

        # Read object and human detections
        for c in range(2, len(classes)):
            for detection in det_res[img_fn][c-1]:
                if detection[4] > 0.7:
                    y0,x0,y1,x1 = detection[0], detection[1], detection[2], detection[3]
                    # plt.plot([y0,y0,y1,y1,y0], [x0,x1,x1,x0,x0], c='red')
                    obj_boxes_all = np.vstack((obj_boxes_all, np.array(detection[:4])[np.newaxis, ...]))
                    obj_classes_all.append(c-1)
        if len(obj_classes_all) == 0:
            continue

        boxes, bodies = pickle.load(open(os.path.join('/home/tengyu/Documents/densepose/DensePoseData/infer_out/hico-det/%s/%s.pkl'%(imageset, img_fn)), 'rb'), encoding='latin-1')

        for human_id in range(len(boxes[1])):
            if boxes[1][human_id][4] < 0.7:
                continue
            for part_id, part_name in enumerate(part_names):
                x, y = np.where(np.isin(bodies[1][human_id], part_ids[part_name]))
                x = x + boxes[1][human_id][1]
                y = y + boxes[1][human_id][0]
                if len(x) > 0:
                    x0, x1, y0, y1 = x.min(), x.max(), y.min(), y.max()
                    part_boxes_all = np.vstack([part_boxes_all, np.array([[y0,x0,y1,x1]])])
                    part_classes_all.append(part_id)
                    part_human_id.append(human_id)

                    # plt.plot([y0,y0,y1,y1,y0], [x0,x1,x1,x0,x0], c='blue')

                    # Add edges
                    for obj_box in obj_boxes_all:
                        edge_box = combine_box(obj_box, part_boxes_all[-1,:])
                        edge_human_id.append(human_id)
                        edge_boxes_all = np.vstack([edge_boxes_all, [edge_box]])

        # plt.show()

        # Get image feature by applying VGG to ROI (roi_vgg)
        if feature_mode == 'resnet':
            # roi_features = np.empty((det_boxes_all.shape[0], 512, 7, 7))
            obj_features = np.zeros((obj_boxes_all.shape[0], 200))
            part_features = np.zeros((part_boxes_all.shape[0], 200))
            edge_features = np.zeros((edge_boxes_all.shape[0], 200))
            for i_box in range(obj_boxes_all.shape[0]):
                obj = obj_boxes_all[i_box, :].astype(int)
                obj_image = original_img[obj[1]:obj[3]+1, obj[0]:obj[2]+1, :]
                # print(classes[obj_classes_all[i_box]+1])
                # plt.imshow(obj_image)
                # plt.show()
                obj_image = transform(cv2.resize(obj_image, (input_h, input_w), interpolation=cv2.INTER_LINEAR))
                obj_image = torch.autograd.Variable(obj_image.unsqueeze(0)).cuda()
                feat, pred = vgg16(obj_image)
                obj_features[i_box, ...] = feat.data.cpu().numpy()

            for i_box in range(part_boxes_all.shape[0]):
                part = part_boxes_all[i_box, :].astype(int)
                part_image = original_img[part[1]:part[3]+1, part[0]:part[2]+1, :]
                # print(part_names[part_classes_all[i_box]])
                # plt.imshow(part_image)
                # plt.show()
                part_image = transform(cv2.resize(part_image, (input_h, input_w), interpolation=cv2.INTER_LINEAR))
                part_image = torch.autograd.Variable(part_image.unsqueeze(0)).cuda()
                feat, pred = vgg16(part_image)
                part_features[i_box, ...] = feat.data.cpu().numpy()

            for i_box in range(edge_boxes_all.shape[0]):
                edge = edge_boxes_all[i_box, :].astype(int)
                edge_image = original_img[edge[1]:edge[3]+1, edge[0]:edge[2]+1, :]
                # plt.imshow(roi_image)
                # plt.show()
                edge_image = transform(cv2.resize(edge_image, (input_h, input_w), interpolation=cv2.INTER_LINEAR))
                edge_image = torch.autograd.Variable(edge_image.unsqueeze(0)).cuda()
                feat, pred = vgg16(edge_image)
                edge_features[i_box, ...] = feat.data.cpu().numpy()

        # Get image feature by extracting ROI from VGG feature (vgg_roi)
        if feature_mode == 'vgg_roi':
            scale_h = feature_size[0]/float(original_img.shape[0])
            scale_w = feature_size[1]/float(original_img.shape[1])
            transform_img = transform(cv2.resize(original_img, (input_h, input_w), interpolation=cv2.INTER_LINEAR))
            transform_img = torch.autograd.Variable(transform_img.unsqueeze(0)).cuda()
            img_feature = vgg16(transform_img)

            obj_features = np.empty((obj_boxes_all.shape[0], 512, feature_size[0], feature_size[1]))
            objs = np.copy(obj_boxes_all)
            objs[:, 0] *= scale_w
            objs[:, 2] *= scale_w
            objs[:, 1] *= scale_h
            objs[:, 3] *= scale_h

            for i_box in range(objs.shape[0]):
                obj = objs[i_box, :].astype(int)
                obj_feature = adaptive_max_pool(img_feature[..., obj[1]:(obj[3] + 1), obj[0]:(obj[2] + 1)])
                obj_features[i_box, :, :, :] = obj_feature.data.cpu().numpy()

            part_features = np.empty((part_boxes_all.shape[0], 512, feature_size[0], feature_size[1]))
            parts = np.copy(part_boxes_all)
            parts[:, 0] *= scale_w
            parts[:, 2] *= scale_w
            parts[:, 1] *= scale_h
            parts[:, 3] *= scale_h

            for i_box in range(parts.shape[0]):
                part = parts[i_box, :].astype(int)
                part_feature = adaptive_max_pool(img_feature[..., part[1]:(part[3] + 1), part[0]:(part[2] + 1)])
                part_features[i_box, :, :, :] = part_feature.data.cpu().numpy()

            edge_features = np.empty((edge_boxes_all.shape[0], 512, feature_size[0], feature_size[1]))
            edges = np.copy(edge_boxes_all)
            edges[:, 0] *= scale_w
            edges[:, 2] *= scale_w
            edges[:, 1] *= scale_h
            edges[:, 3] *= scale_h

            for i_box in range(edges.shape[0]):
                edge = edges[i_box, :].astype(int)
                edge_feature = adaptive_max_pool(img_feature[..., edge[1]:(edge[3] + 1), edge[0]:(edge[2] + 1)])
                edge_features[i_box, :, :, :] = edge_feature.data.cpu().numpy()

        np.save(os.path.join(feature_path, '{}_obj_classes'.format(img_fn)), obj_classes_all)
        np.save(os.path.join(feature_path, '{}_obj_boxes'.format(img_fn)), obj_boxes_all)
        np.save(os.path.join(feature_path, '{}_obj_features'.format(img_fn)), obj_features)

        np.save(os.path.join(feature_path, '{}_part_classes'.format(img_fn)), part_classes_all)
        np.save(os.path.join(feature_path, '{}_part_boxes'.format(img_fn)), part_boxes_all)
        np.save(os.path.join(feature_path, '{}_part_human_id'.format(img_fn)), part_human_id)
        np.save(os.path.join(feature_path, '{}_part_features'.format(img_fn)), part_features)

        np.save(os.path.join(feature_path, '{}_edge_human_id'.format(img_fn)), edge_human_id)
        np.save(os.path.join(feature_path, '{}_edge_boxes'.format(img_fn)), edge_boxes_all)
        np.save(os.path.join(feature_path, '{}_edge_features'.format(img_fn)), edge_features)

        t1 = time.time()
        elapsed = t1 - t0
        ETA = elapsed / (i_image + 1) * (total - i_image - 1)
        print('\r\t%d/%d, ETA: %s'%(i_image, total, str(datetime.timedelta(seconds=ETA))), end='')
        # break

def main():
    paths = hico_config.Paths()
    imagesets = ['train', 'test']
    for imageset in imagesets:
        extract_features(paths, imageset)
        # break

if __name__ == '__main__':
    main()


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

import pycocotools.coco as coco
import numpy as np
import torch
import json
import cv2
import os
import math

from ..generic_dataset import GenericDataset
from utils.ddd_utils import compute_box_3d, project_to_image

class KITTITrackingCustom(GenericDataset):
  num_categories = 2
  default_resolution = [384, 1280]
  class_name = ['Pedestrian', 'Car']
  # negative id is for "not as negative sample for abs(id)".
  # 0 for ignore losses for all categories in the bounding box region
  # ['Pedestrian', 'Car', 'Cyclist', 'Van', 'Truck',  'Person_sitting',
  #       'Tram', 'Misc', 'DontCare']
  cat_ids = {1:1, 2:2, 3:-2, 4:-2, 5:-2, 6:-1, 7:-9999, 8:-9999, 9:0}
  max_objs = 50
  def __init__(self, opt, split):
    data_dir = os.path.join(opt.data_dir, 'kitti_tracking_custom')
    split_ = 'train' if opt.dataset_version != 'test' else 'test' #'test'
    img_dir = os.path.join(
      data_dir, 'data_tracking_image_2', '{}ing'.format(split_), 'image_02')
    ann_file_ = split_ if opt.dataset_version == '' else opt.dataset_version
    print('Warning! opt.dataset_version is not set')
    ann_path = os.path.join(
      data_dir, 'annotations', 'tracking_{}.json'.format(
        split))
    self.images = None
    super(KITTITrackingCustom, self).__init__(opt, split, ann_path, img_dir)
    self.alpha_in_degree = False
    self.num_samples = len(self.images)

    print('Loaded {} {} samples'.format(split, self.num_samples))

  def __len__(self):
    return self.num_samples

  def _to_float(self, x):
    return float("{:.2f}".format(x))

  def save_results(self, results, save_dir):
    results_dir = os.path.join(save_dir, 'results_kitti_tracking')
    if not os.path.exists(results_dir):
      os.mkdir(results_dir)

    for video in self.coco.dataset['videos']:
      video_id = video['id']
      file_name = video['file_name']
      out_path = os.path.join(results_dir, '{}.txt'.format(file_name))
      f = open(out_path, 'w')
      images = self.video_to_images[video_id]

      for image_info in images:
        img_id = image_info['id']
        if not (img_id in results):
          continue
        frame_id = image_info['frame_id']
        for i in range(len(results[img_id])):
          item = results[img_id][i]
          category_id = item['class']
          cls_name_ind = category_id
          class_name = self.class_name[cls_name_ind - 1]
          if not ('alpha' in item):
            item['alpha'] = -1
          if not ('rot_y' in item):
            item['rot_y'] = -10
          if 'dim' in item:
            item['dim'] = [max(item['dim'][0], 0.01),
              max(item['dim'][1], 0.01), max(item['dim'][2], 0.01)]
          if not ('dim' in item):
            item['dim'] = [-1, -1, -1]
          if not ('loc' in item):
            item['loc'] = [-1000, -1000, -1000]

          track_id = item['tracking_id'] if 'tracking_id' in item else -1
          f.write('{} {} {} -1 -1'.format(frame_id - 1, track_id, class_name))
          f.write(' {:d}'.format(int(item['alpha'])))
          f.write(' {:.2f} {:.2f} {:.2f} {:.2f}'.format(
            item['bbox'][0], item['bbox'][1], item['bbox'][2], item['bbox'][3]))

          f.write(' {:d} {:d} {:d}'.format(
            int(item['dim'][0]), int(item['dim'][1]), int(item['dim'][2])))
          f.write(' {:d} {:d} {:d}'.format(
            int(item['loc'][0]), int(item['loc'][1]), int(item['loc'][2])))
          f.write(' {:d} {:.2f}\n'.format(int(item['rot_y']), item['score']))
      f.close()

  def run_eval(self, results, save_dir):
    self.save_results(results, save_dir)
    os.system('python tools/eval_kitti_track/evaluate_tracking.py ' + \
              '{}/results_kitti_tracking/ {}'.format(
                save_dir, self.opt.dataset_version))