data_helper.py

#! /usr/bin/env python
# -*- coding: utf-8 -*-
import os
os.environ['QT_API'] = 'pyqt5'

import scipy.ndimage as nd
import scipy.io as io
import scipy.stats as stats
import numpy as np
from utils import *

try:
    import trimesh
    from stl import mesh
except:
    print 'All dependencies not loaded, some functionality may not work'

class data_runner(object):
    ##############################################################################
    # Functions to initialize the data helper
    ##############################################################################
    def __init__(self, config_dict, for_training):
        self.for_training = for_training

        self.load_config_info(config_dict=config_dict)

        if for_training:
            self.load_dataset(config_dict)  # load the dataset for a certain kind

    def load_config_info(self, config_dict):
        self.config_dict = config_dict

        self.cube_len = self.config_dict['CUBE_LEN']
        self.bbox_size = self.config_dict['BOUNDING_BOX_SIZE']
        self.embedding_size = self.config_dict['EMBEDDING_VECOTR_SIZE']
        self.max_part_size = self.config_dict['MAX_PART_SIZE']

        self.input_pls = {}
        self.obj_num = 0

        if self.for_training:
            self.gpu_nums = self.config_dict['TRAIN']['NUM_GPUS']
            self.batch_size = self.config_dict['TRAIN']['BATCH_SIZE']

            self.shape_name = self.config_dict['TRAIN']['SHAPE_NAME']

            self.voxel_array = np.array([])  # used to store the voxel information in a list
            self.bbox_array = np.array([])  # used to store the baounding box information in a list
            self.part_obj_index_list = []  # the index of part_obj_index_list stands for the index of a certain object and the value is the index of a part
            self.obj_part_index_list = []  # the index of this list is the index of an object and the value stands for the index of a certain part

            self.global_steps = tf.placeholder(tf.float32, shape=(), name="global_steps")

            initial_vert_lr = config_dict['TRAIN']['VERT_LEARNING_RATE']
            initial_edge_lr = config_dict['TRAIN']['EDGE_LEARNING_RATE']
            initial_graph_gen_lr = config_dict['TRAIN']['GRAPH_GEN_LEARNING_RATE']

            vert_gamma = config_dict['TRAIN']['VERT_GAMMA']
            edge_gamma = config_dict['TRAIN']['EDGE_GAMMA']
            graph_gen_gamma = config_dict['TRAIN']['GRAPH_GEN_GAMMA']

            vert_lr_decay_step = config_dict['TRAIN']['VERT_LR_DECAY_STEP']
            edge_lr_decay_step = config_dict['TRAIN']['EDGE_LR_DECAY_STEP']
            graph_gen_lr_decay_step = config_dict['TRAIN']['GRAPH_GEN_LR_DECAY_STEP']

            self.vert_lr = tf.train.exponential_decay(initial_vert_lr,
                                                      global_step=self.global_steps,
                                                      decay_steps=vert_lr_decay_step,
                                                      decay_rate=vert_gamma,
                                                      staircase=True)
            self.edge_lr = tf.train.exponential_decay(initial_edge_lr,
                                                      global_step=self.global_steps,
                                                      decay_steps=edge_lr_decay_step,
                                                      decay_rate=edge_gamma,
                                                      staircase=True)
            self.graph_gen_lr = tf.train.exponential_decay(initial_graph_gen_lr,
                                                           global_step=self.global_steps,
                                                           decay_steps=graph_gen_lr_decay_step,
                                                           decay_rate=graph_gen_gamma,
                                                           staircase=True)

            initial_kl_loss_ratio = config_dict['TRAIN']['RECON_GEN_INITIAL_LOSS_RATIO']
            kl_loss_gamma = config_dict['TRAIN']['RECON_GEN_RATIO_GAMMA']
            kl_loss_decay_step = config_dict['TRAIN']['RECON_GEN_DECAY_STEP']
            self.kl_loss_ratio = tf.train.exponential_decay(initial_kl_loss_ratio,
                                                            global_step=self.global_steps,
                                                            decay_steps=kl_loss_decay_step,
                                                            decay_rate=kl_loss_gamma,
                                                            staircase=False)
        else:
            self.gpu_nums = 1
            self.batch_size = self.config_dict['TRAIN']['BATCH_SIZE']

            self.shape_name = config_dict['model_info']['model_class']

            self.visible_part_indexes_array = config_dict['model_info']['part_masks']
            self.obj_num = len(config_dict['model_info']['part_masks'])

            self.bbox_means_array = config_dict['model_info']['bbox_info']['bbox_mean']
            self.bbox_stds_array = config_dict['model_info']['bbox_info']['bbox_variance']
            self.bbox_benchmark_array = config_dict['model_info']['bbox_info']['bbox_benchmark']


    ##############################################################################
    # Functions to load the dataset
    ##############################################################################
    def get_voxels_from_file(self, path, cube_len=32):
        voxels = io.loadmat(path)['voxels3D']
        if cube_len != 32 and cube_len == 64:
            voxels = nd.zoom(voxels, (2, 2, 2), mode='constant', order=0)
        return voxels

    def get_bboxs_from_file(self, obj_path):
        found_bbox = False

        bbox_file = open(obj_path, 'r')
        all_lines = bbox_file.readlines()
        for c_line in all_lines:
            if c_line.find('Min_Max') != -1:
                found_bbox = True
                continue

            if found_bbox:
                cur_line = c_line.strip()
                bbox_list = map(float, cur_line.split())
                bbox_list = np.array(bbox_list)
                return bbox_list

    def get_part_label(self, file_path):
        if file_path.find('.mat') == len(file_path) - 4:
            label_num = int(file_path[len(file_path) - 5 : len(file_path) - 4])
            return label_num

    def get_part_info(self, obj_path, cube_len=32):
        voxel_file_list = [f for f in os.listdir(obj_path) if f.endswith('.mat')]
        bbox_file_list = [f for f in os.listdir(obj_path) if f.endswith('.txt')]

        voxel_file_list.sort()
        bbox_file_list.sort()

        voxel_default_value = self.config_dict['TRAIN']['VOXEL_DEFAULT_VALUE']

        if len(voxel_file_list) != len(bbox_file_list):
            error_str = "the file number for voxel maps should correspond to the number of bounding boxes."
            raise KeyError(error_str)

        visible_part_index_list = np.asarray([self.get_part_label(f_name) for f_name in voxel_file_list], dtype=np.int).tolist()
        voxels_list = np.asarray([self.get_voxels_from_file(obj_path + f_name, cube_len) for f_name in voxel_file_list], dtype=np.float32)
        bboxs_list = np.array([self.get_bboxs_from_file(os.path.join(obj_path, f_name)) for f_name in bbox_file_list])

        final_part_mask_list = [[]] * self.max_part_size
        final_part_voxel_list = [[]] * self.max_part_size
        final_part_bbox_list = [[]] * self.max_part_size

        for index, label in enumerate(visible_part_index_list):
            final_part_voxel_list[label] = voxels_list[index]
            final_part_bbox_list[label] = bboxs_list[index]
            final_part_mask_list[label] = 1

        for part_index, mask in enumerate(final_part_mask_list):
            if not mask:
                # final_part_voxel_list[part_index] = np.zeros_like(voxels_list[0])
                final_part_voxel_list[part_index] = np.full(voxels_list[0].shape, float(voxel_default_value))
                final_part_bbox_list[part_index] = np.zeros_like(bboxs_list[0])
                final_part_mask_list[part_index] = 0

        voxels_array = np.asarray(final_part_voxel_list, dtype=np.float32)
        bboxs_array = np.asarray(final_part_bbox_list)
        part_mask_array = np.asarray(final_part_mask_list)

        return voxels_array, bboxs_array, part_mask_array

    def load_dataset(self, config_dict):
        data_path = config_dict['TRAIN']['DIR_PATH']
        data_path = os.path.join(data_path, self.shape_name)

        # fn_list: a list for file names
        fn_list = [f_path for f_path in os.listdir(data_path) if os.path.isdir(os.path.join(data_path, f_path))]
        # fp_list: a list for file paths
        fp_list = [os.path.join(data_path, f_path) for f_path in fn_list if os.path.isdir(os.path.join(data_path, f_path))]

        obj_num = len(fp_list)  # Total number of 3D shapes
        all_voxels_list = []    # a list for all the voxel maps
        all_bboxs_list = []     # a list for all the bounding boxes
        all_visible_part_indexes_list = []    # a list for the visible part indexes. Two-dimensional list: (shape_index, [visible_part_indexes])

        """part_obj_index_list is used to store the index of certain object the index of relation_list indicates a part
        in all_voxels_list or all_bboxs_list and the value is the index of certain object"""
        part_index = 0
        obj_part_index_list = []    # two-dimensional list
        part_obj_index_list = []    # one-dimensional list

        for obj_index, f_path in enumerate(fp_list):
            cur_voxels, cur_bboxs, cur_visible_part_index_list = self.get_part_info(f_path + "/")
            if cur_voxels.shape[0] != cur_bboxs.shape[0]:
                raise KeyError("Wrong sizes between the current voxel map arrays and bounding box arrays!!!")

            for c_voxel in cur_voxels:
                all_voxels_list.append(c_voxel)
                part_obj_index_list.append(obj_index)

            cur_part_index_list = []
            for c_bbox in cur_bboxs:
                all_bboxs_list.append(c_bbox)
                cur_part_index_list.append(part_index)

                part_index = part_index + 1
            obj_part_index_list.append(cur_part_index_list)

            all_visible_part_indexes_list.append(cur_visible_part_index_list)

        obj_part_index_array = np.array(obj_part_index_list)
        file_name_array = np.array(fn_list)

        all_voxels_array = np.array(all_voxels_list, dtype=np.float32)
        all_bboxs_array = np.array(all_bboxs_list)
        visible_part_index_array = np.array(all_visible_part_indexes_list)

        all_bboxs_array = self.preprocess_bbox_data(all_bboxs_array, visible_part_indexes=visible_part_index_array)

        self.obj_num, self.voxel_array, self.bbox_array, self.visible_part_indexes_array,\
        self.part_obj_index_list, self.obj_part_index_array, self.fn_list = \
            (obj_num, all_voxels_array, all_bboxs_array, visible_part_index_array, part_obj_index_list, obj_part_index_array, file_name_array)

        return obj_num, all_voxels_array, all_bboxs_array, all_visible_part_indexes_list, part_obj_index_list, obj_part_index_array


    ##############################################################################
    # Functions to preprocess the original input data
    ##############################################################################
    def preprocess_voxel_data(self, voxel_input):
        voxel_data = np.reshape(voxel_input, [-1, self.max_part_size, self.cube_len*self.cube_len*self.cube_len])
        voxel_data = voxel_data.astype(float)

        final_voxel_mean_list = []
        final_voxel_std_list = []
        final_voxel_list = []

        for part_index in range(self.max_part_size):
            cur_part_data = voxel_data[:, part_index, :]
            cur_part_data = np.reshape(cur_part_data, [-1])

            cur_mean = np.mean(cur_part_data)
            final_voxel_mean_list.append(cur_mean)

            cur_part_data -= cur_mean

            cur_std = np.std(cur_part_data)
            final_voxel_std_list.append(cur_std)

            cur_part_data /= cur_std

            cur_part_data = np.reshape(cur_part_data, [-1, self.cube_len*self.cube_len*self.cube_len])
            final_voxel_list.append(cur_part_data)

        self.voxel_means_array = np.array(final_voxel_mean_list)
        self.voxel_stds_array = np.array(final_voxel_std_list)

        if len(self.voxel_means_array) != len(self.voxel_stds_array):
            raise KeyError("The voxel means array and stds array don't have the same shape")

        final_data_array = np.stack(final_voxel_list, axis=1)
        final_data_array = np.reshape(final_data_array, [-1, self.cube_len, self.cube_len, self.cube_len])
        return final_data_array

    def preprocess_bbox_data(self, bbox_input, visible_part_indexes=None):
        """There are two steps in preprocessing bbox information. One is to transform the bounding boxes into offset forms. And the second step
            is to scale the processed bounding boxes with extracted means and variances."""
        bbox_data = np.reshape(bbox_input, [-1, self.max_part_size, self.bbox_size])

        # calculate average values for bounding boxes
        average_bboxs = self.average_bboxs(bbox_data, remove_zero=True)

        # complete the missing data
        bbox_data = self.compute_missing_data(bbox_data, visible_part_indexes)

        transformed_bbox_list = []
        # transform the bounding boxes into the offset form
        for part_index in range(self.max_part_size):
            cur_part_data = bbox_data[:, part_index, :]

            target_dx = (cur_part_data[:, 0] - average_bboxs[part_index][0]) / average_bboxs[part_index][3]
            target_dy = (cur_part_data[:, 1] - average_bboxs[part_index][1]) / average_bboxs[part_index][4]
            target_dz = (cur_part_data[:, 2] - average_bboxs[part_index][2]) / average_bboxs[part_index][5]
            target_dl = np.log(cur_part_data[:, 3] / average_bboxs[part_index][3])
            target_dw = np.log(cur_part_data[:, 4] / average_bboxs[part_index][4])
            target_dh = np.log(cur_part_data[:, 5] / average_bboxs[part_index][5])

            normalized_part_array = np.stack([target_dx, target_dy, target_dz, target_dl, target_dw, target_dh], axis=1)
            transformed_bbox_list.append(normalized_part_array)

        bboxs_array = np.stack(transformed_bbox_list, axis=1)
        bboxs_array = np.reshape(bboxs_array, [-1, self.bbox_size])

        # scale the bounding boxes with computed means and variances
        final_data_array = self.scale_bbox_data(bboxs_array)

        return final_data_array

    def compute_missing_data(self, bbox_input, visible_part_indexes=None):
        bbox_data = np.reshape(bbox_input, [-1, self.max_part_size, self.bbox_size])

        for obj_ind, part_masks in enumerate(visible_part_indexes):
            for p_ind, p_mask in enumerate(part_masks):
                if not p_mask:
                    bbox_data[obj_ind][p_ind] = self.bbox_benchmark_array[p_ind]

        return bbox_data

    def average_bboxs(self, bbox_input, remove_zero=True):
        """The benchmark array indicates the mean value for each part"""
        bboxs_mean_list = []
        for part_index in range(self.max_part_size):
            cur_part_data = bbox_input[:, part_index, :]

            if remove_zero:
                cur_mean_list = np.ma.masked_equal(cur_part_data, 0.0).mean(axis=0)
            else:
                cur_mean_list = np.mean(cur_part_data, axis=0)
            bboxs_mean_list.append(cur_mean_list)

        # The benchmark array indicates the mean value for each part
        self.bbox_benchmark_array = np.array(bboxs_mean_list)

        return self.bbox_benchmark_array

    def bboxs_mode(self, bbox_input):
        bbox_data = np.reshape(bbox_input, [-1, self.max_part_size, self.bbox_size])

        bboxs_mode_list = []
        for part_index in range(self.max_part_size):
            cur_part_data = bbox_data[:, part_index, :]

            cur_mode_list = stats.mode(cur_part_data, axis=0)
            bboxs_mode_list.append(cur_mode_list)

        self.bbox_benchmark_array = np.array(bboxs_mode_list)

        return self.bbox_benchmark_array

    def bboxs_median(self, bbox_input):
        bbox_data = np.reshape(bbox_input, [-1, self.max_part_size, self.bbox_size])

        bboxs_median_list = []
        for part_index in range(self.max_part_size):
            cur_part_data = bbox_data[:, part_index, :]

            mask_part_data = np.ma.masked_where(cur_part_data == 0, cur_part_data)
            cur_median_list = np.ma.median(mask_part_data, axis=0).filled(0)

            # cur_median_list = np.median(cur_part_data, axis=0)
            bboxs_median_list.append(cur_median_list)

        self.bbox_benchmark_array = np.array(bboxs_median_list)

        return self.bbox_benchmark_array

    def scale_bbox_data(self, bbox_input):
        bbox_data = np.reshape(bbox_input, [-1, self.max_part_size, self.bbox_size])

        final_bbox_mean_list = []
        final_bbox_std_list = []
        final_data_list = []

        for part_index in range(self.max_part_size):
            cur_part_data = bbox_data[:, part_index, :]

            cur_mean_list = np.mean(cur_part_data, axis=0)
            final_bbox_mean_list.append(cur_mean_list)

            cur_part_data -= cur_mean_list

            cur_std_list = np.std(cur_part_data, axis=0)
            final_bbox_std_list.append(cur_std_list)

            cur_std_array = np.array(cur_std_list)
            cur_part_data /= cur_std_array

            final_data_list.append(cur_part_data)

        self.bbox_means_array = np.array(final_bbox_mean_list)
        self.bbox_stds_array = np.array(final_bbox_std_list)

        if self.bbox_means_array.shape[0] != self.bbox_stds_array.shape[0] or\
                self.bbox_means_array.shape[1] != self.bbox_stds_array.shape[1]:
            raise KeyError("The bbox means array and bbox stds array don't have the same shape")

        final_data_array = np.stack(final_data_list, axis=1)
        final_data_array = np.reshape(final_data_array, [-1, self.bbox_size])

        return final_data_array


    ##############################################################################
    # Functions to prepare the data for next batch
    ##############################################################################
    def get_placeholders(self):
        if self.for_training:
            input_pls = self.get_placeholders_for_training()
        else:
            input_pls = self.get_placeholders_for_testing()

        return input_pls

    def get_placeholders_for_testing(self):
        embedding_size = self.config_dict['EMBEDDING_VECOTR_SIZE']

        input_pls = {
            'part_visible_masks': tf.placeholder(dtype=tf.int32, shape=[None, self.batch_size, self.max_part_size], name='part_visible_masks'),  # the mask array for visible parts
            'latent_codes': tf.placeholder(dtype=tf.float32, shape=[self.batch_size, embedding_size], name='latent_codes'),  # the mask array for visible parts
        }
        return input_pls

    def get_placeholders_for_training(self):
        cube_len = self.cube_len
        bbox_size = self.bbox_size

        input_pls = {
            'part_voxels': tf.placeholder(dtype=tf.float32, shape=[None, None, cube_len, cube_len, cube_len, 1], name='cur_part_voxels'),
            'part_bbox': tf.placeholder(shape=[None, None, bbox_size], dtype=tf.float32, name='part_bbox'),
            'gaussian_noise': tf.placeholder(shape=[None, None, self.embedding_size], dtype=tf.float32, name='gaussian_noise'),
            'labels': tf.placeholder(dtype=tf.int32, shape=[None], name='labels'),  # label of each part per batch
            'rel_pair_mask_inds': tf.placeholder(dtype=tf.int32, shape=[None, None, 2], name='rel_pair_mask_inds'), # the indexes of two nodes for a bounding box pair
            'part_visible_masks': tf.placeholder(dtype=tf.int32, shape=[None, self.batch_size, self.max_part_size], name='part_visible_masks'),  # the mask array for visible parts
            # the weight vector for a part's loss
            'part_voxel_loss_weights': tf.placeholder(dtype=tf.float32, shape=[None, self.max_part_size], name='part_voxel_loss_weights'),
            'part_bbox_loss_masks': tf.placeholder(dtype=tf.float32, shape=[None, None, self.max_part_size * (self.max_part_size - 1) / 2], name='part_bbox_loss_masks'),  # the loss mask for a part's loss
            'vert_lr': tf.placeholder(dtype=tf.float32, shape=(), name='vert_learning_rate'),
            'edge_lr': tf.placeholder(dtype=tf.float32, shape=(), name='edge_learning_rate'),
            'graph_gen_lr': tf.placeholder(dtype=tf.float32, shape=(), name='graph_gen_learning_rate'),
            'recon_gen_loss_ratio': tf.placeholder(dtype=tf.float32, shape=(), name='recon_generation_loss_ratio'),
            'voxel_bbox_ratio': tf.placeholder(dtype=tf.float32, shape=(), name='voxel_bbox_ratio'),
            'g_rec_kl_loss_ratio': tf.placeholder(dtype=tf.float32, shape=(), name='graph_recon_kl_loss_ratio'),
            'max_gradient_norm': tf.placeholder(dtype=tf.float32, shape=(), name='max_gradient_norm')
        }

        self.input_pls = input_pls

        return dict(input_pls)

    def get_inputs_for_testing(self, cur_iter):
        cur_selected_obj_idx = np.random.randint(self.obj_num, size=[self.gpu_nums, self.batch_size])

        current_visible_part_index = self.visible_part_indexes_array[cur_selected_obj_idx]
        gaussian_noise = np.random.normal(0, 1, [self.gpu_nums, self.batch_size, self.embedding_size])
        # gaussian_noise = np.full((self.gpu_nums, self.batch_size, self.embedding_size), 1) * (cur_iter + 1)

        print "Gaussian noise:\n"
        print gaussian_noise

        input_pls = {}
        input_pls['gaussian_noise'] = gaussian_noise
        input_pls['visible_part_index'] = current_visible_part_index
        input_pls['selected_model_idx'] = cur_selected_obj_idx

        return input_pls

    def get_next_minibatch(self, object_index=None, cur_iter=-1):
        """We can set the object_index to be not None to retrieve certain objects"""
        if object_index is not None:
            cur_selected_obj_idx = object_index
        else:
            cur_selected_obj_idx = np.random.randint(self.obj_num, size=[self.gpu_nums, self.batch_size])

        selected_part_index = self.obj_part_index_array[cur_selected_obj_idx]   # two dimensional array
        selected_visible_part_index = self.visible_part_indexes_array[cur_selected_obj_idx]

        final_selected_part_index_list = np.reshape(selected_part_index, [len(selected_part_index), -1])

        # get the voxel map/bounding boxes/part label for current mini batch
        selected_part_voxels_list, selected_part_bboxs_list = self.prepare_train_batch(final_selected_part_index_list)

        # Compute the edge indexes for each bounding box pair
        total_edge_mask_inds = []
        for gpu_iter in range(self.gpu_nums):
            current_begin_index = 0
            cur_edge_pair_mask_inds = []
            for s_index in range(self.batch_size):
                for f_offset in range(self.max_part_size):
                    for s_offset in range(f_offset + 1, self.max_part_size):
                        tmp_list = [f_offset + current_begin_index, s_offset + current_begin_index]
                        cur_edge_pair_mask_inds.append(tmp_list)
                current_begin_index = current_begin_index + self.max_part_size
            total_edge_mask_inds.append(cur_edge_pair_mask_inds)

        part_voxel_arr = np.array(selected_part_voxels_list)
        part_voxel_arr = part_voxel_arr[:, :, :, :, :, np.newaxis]
        part_bbox_arr = np.array(selected_part_bboxs_list)

        # Compute the loss mask for each bounding box pair
        cur_bbox_loss_masks = []
        for g_part_visible_masks in selected_visible_part_index:    # per GPU
            cur_g_loss_masks = []
            for b_part_visible_masks in g_part_visible_masks:   # per batch
                cur_batch_loss_masks = []
                for f_p_index in range(len(b_part_visible_masks)):
                    for s_p_index in range(f_p_index + 1, len(b_part_visible_masks)):
                        # Only compute the bounding box pair losses when at least one part is visible
                        if b_part_visible_masks[f_p_index] + b_part_visible_masks[s_p_index] > 0:
                            cur_batch_loss_masks.append(1)
                        else:
                            cur_batch_loss_masks.append(0)
                cur_g_loss_masks.append(cur_batch_loss_masks)
            cur_bbox_loss_masks.append(cur_g_loss_masks)

        bbox_loss_mask_array = np.array(cur_bbox_loss_masks)

        voxel_loss_weights = None
        shape_name = self.config_dict['TRAIN']['SHAPE_NAME']
        if shape_name == 'motorbike':
            voxel_loss_weights = np.array([[1, 1, 1, 1, 1], ] * self.batch_size)  # motorbike
        elif shape_name == 'chair':
            voxel_loss_weights = np.array([[40, 5, 5, 1, 4], ] * self.batch_size)  # chair
        elif shape_name == 'airplane':
            voxel_loss_weights = np.array([[1, 1, 1, 1, 1, 1], ] * self.batch_size)  # airplane
        elif shape_name == 'guitar':
            voxel_loss_weights = np.array([[1, 1, 1], ] * self.batch_size)  # guitar
        elif shape_name == 'lamp':
            voxel_loss_weights = np.array([[5, 20, 1, 1], ] * self.batch_size)  # lamp
        elif shape_name == 'toy_examples':
            voxel_loss_weights = np.array([[1, 20], ] * self.batch_size)  # toy_examples

        data_pls = self.get_batch_config_dict(config_dict=self.config_dict, cur_iter=cur_iter)
        data_pls['part_voxels'] = part_voxel_arr
        data_pls['part_bbox'] = part_bbox_arr
        data_pls['rel_pair_mask_inds'] = total_edge_mask_inds
        data_pls['selected_model_idx'] = cur_selected_obj_idx
        data_pls['part_visible_masks'] = selected_visible_part_index
        data_pls['part_voxel_loss_weights'] = voxel_loss_weights
        data_pls['part_bbox_loss_masks'] = bbox_loss_mask_array

        return dict(data_pls)

    def get_batch_config_dict(self, config_dict, cur_iter=-1):
        # Compute the current learning rate for different components in the network
        f_dict = {}
        f_dict['vert_lr'] = self.vert_lr.eval(feed_dict={self.global_steps: cur_iter})
        f_dict['edge_lr'] = self.edge_lr.eval(feed_dict={self.global_steps: cur_iter})
        f_dict['graph_gen_lr'] = self.graph_gen_lr.eval(feed_dict={self.global_steps: cur_iter})

        kl_anneal_iter = config_dict['TRAIN']['KL_ANNEAL_ITER']
        if cur_iter > kl_anneal_iter:
            final_kl_ratio = config_dict['TRAIN']['FINAL_KL_RATIO']

            kl_loss_ratio = self.kl_loss_ratio.eval(feed_dict={self.global_steps: (cur_iter - kl_anneal_iter)})
            kl_loss_rate = max(kl_loss_ratio, final_kl_ratio)
        else:
            kl_loss_rate = 1.0

        print "vert_lr: %6f, edge_lr: %6f, graph_gen_lr: %6f, kl_loss_ratio: %6f" % \
              (f_dict['vert_lr'], f_dict['edge_lr'], f_dict['graph_gen_lr'], kl_loss_rate)

        f_dict['recon_gen_loss_ratio'] = kl_loss_rate

        f_dict['voxel_bbox_ratio'] = config_dict['TRAIN']['VOXEL_BBOX_LOSS_RATIO']

        f_dict['g_rec_kl_loss_ratio'] = config_dict['TRAIN']['GRAPH_REC_KL_LOSS_RATIO']

        embedding_size = config_dict['EMBEDDING_VECOTR_SIZE']
        f_dict['gaussian_noise'] = np.random.normal(0, 1, [self.gpu_nums, self.batch_size, embedding_size])

        f_dict['max_gradient_norm'] = float(config_dict['TRAIN']['MAX_GRADIENT_NORM'])

        return f_dict

    def prepare_train_batch(self, selected_part_index_list):
        total_selected_voxels_list = []
        total_selected_bboxs_list = []

        for s_part_indexs in selected_part_index_list:
            selected_part_voxels_list = []
            selected_part_bboxs_list = []

            for part_index in s_part_indexs:
                if part_index != -1:
                    selected_part_voxels_list.append(self.voxel_array[part_index])
                    selected_part_bboxs_list.append(self.bbox_array[part_index])
                else:
                    zero_voxel_array = np.zeros([self.cube_len, self.cube_len, self.cube_len], dtype=np.float32)
                    zero_bbox_array = np.zeros(self.bbox_size)

                    selected_part_voxels_list.append(zero_voxel_array)
                    selected_part_bboxs_list.append(zero_bbox_array)

            total_selected_voxels_list.append(selected_part_voxels_list)
            total_selected_bboxs_list.append(selected_part_bboxs_list)

        return total_selected_voxels_list, total_selected_bboxs_list


    ##############################################################################
    # Other functions for data processing
    ##############################################################################
    def clip_voxels_data(self, voxel_data):
        """Clip the voxel value to zero or one by evaluating the threshold 0.5"""
        result_voxels = np.copy(voxel_data)

        for x_ind in range(len(voxel_data)):
            for y_ind in range(len(voxel_data[0])):
                for z_ind in range(len(voxel_data[0][0])):
                    cur_value = voxel_data[x_ind][y_ind][z_ind]
                    if cur_value >= 0 and cur_value <= 0.5:
                        result_voxels[x_ind][y_ind][z_ind] = 0
                    elif cur_value > 0.5 and cur_value <= 1:
                        result_voxels[x_ind][y_ind][z_ind] = 1

        return result_voxels

    def process_voxel_data(self, voxel_data, part_visible_masks):
        """This method is used to turn the voxel value to zero or one"""
        voxel_data = np.reshape(voxel_data, [-1, self.cube_len, self.cube_len, self.cube_len])

        voxels_list = []
        start_index = 0

        for part_masks in part_visible_masks:
            cur_object_voxel_list = []
            for p_ind, p_mask in enumerate(part_masks):  # per part
                if int(p_mask) == 1:
                    voxel_index = start_index + p_ind
                    cur_voxel_map = voxel_data[voxel_index]
                    result_voxels = self.clip_voxels_data(cur_voxel_map)
                    cur_object_voxel_list.append(result_voxels)
            voxels_list.append(cur_object_voxel_list)
            start_index = start_index + self.max_part_size

        return voxels_list

    def process_bbox_data(self, bboxs_data, part_visible_masks):
        """"Split a 12D feature vector into two 6D vectors and then compute the final bounding box output by averaging
        the (k-1) corresponding bounding boxes for a part."""
        max_ctx_size = self.max_part_size * (self.max_part_size - 1) / 2

        bboxs_list = []
        start_index = 0
        for part_masks in part_visible_masks:
            object_ctx_arr = bboxs_data[start_index:start_index + max_ctx_size, :]
            object_ctx_arr = np.split(object_ctx_arr, 2, axis=1)

            start_index = start_index + max_ctx_size

            # indicates the first bounding box in a box pair
            f_bbox_arr = np.zeros([self.max_part_size, self.bbox_size], dtype=np.float32)
            # indicates the second bounding box in a box pair
            s_bbox_arr = np.zeros([self.max_part_size, self.bbox_size], dtype=np.float32)

            cur_part_index = 0
            for f_offset in range(self.max_part_size):
                for s_offset in range(f_offset + 1, self.max_part_size):
                    cur_first_bbox_value = object_ctx_arr[0][cur_part_index]
                    cur_second_bbox_value = object_ctx_arr[1][cur_part_index]

                    f_bbox_arr[f_offset] = f_bbox_arr[f_offset] + cur_first_bbox_value
                    s_bbox_arr[s_offset] = s_bbox_arr[s_offset] + cur_second_bbox_value
                    cur_part_index = cur_part_index + 1

            if self.max_part_size > 1:
                total_bbox_arr = (f_bbox_arr + s_bbox_arr) / (self.max_part_size - 1)
            else:
                total_bbox_arr = f_bbox_arr + s_bbox_arr

            cur_object_ctx_list = []  # store all the part bbox information for one certain object
            for p_ind, p_mask in enumerate(part_masks):
                if int(p_mask) == 1:
                    cur_object_ctx_list.append(total_bbox_arr[p_ind])
            bboxs_list.append(cur_object_ctx_list)

        return bboxs_list


    ##############################################################################
    # Functions to output results to files
    ##############################################################################
    def write_output_to_file(self, voxels_list=None, bboxs_list=None, part_visible_masks=None, input_info_dict=None,
                             output_dir=None, iter_n=None):
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        selected_model_idx = input_info_dict['selected_model_idx']

        if len(voxels_list[0]) != len(bboxs_list[0]):
            raise KeyError("The voxels list and the bboxs list should share the same size when outputting")
        if len(voxels_list) != len(selected_model_idx[self.gpu_nums-1]):
            raise KeyError("The voxels list and the selected-model-idx list should share the same size when outputting")

        for obj_index, (iter_model_voxels, iter_model_bboxs, cur_model_idx, p_masks) in \
                enumerate(zip(voxels_list, bboxs_list, selected_model_idx[self.gpu_nums-1], part_visible_masks)):
            if len(iter_model_voxels) != len(iter_model_bboxs):
                raise KeyError("The iter_model_voxels and the iter_model_bboxs don't have the same size when outputting")

            if self.for_training:
                cur_model_name = self.fn_list[cur_model_idx]

                dir_name = "%s_iter_%d" % (cur_model_name, iter_n)
                output_model_dir = os.path.join(output_dir, dir_name)
                if not os.path.exists(output_model_dir):
                    os.makedirs(output_model_dir)

                file_name = dir_name
            else:
                cur_output_index = iter_n * self.batch_size + obj_index

                output_model_dir = os.path.join(output_dir, str(cur_output_index))
                if not os.path.exists(output_model_dir):
                    os.makedirs(output_model_dir)

                file_name = str(cur_output_index)

            start_index = 0
            for m_idx, p_mask in enumerate(p_masks):
                if p_mask:
                    self.write_voxels_to_file(iter_model_voxels[start_index], file_name, m_idx, output_model_dir)

                    cur_model_bbox = iter_model_bboxs[start_index] * self.bbox_stds_array[m_idx] + self.bbox_means_array[m_idx]
                    cur_model_bbox[0] = self.bbox_benchmark_array[m_idx][3] * cur_model_bbox[0] + self.bbox_benchmark_array[m_idx][0]
                    cur_model_bbox[1] = self.bbox_benchmark_array[m_idx][4] * cur_model_bbox[1] + self.bbox_benchmark_array[m_idx][1]
                    cur_model_bbox[2] = self.bbox_benchmark_array[m_idx][5] * cur_model_bbox[2] + self.bbox_benchmark_array[m_idx][2]
                    cur_model_bbox[3] = np.exp(cur_model_bbox[3]) * self.bbox_benchmark_array[m_idx][3]
                    cur_model_bbox[4] = np.exp(cur_model_bbox[4]) * self.bbox_benchmark_array[m_idx][4]
                    cur_model_bbox[5] = np.exp(cur_model_bbox[5]) * self.bbox_benchmark_array[m_idx][5]
                    self.write_bboxs_to_file(cur_model_bbox, file_name, m_idx, output_model_dir)
                    start_index = start_index + 1

                    print file_name + "_" + str(m_idx)

    def write_voxels_to_file(self, voxels_info, file_name, seq_num, output_dir):
        """The parameter dual_iter_index stands for the dual iteration index while the iter_index is the
           index of certain epoch and the seq_num stands for the index of a part inside certain object"""

        output_file_name = "%s_%d.mat" % (file_name, seq_num)
        output_file_path = os.path.join(output_dir, output_file_name)

        io.savemat(output_file_path, {'voxels3D': voxels_info})

    def write_bboxs_to_file(self, bboxs_info, file_name, seq_num, output_dir):
        """the parameter dual_iter_index stands for the dual iteration index while the iter_index is the
        index of certain epoch and the seq_num stands for the index of a part inside certain object"""

        output_file_name = "%s_%d_transform_info.txt" % (file_name, seq_num)
        output_file_path = os.path.join(output_dir, output_file_name)

        with open(output_file_path, "w") as out_f:
            out_f.write("6D vector for representation\n")
            for value in bboxs_info:
                output_content = "%f " % (value)
                out_f.write(output_content)
            out_f.write("\n")


    ##############################################################################
    # Functions to output model config info
    ##############################################################################
    def write_model_info_to_file(self, output_dir=None):
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        output_file_name = "%s_model_info.txt" % self.shape_name
        output_file_path = os.path.join(output_dir, output_file_name)

        with open(output_file_path, "w") as out_f:
            out_f.write("model_class\n")
            out_f.write("%s\n" % self.shape_name)

            out_f.write("\nbbox_mean\n")
            for p_means in self.bbox_means_array:
                for value in p_means:
                    output_content = "%f " % (value)
                    out_f.write(output_content)
                out_f.write("\n")

            out_f.write("\nbbox_variance\n")
            for p_stds in self.bbox_stds_array:
                for value in p_stds:
                    output_content = "%f " % (value)
                    out_f.write(output_content)
                out_f.write("\n")

            out_f.write("\nbbox_benchmark\n")
            for p_benchmark in self.bbox_benchmark_array:
                for value in p_benchmark:
                    output_content = "%f " % (value)
                    out_f.write(output_content)
                out_f.write("\n")

        mask_file_name = "%s_mask.mat" % self.shape_name
        mask_file_path = os.path.join(output_dir, mask_file_name)

        io.savemat(mask_file_path, {'masks': self.visible_part_indexes_array})