sf_gru.py

"""
The code implementation of the paper:

A. Rasouli, I. Kotseruba, and J. K. Tsotsos, "Pedestrian Action Anticipation using
Contextual Feature Fusion in Stacked RNNs", BMVC 2019.

Copyright (c) 2019 Amir Rasouli

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""

import time
import pickle

from utils import *

from PIL import Image, ImageDraw
from keras.layers import Input, Concatenate, Dense
from keras.layers.recurrent import GRU
from keras.models import Model, load_model
from keras.applications import vgg16
from keras.optimizers import Adam
from keras.layers.core import regularizers
from keras.preprocessing.image import img_to_array, load_img
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
from sklearn.metrics import roc_auc_score, roc_curve, precision_recall_curve


class SFGRU(object):
    """
     An encoder decoder model for pedestrian trajectory prediction

     Attributes:
        _num_hidden_units: Number of LSTM hidden units
        _regularizer_value: The value of L2 regularizer for training
        _regularizer: Training regularizer set as L2
        self._global_pooling: The pulling method for visual features. Options are: 'avg', 'max', 'none' (will return
                              flattened output

     Methods:
        load_images_crop_and_process: Reads the images and generate feature suquences for training
        get_poses: gets the poses for PIE dataset
        flip_pose: Flips the pose joint coordinates
        get_data_sequence: Generates data sequences
        get_data_sequence_balance: Generates data sequences and balances positive and negative samples by augmentations
        get_data: Receives the data sequences generated by the dataset interface and returns train/test data according
                  to model specifications.
        log_configs: Writes model and training configurations to a file
        train: Trains the model
        test: Tests the model
        stacked_rnn: Generates the network model
        _gru: A helper function for creating a GRU unit
     """
    def __init__(self,
                 num_hidden_units=256,
                 global_pooling='avg',
                 regularizer_val=0.0001):

        # Network parameters
        self._num_hidden_units = num_hidden_units
        self._regularizer_value = regularizer_val
        self._regularizer = regularizers.l2(regularizer_val)
        self._global_pooling = global_pooling

    # Processing images anf generate features
    def load_images_crop_and_process(self, img_sequences, bbox_sequences,
                                     ped_ids, save_path,
                                     data_type='train',
                                     crop_type='none',
                                     crop_mode='warp',
                                     crop_resize_ratio=2,
                                     regen_data=False):
        """
        Generate visual feature seuqences by reading and processing images
        :param img_sequences: Sequences of image names
        :param bbox_sequences: Sequences of bounding boxes
        :param ped_ids: Sequences of pedestrian ids
        :param save_path: The path to save the features
        :param data_type: Whether data is for training or testing
        :param crop_type: The method to crop the bounding boxes from the images
                          Options: 'bbox' crops using bounding box coordinats
                                   'context' crops using an enlarged ratio (specified
                                             by 'crop_resize_ratio') of
                                             bounding box coordinates
                                   'surround' similar to context with the difference of
                                              suppressing (by setting to gray value) areas
                                              within the original bounding box coordinate

        :param crop_mode: How the cropped image resized and padded to match the input of
                          processing network. Options are 'warp', 'same', 'pad_same',
                          'pad_resize', 'pad_fit' (see utils.py:img_pad() for more details)
        :param crop_resize_ratio: The ratio by which the image is enlarged to capture the context
                                  Used by crop types 'context' and 'surround'.
        :param regen_data: Whether regenerate the currently saved data.
        :return: Sequences of visual features
        """
        # load the feature files if exists
        print("Generating {} features crop_type={} crop_mode={}\
              \nsave_path={}, ".format(data_type, crop_type, crop_mode,
              save_path))
        convnet = vgg16.VGG16(input_shape=(224, 224, 3),
                              include_top=False, weights='imagenet')
        sequences = []
        bbox_seq = bbox_sequences.copy()
        i = -1
        for seq, pid in zip(img_sequences, ped_ids):
            i += 1
            update_progress(i / len(img_sequences))
            img_seq = []
            for imp, b, p in zip(seq, bbox_seq[i], pid):
                flip_image = False
                set_id = imp.split('/')[-3]
                vid_id = imp.split('/')[-2]
                img_name = imp.split('/')[-1].split('.')[0]

                img_save_folder = os.path.join(save_path, set_id, vid_id)
                if crop_type == 'none':
                    img_save_path = os.path.join(img_save_folder, img_name + '.pkl')
                else:
                    img_save_path = os.path.join(img_save_folder, img_name + '_' + p[0] + '.pkl')

                if os.path.exists(img_save_path) and not regen_data:
                    with open(img_save_path, 'rb') as fid:
                        try:
                            img_features = pickle.load(fid)
                        except:
                            img_features = pickle.load(fid, encoding='bytes')
                else:
                    if 'flip' in imp:
                        imp = imp.replace('_flip', '')
                        flip_image = True
                    if crop_type == 'none':
                        img_data = load_img(imp, target_size=(224, 224))
                        if flip_image:
                            img_data = img_data.transpose(Image.FLIP_LEFT_RIGHT)
                    else:
                        img_data = load_img(imp)

                        if flip_image:
                            img_data = img_data.transpose(Image.FLIP_LEFT_RIGHT)
                        if crop_type == 'bbox':
                            cropped_image = img_data.crop(list(map(int, b[0:4])))
                            img_data = img_pad(cropped_image, mode=crop_mode, size=224)
                        elif 'context' in crop_type:
                            bbox = jitter_bbox(imp, [b], 'enlarge', crop_resize_ratio)[0]
                            bbox = squarify(bbox, 1, img_data.size[0])
                            bbox = list(map(int, bbox[0:4]))
                            cropped_image = img_data.crop(bbox)
                            img_data = img_pad(cropped_image, mode='pad_resize', size=224)
                        elif 'surround' in crop_type:
                            b_org = [b[0], b[1], b[2], b[3]]
                            bbox = jitter_bbox(imp, [b], 'enlarge', crop_resize_ratio)[0]
                            bbox = squarify(bbox, 1, img_data.size[0])
                            bbox = list(map(int, bbox[0:4]))
                            draw = ImageDraw.Draw(img_data)
                            draw.rectangle([b_org[0], b_org[1], b_org[2], b_org[3]],
                                           fill=(128, 128, 128))
                            del draw
                            cropped_image = img_data.crop(bbox)
                            img_data = img_pad(cropped_image, mode='pad_resize', size=224)
                        else:
                            raise ValueError('ERROR: Undefined value for crop_type {}!'.format(crop_type))
                    image_array = img_to_array(img_data)
                    preprocessed_img = vgg16.preprocess_input(image_array)
                    expanded_img = np.expand_dims(preprocessed_img, axis=0)
                    img_features = convnet.predict(expanded_img)
                    if not os.path.exists(img_save_folder):
                        os.makedirs(img_save_folder)
                    with open(img_save_path, 'wb') as fid:
                        pickle.dump(img_features, fid, pickle.HIGHEST_PROTOCOL)

                if self._global_pooling == 'max':
                    img_features = np.squeeze(img_features)
                    img_features = np.amax(img_features, axis=0)
                    img_features = np.amax(img_features, axis=0)
                elif self._global_pooling == 'avg':
                    img_features = np.squeeze(img_features)
                    img_features = np.average(img_features, axis=0)
                    img_features = np.average(img_features, axis=0)
                else:
                    img_features = img_features.ravel()

                img_seq.append(img_features)
            sequences.append(img_seq)
        sequences = np.array(sequences)
        return sequences

    def get_pose(self, img_sequences,
                      ped_ids, file_path,
                      data_type='train'):
        """
        Reads the pie poses from saved .pkl files
        :param img_sequences: Sequences of image names
        :param ped_ids: Sequences of pedestrian ids
        :param file_path: Path to where poses are saved
        :param data_type: Whether it is for training or testing
        :return: Sequences of poses
        """

        print('\n#####################################')
        print('Getting poses %s' % data_type)
        print('#####################################')
        poses_all = []
        set_poses_list = os.listdir(file_path)
        set_poses = {}
        for s in set_poses_list:
            with open(os.path.join(file_path, s), 'rb') as fid:
                try:
                    p = pickle.load(fid)
                except:
                    p = pickle.load(fid, encoding='bytes')
            set_poses[s.split('.pkl')[0].split('_')[-1]] = p
        print(set_poses.keys())
        i = -1
        for seq, pid in zip(img_sequences, ped_ids):
            i += 1
            update_progress(i / len(img_sequences))
            pose = []
            for imp, p in zip(seq, pid):
                flip_image = False
                set_id = imp.split('/')[-3]
                vid_id = imp.split('/')[-2]
                img_name = imp.split('/')[-1].split('.')[0]
                if 'flip' in img_name:
                    img_name = img_name.replace('_flip', '')
                    flip_image = True
                k = img_name + '_' + p[0]
                if k in set_poses[set_id][vid_id].keys():
                    # [nose, neck, Rsho, Relb, Rwri, Lsho, Lelb, Lwri, Rhip, Rkne,
                    #  Rank, Lhip, Lkne, Lank, Leye, Reye, Lear, Rear, pt19]
                    if flip_image:
                        pose.append(self.flip_pose(set_poses[set_id][vid_id][k]))
                    else:
                        pose.append(set_poses[set_id][vid_id][k])
                else:
                    pose.append([0] * 36)
            poses_all.append(pose)
        poses_all = np.array(poses_all)
        return poses_all

    def flip_pose(self, pose):
        """
        Flips a given pose coordinates
        :param pose: The original pose
        :return: Flipped pose
        """
        # [nose(0,1), neck(2,3), Rsho(4,5),   Relb(6,7),   Rwri(8,9),
        # 						 Lsho(10,11), Lelb(12,13), Lwri(14,15),
        #						 Rhip(16,17), Rkne(18,19), Rank(20,21),
        #                        Lhip(22,23), Lkne(24,25), Lank(26,27),
        #						 Leye(28,29), Reye (30,31),
        #						 Lear(32,33), Rear(34,35)]
        flip_map = [0, 1, 2, 3, 10, 11, 12, 13, 14, 15, 4, 5, 6, 7, 8, 9, 22, 23, 24, 25,
                    26, 27, 16, 17, 18, 19, 20, 21, 30, 31, 28, 29, 34, 35, 32, 33]
        new_pose = pose.copy()
        flip_pose = [0] * len(new_pose)
        for i in range(len(new_pose)):
            if i % 2 == 0 and new_pose[i] != 0:
                new_pose[i] = 1 - new_pose[i]
            flip_pose[flip_map[i]] = new_pose[i]
        return flip_pose

    def get_data_sequence(self, data_raw, obs_length, time_to_event, normalize):
        """
        Generates data sequences according to the length of the observation and time to event
        :param data_raw: The data sequences from the dataset
        :param obs_length: Observation length
        :param time_to_event: Time (number of frames) to event
        :param normalize: Whether to normalize the bounding box coordinates
        :return: Processed data sequences
        """
        print('\n#####################################')
        print('Generating raw data')
        print('#####################################')
        d = {'center': data_raw['center'].copy(),
             'box': data_raw['bbox'].copy(),
             'box_org': data_raw['bbox'].copy(),
             'ped_id': data_raw['pid'].copy(),
             'acts': data_raw['activities'].copy(),
             'image': data_raw['image'].copy()}

        try:
            d['speed'] = data_raw['obd_speed'].copy()
        except KeyError:
            d['speed'] = data_raw['vehicle_act'].copy()
            print('Jaad dataset does not have speed information')
            print('Vehicle actions are used instead')

        for i in range(len(d['box'])):
            d['box'][i] = d['box'][i][- obs_length - time_to_event:-time_to_event]
            d['center'][i] = d['center'][i][- obs_length - time_to_event:-time_to_event]
            if normalize:
                d['box'][i] = np.subtract(d['box'][i][1:], d['box'][i][0]).tolist()
                d['center'][i] = np.subtract(d['center'][i][1:], d['center'][i][0]).tolist()

        if normalize:
            obs_length -= 1

        for k in d.keys():
            if k != 'box' and k != 'center':
                for i in range(len(d[k])):
                    d[k][i] = d[k][i][- obs_length - time_to_event:-time_to_event]
                d[k] = np.array(d[k])
            else:
                d[k] = np.array(d[k])
        d['acts'] = d['acts'][:, 0, :]
        return d

    def get_data_sequence_balance(self, data_raw, obs_length, time_to_event, normalize):
        """
        Generates data sequences according to the length of the observation and time to event.
        The number of positive and negative sequences are balanced. Add flipped version of underrepresented
        sequences and subsamples from the overrepresented samples to match the number of samples.
        :param dataset: The data sequences from the dataset
        :param obs_length: Observation length
        :param time_to_event: Time (number of frames) to event
        :param normalize: Whether to normalize the bounding box coordinates
        :return: Processed data sequences
        """
        print('\n#####################################')
        print('Generating balanced raw data')
        print('#####################################')
        d = {'center': data_raw['center'].copy(),
             'box': data_raw['bbox'].copy(),
             'ped_id': data_raw['pid'].copy(),
             'acts': data_raw['activities'].copy(),
             'image': data_raw['image'].copy()}

        try:
            d['speed'] = data_raw['obd_speed'].copy()
        except:
            d['speed'] = data_raw['vehicle_act'].copy()
            print('Jaad dataset does not have speed information')
            print('Vehicle actions are used instead')

        gt_labels = [gt[0] for gt in d['acts']]
        num_pos_samples = np.count_nonzero(np.array(gt_labels))
        num_neg_samples = len(gt_labels) - num_pos_samples

        # finds the indices of the samples with larger quantity
        if num_neg_samples == num_pos_samples:
            print('Positive and negative samples are already balanced')
        else:
            print('Unbalanced: \t Positive: {} \t Negative: {}'.format(num_pos_samples, num_neg_samples))
            if num_neg_samples > num_pos_samples:
                gt_augment = 1
            else:
                gt_augment = 0

            img_width = data_raw['image_dimension'][0]
            num_samples = len(d['ped_id'])
            for i in range(num_samples):
                if d['acts'][i][0][0] == gt_augment:
                    flipped = d['center'][i].copy()
                    flipped = [[img_width - c[0], c[1]]
                               for c in flipped]
                    d['center'].append(flipped)
                    flipped = d['box'][i].copy()

                    flipped = [np.array([img_width - c[2], c[1], img_width - c[0], c[3]])
                               for c in flipped]
                    d['box'].append(flipped)

                    d['ped_id'].append(data_raw['pid'][i].copy())
                    d['acts'].append(d['acts'][i].copy())
                    flipped = d['image'][i].copy()
                    flipped = [c.replace('.png', '_flip.png') for c in flipped]

                    d['image'].append(flipped)
                    if 'speed' in d.keys():
                        d['speed'].append(d['speed'][i].copy())
            gt_labels = [gt[0] for gt in d['acts']]
            num_pos_samples = np.count_nonzero(np.array(gt_labels))
            num_neg_samples = len(gt_labels) - num_pos_samples
            if num_neg_samples > num_pos_samples:
                rm_index = np.where(np.array(gt_labels) == 0)[0]
            else:
                rm_index = np.where(np.array(gt_labels) == 1)[0]

            # Calculate the difference of sample counts
            dif_samples = abs(num_neg_samples - num_pos_samples)
            # shuffle the indices
            np.random.seed(42)
            np.random.shuffle(rm_index)
            # reduce the number of indices to the difference
            rm_index = rm_index[0:dif_samples]

            # update the data
            for k in d:
                seq_data_k = d[k]
                d[k] = [seq_data_k[i] for i in range(0, len(seq_data_k)) if i not in rm_index]

            new_gt_labels = [gt[0] for gt in d['acts']]
            num_pos_samples = np.count_nonzero(np.array(new_gt_labels))
            print('Balanced:\t Positive: %d  \t Negative: %d\n'
                  % (num_pos_samples, len(d['acts']) - num_pos_samples))

        d['box_org'] = d['box'].copy()

        for i in range(len(d['box'])):
            d['box'][i] = d['box'][i][- obs_length - time_to_event:-time_to_event]
            d['center'][i] = d['center'][i][- obs_length - time_to_event:-time_to_event]
            if normalize:
                d['box'][i] = np.subtract(d['box'][i][1:], d['box'][i][0]).tolist()
                d['center'][i] = np.subtract(d['center'][i][1:], d['center'][i][0]).tolist()
        if normalize:
            obs_length -= 1
        for k in d.keys():
            if k != 'box' and k != 'center':
                for i in range(len(d[k])):
                    d[k][i] = d[k][i][- obs_length - time_to_event:-time_to_event]
                d[k] = np.array(d[k])
            else:
                d[k] = np.array(d[k])

        d['acts'] = d['acts'][:, 0, :].copy()
        return d

    def get_model_opts(self, model_opts):
        default_opts =  {'obs_input_type': ['local_box', 'local_context', 'pose', 'box','speed'],
                      'enlarge_ratio': 1.5,
                      'pred_target_type': ['crossing'],
                      'obs_length': 15,
                      'time_to_event': 60,
                      'dataset': 'pie',
                      'normalize_boxes': True}
        default_opts.update(model_opts)
        return default_opts
    def get_data(self, data_raw, model_opts):
        """
        Generates train/test data
        :param data_raw: The sequences received from the dataset interface
        :param model_opts: Model options:
                            'obs_input_type': The types of features to be used for train/test. The order
                                            in which features are named in the list defines at what level
                                            in the network the features are processed. e.g. ['local_context',
                                            pose] would behave different to ['pose', 'local_context']
                            'enlarge_ratio': The ratio (with respect to bounding boxes) that is used for processing
                                           context surrounding pedestrians.
                            'pred_target_type': Learning target objective. Currently only supports 'crossing'
                            'obs_length': Observation length prior to reasoning
                            'time_to_event': Number of frames until the event occurs
                            'dataset': Name of the dataset

        :return: Train/Test data
        """
        data = {}
        data_type_sizes_dict = {}

        model_opts = self.get_model_opts(model_opts)

        obs_length = model_opts['obs_length']
        time_to_event = model_opts['time_to_event']
        dataset = model_opts['dataset']
        eratio = model_opts['enlarge_ratio']
        data_type_keys = sorted(data_raw.keys())

        for k in data_type_keys:
            if k == 'test':
                data[k] = self.get_data_sequence(data_raw[k], obs_length, time_to_event, model_opts['normalize_boxes'])
            else:
                data[k] = self.get_data_sequence_balance(data_raw[k], obs_length, time_to_event, model_opts['normalize_boxes'])
            data_type_sizes_dict['box'] = data[k]['box'].shape[1:]

            if 'speed' in data[k].keys():
                data_type_sizes_dict['speed'] = data[k]['speed'].shape[1:]
            # Poses
            if 'pose' in model_opts['obs_input_type']:
                # get Poses

                path_to_pose, _ = get_path(save_folder='poses',
                                           dataset=dataset,
                                           save_root_folder='data/features')
                print(path_to_pose)
                data[k]['pose'] = self.get_pose(data[k]['image'],
                                           data[k]['ped_id'], data_type=k,
                                           file_path=path_to_pose)
                data_type_sizes_dict['pose'] = data[k]['pose'].shape[1:]

            # crop only bounding boxes
            if 'local_box' in model_opts['obs_input_type']:
                print('\n#####################################')
                print('Generating local box %s' % k)
                print('#####################################')
                path_to_local_boxes, _ = get_path(save_folder='local_box',
                                                  dataset=dataset,
                                                  save_root_folder='data/features')
                data[k]['local_box'] = self.load_images_crop_and_process(data[k]['image'],
                                                                         data[k]['box_org'], data[k]['ped_id'],
                                                                         data_type=k,
                                                                         save_path=path_to_local_boxes,
                                                                         crop_type='bbox',
                                                                         crop_mode='pad_resize')
                data_type_sizes_dict['local_box'] = data[k]['local_box'].shape[1:]
            # local context 2d features
            if 'local_context' in model_opts['obs_input_type']:
                print('\n#####################################')
                print('Generating local context %s' % k)
                print('#####################################')

                path_to_local_context, _ = get_path(save_folder='local_context',
                                                      dataset=dataset,
                                                      save_root_folder='data/features')
                data[k]['local_context'] = self.load_images_crop_and_process(data[k]['image'],
                                                                             data[k]['box_org'], data[k]['ped_id'],
                                                                             data_type=k,
                                                                             save_path=path_to_local_context,
                                                                             crop_type='surround',
                                                                             crop_resize_ratio=eratio)
                data_type_sizes_dict['local_context'] = data[k]['local_context'].shape[1:]

        # Create a empty dict for storing the data
        train_test_data = {}
        data_final_keys = sorted(data.keys())
        for k in data_final_keys:
            train_test_data[k] = []

        # Store the type and size of each image
        data_sizes = []
        data_types = []

        for d_type in model_opts['obs_input_type']:
            for k in data.keys():
                train_test_data[k].append(data[k][d_type])
            data_sizes.append(data_type_sizes_dict[d_type])
            data_types.append(d_type)

        # create the final data file to be returned
        for k in data_final_keys:
            train_test_data[k] = (train_test_data[k], data[k]['acts'])

        return train_test_data, data_types, data_sizes

    def log_configs(self, config_path, batch_size, epochs,
                    lr,  opts):
        """
        Logs the parameters of the model and training
        :param config_path: The path to save the file
        :param batch_size: Batch size of training
        :param epochs: Number of epochs for training
        :param lr: Learning rate of training
        :param opts: Data generation parameters (see get_data)
        """
        # Save config and training param files
        with open(config_path, 'wt') as fid:
            fid.write("####### Model options #######\n")
            for k in opts:
                fid.write("%s: %s\n" % (k, str(opts[k])))

            fid.write("\n####### Network config #######\n")
            fid.write("%s: %s\n" % ('hidden_units', str(self._num_hidden_units)))
            fid.write("%s: %s\n" % ('reg_value ', str(self._regularizer_value)))

            fid.write("\n####### Training config #######\n")
            fid.write("%s: %s\n" % ('batch_size', str(batch_size)))
            fid.write("%s: %s\n" % ('epochs', str(epochs)))
            fid.write("%s: %s\n" % ('lr', str(lr)))

        print('Wrote configs to {}'.format(config_path))

    def train(self, data_train,
              batch_size=32,
              epochs=60,
              lr=0.000005,
              model_opts=None):
        """
        Train function
        :param data_train: The raw training data from the dataset interface
        :param batch_size: Number of batches
        :param epochs: Number of epochs
        :param lr: Learning rate
        :param model_opts: Model options (see  get_data() for more details)
        :return: The path to where the final model is saved.
        """

        # Set the path for saving models
        model_folder_name = time.strftime("%d%b%Y-%Hh%Mm%Ss")
        model_path, _ = get_path(save_folder=model_folder_name,
                                 save_root_folder='data/models',
                                 file_name='model.h5')

        # Read train data
        train_val_data, data_types, data_sizes = self.get_data({'train': data_train}, model_opts)
        train_data = train_val_data['train']

        # Crate model
        train_model = self.stacked_rnn(data_types, data_sizes)

        # Train the model
        optimizer = Adam(lr=lr)
        train_model.compile(loss='binary_crossentropy', optimizer=optimizer, metrics=['accuracy'])
        history = train_model.fit(x=train_data[0],
                                  y=train_data[1],
                                  batch_size=batch_size,
                                  epochs=epochs,
                                  verbose=1)
        print('Train model is saved to {}'.format(model_path))
        train_model.save(model_path)

        # Save data options and configurations
        model_opts_path, _ = get_path(save_folder=model_folder_name,
                                      save_root_folder='data/models',
                                      file_name='model_opts.pkl')
        with open(model_opts_path, 'wb') as fid:
            pickle.dump(model_opts, fid, pickle.HIGHEST_PROTOCOL)

        config_path, _ = get_path(save_folder=model_folder_name,
                                  save_root_folder='data/models',
                                  file_name='configs.txt')
        self.log_configs(config_path, batch_size, epochs,
                         lr, model_opts)

        # Save training history
        history_path, saved_files_path = get_path(save_folder=model_folder_name,
                                                  save_root_folder='data/models',
                                                  file_name='history.pkl')
        with open(history_path, 'wb') as fid:
            pickle.dump(history.history, fid, pickle.HIGHEST_PROTOCOL)

        return saved_files_path

    # Test Functions
    def test(self, data_test, model_path=''):
        """
        Test function
        :param data_test: The raw data received from the dataset interface
        :param model_path: The path to the folder where the model and config files are saved.
        :return: The following performance metrics: acc, auc, f1, precision, recall
        """
        with open(os.path.join(model_path, 'model_opts.pkl'), 'rb') as fid:
            try:
                model_opts = pickle.load(fid)
            except:
                model_opts = pickle.load(fid, encoding='bytes')

        test_model = load_model(os.path.join(model_path, 'model.h5'))
        test_model.summary()

        test_data, _, _ = self.get_data({'test': data_test}, model_opts)
        test_results = test_model.predict(test_data['test'][0],
                                          batch_size=32, verbose=1)

        acc = accuracy_score(test_data['test'][1], np.round(test_results))
        f1 = f1_score(test_data['test'][1], np.round(test_results))
        auc = roc_auc_score(test_data['test'][1], np.round(test_results))
        roc = roc_curve(test_data['test'][1], test_results)
        precision = precision_score(test_data['test'][1], np.round(test_results))
        recall = recall_score(test_data['test'][1], np.round(test_results))
        pre_recall = precision_recall_curve(test_data['test'][1], test_results)

        print('acc:{} auc:{} f1:{} precision:{} recall:{}'.format(acc, auc, f1, precision, recall))

        save_results_path = os.path.join(model_path, '{:.2f}'.format(acc) + '.pkl')

        if not os.path.exists(save_results_path):
           results = {'results': test_results,
                       'data': test_data,
                       'acc': acc,
                       'auc': auc,
                       'f1': f1,
                       'roc': roc,
                       'precision': precision,
                       'recall': recall,
                       'pre_recall_curve': pre_recall}

        with open(save_results_path, 'wb') as fid:
            pickle.dump(test_results, fid, pickle.HIGHEST_PROTOCOL)
        return acc, auc, f1, precision, recall

    # Custom layers
    def stacked_rnn(self, data_types, data_sizes):
        """
        Generates SF-GRU model
        :param data_types: Data types, used for naming the layers of network.
        :param data_sizes: The sizes of each data type used for setting up Input layers
        :return: The model
        """
        network_inputs = []
        return_sequence = True
        num_layers = len(data_sizes)

        for i in range(num_layers):
            network_inputs.append(Input(shape=data_sizes[i], name='input_' + data_types[i]))

            if i == num_layers - 1:
                return_sequence = False

            if i == 0:
                x = self._gru(name='enc_' + data_types[i], r_sequence=return_sequence)(network_inputs[i])
            else:
                x = Concatenate(axis=2)([x, network_inputs[i]])

                x = self._gru(name='enc_' + data_types[i], r_sequence=return_sequence)(x)

        model_output = Dense(1, activation='sigmoid', name='output_dense')(x)

        net_model = Model(inputs=network_inputs, outputs=model_output)
        net_model.summary()

        return net_model

    def _gru(self, name='gru', r_state=False, r_sequence=False):
        """
        A helper function to creat a single GRU unit
        :param name: Name of the layer
        :param r_state: Whether to return the states of the GRU
        :param r_sequence: Whether to return sequence
        :return: A GRU unit
        """
        return GRU(units=self._num_hidden_units,
                   return_state=r_state,
                   return_sequences=r_sequence,
                   stateful=False,
                   kernel_regularizer=self._regularizer,
                   recurrent_regularizer=self._regularizer,
                   bias_regularizer=self._regularizer,
                   name=name)