data_process.py

'''
Parallel data loading functions
'''
import sys
import time
import theano
import numpy as np
import traceback
from PIL import Image
from six.moves import queue
from multiprocessing import Process, Event

from lib.config import cfg
from lib.data_augmentation import preprocess_img
from lib.data_io import get_voxel_file, get_rendering_file
from lib.binvox_rw import read_as_3d_array


def print_error(func):
    '''Flush out error messages. Mainly used for debugging separate processes'''

    def func_wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except:
            traceback.print_exception(*sys.exc_info())
            sys.stdout.flush()

    return func_wrapper


class DataProcess(Process):

    def __init__(self, data_queue, data_paths, repeat=True):
        '''
        data_queue : Multiprocessing queue
        data_paths : list of data and label pair used to load data
        repeat : if set True, return data until exit is set
        '''
        super(DataProcess, self).__init__()
        # Queue to transfer the loaded mini batches
        self.data_queue = data_queue
        self.data_paths = data_paths
        self.num_data = len(data_paths)
        self.repeat = repeat

        # Tuple of data shape
        self.batch_size = cfg.CONST.BATCH_SIZE
        self.exit = Event()
        self.shuffle_db_inds()

    def shuffle_db_inds(self):
        # Randomly permute the training roidb
        if self.repeat:
            self.perm = np.random.permutation(np.arange(self.num_data))
        else:
            self.perm = np.arange(self.num_data)
        self.cur = 0

    def get_next_minibatch(self):
        if (self.cur + self.batch_size) >= self.num_data and self.repeat:
            self.shuffle_db_inds()

        db_inds = self.perm[self.cur:min(self.cur + self.batch_size, self.num_data)]
        self.cur += self.batch_size
        return db_inds

    def shutdown(self):
        self.exit.set()

    @print_error
    def run(self):
        iteration = 0
        # Run the loop until exit flag is set
        while not self.exit.is_set() and self.cur <= self.num_data:
            # Ensure that the network sees (almost) all data per epoch
            db_inds = self.get_next_minibatch()

            data_list = []
            label_list = []
            for batch_id, db_ind in enumerate(db_inds):
                datum = self.load_datum(self.data_paths[db_ind])
                label = self.load_label(self.data_paths[db_ind])

                data_list.append(datum)
                label_list.append(label)

            batch_data = np.array(data_list).astype(np.float32)
            batch_label = np.array(label_list).astype(np.float32)

            # The following will wait until the queue frees
            self.data_queue.put((batch_data, batch_label), block=True)
            iteration += 1

    def load_datum(self, path):
        pass

    def load_label(self, path):
        pass


class ReconstructionDataProcess(DataProcess):

    def __init__(self, data_queue, category_model_pair, background_imgs=[], repeat=True,
                 train=True):
        self.repeat = repeat
        self.train = train
        self.background_imgs = background_imgs
        super(ReconstructionDataProcess, self).__init__(
            data_queue, category_model_pair, repeat=repeat)

    @print_error
    def run(self):
        # set up constants
        img_h = cfg.CONST.IMG_W
        img_w = cfg.CONST.IMG_H
        n_vox = cfg.CONST.N_VOX

        # This is the maximum number of views
        n_views = cfg.CONST.N_VIEWS

        while not self.exit.is_set() and self.cur <= self.num_data:
            # To insure that the network sees (almost) all images per epoch
            db_inds = self.get_next_minibatch()

            # We will sample # views
            if cfg.TRAIN.RANDOM_NUM_VIEWS:
                curr_n_views = np.random.randint(n_views) + 1
            else:
                curr_n_views = n_views

            # This will be fed into the queue. create new batch everytime
            batch_img = np.zeros(
                (curr_n_views, self.batch_size, 3, img_h, img_w), dtype=theano.config.floatX)
            batch_voxel = np.zeros(
                (self.batch_size, n_vox, 2, n_vox, n_vox), dtype=theano.config.floatX)

            # load each data instance
            for batch_id, db_ind in enumerate(db_inds):
                category, model_id = self.data_paths[db_ind]
                image_ids = np.random.choice(cfg.TRAIN.NUM_RENDERING, curr_n_views)

                # load multi view images
                for view_id, image_id in enumerate(image_ids):
                    im = self.load_img(category, model_id, image_id)
                    # channel, height, width
                    batch_img[view_id, batch_id, :, :, :] = \
                        im.transpose((2, 0, 1)).astype(theano.config.floatX)

                voxel = self.load_label(category, model_id)
                voxel_data = voxel.data

                batch_voxel[batch_id, :, 0, :, :] = voxel_data < 1
                batch_voxel[batch_id, :, 1, :, :] = voxel_data

            # The following will wait until the queue frees
            self.data_queue.put((batch_img, batch_voxel), block=True)

        print('Exiting')

    def load_img(self, category, model_id, image_id):
        image_fn = get_rendering_file(category, model_id, image_id)
        im = Image.open(image_fn)

        t_im = preprocess_img(im, self.train)
        return t_im

    def load_label(self, category, model_id):
        voxel_fn = get_voxel_file(category, model_id)
        with open(voxel_fn, 'rb') as f:
            voxel = read_as_3d_array(f)

        return voxel


def kill_processes(queue, processes):
    print('Signal processes')
    for p in processes:
        p.shutdown()

    print('Empty queue')
    while not queue.empty():
        time.sleep(0.5)
        queue.get(False)

    print('kill processes')
    for p in processes:
        p.terminate()


def make_data_processes(queue, data_paths, num_workers, repeat=True, train=True):
    '''
    Make a set of data processes for parallel data loading.
    '''
    processes = []
    for i in range(num_workers):
        process = ReconstructionDataProcess(queue, data_paths, repeat=repeat, train=train)
        process.start()
        processes.append(process)
    return processes


def get_while_running(data_process, data_queue, sleep_time=0):
    while True:
        time.sleep(sleep_time)
        try:
            batch_data, batch_label = data_queue.get_nowait()
        except queue.Empty:
            if not data_process.is_alive():
                break
            else:
                continue
        yield batch_data, batch_label


def test_process():
    from multiprocessing import Queue
    from lib.config import cfg
    from lib.data_io import category_model_id_pair

    cfg.TRAIN.PAD_X = 10
    cfg.TRAIN.PAD_Y = 10

    data_queue = Queue(2)
    category_model_pair = category_model_id_pair(dataset_portion=[0, 0.1])

    data_process = ReconstructionDataProcess(data_queue, category_model_pair)
    data_process.start()
    batch_img, batch_voxel = data_queue.get()

    kill_processes(data_queue, [data_process])


if __name__ == '__main__':
    test_process()