loader/data_loader.py

from functools import partial
import numpy
import os
import re
import random
import signal
import csv
from PIL import Image
import settings
import numpy as np
from collections import OrderedDict
import cv2
# from scipy.misc import imread
from multiprocessing import Pool, cpu_count
from multiprocessing.pool import ThreadPool
from scipy.ndimage.interpolation import zoom
import sys
import pickle

def load_csv(filename, readfields=None):
    def convert(value):
        if re.match(r'^-?\d+$', value):
            try:
                return int(value)
            except:
                pass
        if re.match(r'^-?[\.\d]+(?:e[+=]\d+)$', value):
            try:
                return float(value)
            except:
                pass
        return value

    with open(filename) as f:
        reader = csv.DictReader(f)
        result = [{k: convert(v) for k, v in row.items()} for row in reader]
        if readfields is not None:
            readfields.extend(reader.fieldnames)
    return result

class AbstractSegmentation:
    def all_names(self, category, j):
        raise NotImplementedError
    def size(self, split=None):
        return 0
    def filename(self, i):
        raise NotImplementedError
    def metadata(self, i):
        return self.filename(i)
    @classmethod
    def resolve_segmentation(cls, m):
        return {}

    def name(self, category, i):
        '''
        Default implemtnation for segmentation_data,
        utilizing all_names.
        '''
        all_names = self.all_names(category, i)
        return all_names[0] if len(all_names) else ''

    def segmentation_data(self, category, i, c=0, full=False):
        '''
        Default implemtnation for segmentation_data,
        utilizing metadata and resolve_segmentation.
        '''
        segs = self.resolve_segmentation(
                self.metadata(i), categories=[category])
        if category not in segs:
            return 0
        data = segs[category]
        if not full and len(data.shape) >= 3:
            return data[0]
        return data


class SegmentationData(AbstractSegmentation):
    '''
    Represents and loads a multi-channel segmentation represented with
    a series of csv files: index.csv lists the images together with
    any label data avilable in each category; category.csv lists
    the categories of segmentations available; and label.csv lists the
    numbers used to describe each label class. In addition, the categories
    each have a separate c_*.csv file describing a dense coding of labels.

    isImageSet - if True, duplicate rgb images in index.csv will be removed
    '''

    def __init__(self, directory, categories=None, require_all=False, isImageSet=False):
        directory = os.path.expanduser(directory)
        self.directory = directory
        with open(os.path.join(directory, settings.INDEX_FILE)) as f:
            self.image = [decode_index_dict(r) for r in csv.DictReader(f)]
        # self.actualFeatIdx = None ### use this dict in tally to map duplicates idx to nonduplicates
        # if isImageSet is True:
        #     self.actualFeatIdx = {}
        #     self.newImgSet = []
        #     self.duplicateDict = {}
        #     for imgRGBIdx, imgRGBData in enumerate(self.image):
        #         if imgRGBData["image"] not in self.duplicateDict:
        #             self.newImgSet.append(imgRGBData)
        #             self.duplicateDict[imgRGBData["image"]] = len(self.newImgSet) - 1
        #         # print("self.duplicateDict[imgRGBData[image]]: ", self.duplicateDict[imgRGBData["image"]])
        #         self.actualFeatIdx[imgRGBIdx] = self.duplicateDict[imgRGBData["image"]] ### Start at 0
        #     self.image = self.newImgSet
        # print("data_set.actualFeatIdx: ", self.actualFeatIdx)
        # sys.exit()
        # print("self.image: ", self.image) ### list
        # print("type self.image: ", type(self.image)) ### list
        # print("len self.image: ", len(self.image)) ### total rows in index.csv
        # sys.exit()
        with open(os.path.join(directory, 'category.csv')) as f:
            self.category = OrderedDict()
            for row in csv.DictReader(f):
                if categories and row['name'] in categories:
                    self.category[row['name']] = row
        categories = self.category.keys()
        with open(os.path.join(directory, 'label.csv')) as f:
            label_data = [decode_label_dict(r) for r in csv.DictReader(f)]
        self.label = build_dense_label_array(label_data) ### Len is label_data+1 (from csv), 0 index is None
        # print("self.label[0]: ", self.label[0]) ### None value (no class specified)
        # print("self.label: ", self.label)
        # sys.exit()
        # Filter out images with insufficient data
        filter_fn = partial(
                index_has_all_data if require_all else index_has_any_data,
                categories=categories)
        self.image = [row for row in self.image if filter_fn(row)]
        # Build dense remapping arrays for labels, so that you can
        # get dense ranges of labels for each category.
        self.category_map = {}
        self.category_unmap = {}
        self.category_label = {}
        for cat in self.category:
            with open(os.path.join(directory, 'c_%s.csv' % cat)) as f:
                c_data = [decode_label_dict(r) for r in csv.DictReader(f)]
            self.category_unmap[cat], self.category_map[cat] = (
                    build_numpy_category_map(c_data))
            self.category_label[cat] = build_dense_label_array(
                    c_data, key='code')
        # print("category_unmap: ", self.category_unmap)
        self.labelcat = self.onehot(self.primary_categories_per_index()) # (480,1)
        ### labelcat - all ones

    def primary_categories_per_index(ds):
        '''
        Returns an array of primary category numbers for each label, where the
        first category listed in ds.category_names is given category number 0.
        '''
        catmap = {}
        categories = ds.category_names()
        for cat in categories:
            imap = ds.category_index_map(cat)
            if len(imap) < ds.label_size(None):
                imap = np.concatenate((imap, np.zeros(
                    ds.label_size(None) - len(imap), dtype=imap.dtype)))
            catmap[cat] = imap
        result = []
        for i in range(ds.label_size(None)):
            maxcov, maxcat = max(
                (ds.coverage(cat, catmap[cat][i]) if catmap[cat][i] else 0, ic)
                for ic, cat in enumerate(categories))
            result.append(maxcat)
        return np.array(result)

    def onehot(self, arr, minlength=None):
        '''
        Expands an array of integers in one-hot encoding by adding a new last
        dimension, leaving zeros everywhere except for the nth dimension, where
        the original array contained the integer n.  The minlength parameter is
        used to indcate the minimum size of the new dimension.
        '''
        length = np.amax(arr) + 1
        if minlength is not None:
            length = max(minlength, length)
        result = np.zeros(arr.shape + (length,))
        result[list(np.indices(arr.shape)) + [arr]] = 1
        return result


    def all_names(self, category, j):
        '''All English synonyms for the given label'''
        if category is not None:
            j = self.category_unmap[category][j]
        return [self.label[j]['name']] + self.label[j]['syns']

    def size(self, split=None):
        '''The number of images in this data set.'''
        if split is None:
            return len(self.image)
        return len([im for im in self.image if im['split'] == split])

    def filename(self, i):
        '''The filename of the ith jpeg (original image).'''
        return os.path.join(self.directory, 'images', self.image[i]['image'])

    def split(self, i):
        '''Which split contains item i.'''
        return self.image[i]['split']

    def metadata(self, i):
        '''Extract metadata for image i, For efficient data loading.'''
        return self.directory, self.image[i]

    meta_categories = ['image', 'split', 'ih', 'iw', 'sh', 'sw']

    @classmethod
    def resolve_segmentation(cls, m, categories=None, segm_to_label=None):
        '''
        Resolves a full segmentation, potentially in a differenct process,
        for efficient multiprocess data loading.
        '''
        directory, row = m
        result = {}
        for cat, d in row.items():
            if cat in cls.meta_categories:
                continue
            if not wants(cat, categories):
                continue
            if all(isinstance(data, int) for data in d):
                result[cat] = d
                continue
            out = numpy.empty((len(d), row['sh'], row['sw']), dtype=numpy.int16)
            for i, channel in enumerate(d):
                if isinstance(channel, int):
                    out[i] = channel
                else:
                    segmFilenameSplit = channel.split('/')
                    segmFileToLabelName = segmFilenameSplit[-2] + "/" + segmFilenameSplit[-1]
                    if 'seg_label' not in result:
                        result['seg_label'] = []
                    result['seg_label'].append(segm_to_label[segmFileToLabelName])
                    # print("os.path.join(directory, 'images', channel)): ", os.path.join(directory, 'images', channel))
                    rgb = cv2.resize(cv2.imread(os.path.join(directory, 'images', channel)), (settings.SEGM_SIZE, settings.SEGM_SIZE))
                    rgb[:,:,0] = 0
                    rgb[:,:,2] = np.where(rgb[:,:,2]>0, 235, 0)
                    rgb[:,:,1] = np.where(rgb[:,:,2]>0, 1, 0)
                    out[i] = rgb[:,:,0] + rgb[:,:,1] * 256
            result[cat] = out
        return result, (row['sh'], row['sw'])

    def label_size(self, category=None):
        '''
        Returns the number of distinct labels (plus zero), i.e., one
        more than the maximum label number.  If a category is specified,
        returns the number of distinct labels within that category.
        '''
        if category is None:
            return len(self.label)
        else:
            return len(self.category_unmap[category])

    def name(self, category, j):
        '''
        Returns an English name for the jth label.  If a category is
        specified, returns the name for the category-specific nubmer j.
        If category=None, then treats j as a fully unified index number.
        '''
        if category is not None:
            j = self.category_unmap[category][j]
        return self.label[j]['name']

    def frequency(self, category, j):
        '''
        Returns the number of images for which the label appears.
        '''
        if category is not None:
            return self.category_label[category][j]['frequency']
        return self.label[j]['frequency']

    def coverage(self, category, j):
        '''
        Returns the pixel coverage of the label in units of whole-images.
        '''
        if category is not None:
            return self.category_label[category][j]['coverage']
        return self.label[j]['coverage']

    def category_names(self):
        '''
        Returns the set of category names.
        '''
        return list(self.category.keys())

    def category_frequency(self, category):
        '''
        Returns the number of images touched by a category.
        '''
        return float(self.category[category]['frequency'])

    def primary_categories_per_index(self, categories=None):
        '''
        Returns an array of primary category numbers for each label, where
        catagories are indexed according to the list of categories passed,
        or self.category_names() if none.
        '''
        if categories is None:
            categories = self.category_names()
        # Make lists which are nonzero for labels in a category
        catmap = {}
        for cat in categories:
            imap = self.category_index_map(cat)
            if len(imap) < self.label_size(None):
                imap = numpy.concatenate((imap, numpy.zeros(
                    self.label_size(None) - len(imap), dtype=imap.dtype)))
            catmap[cat] = imap
        # For each label, find the category with maximum coverage.
        result = []
        for i in range(self.label_size(None)):
            maxcov, maxcat = max(
                    (self.coverage(cat, catmap[cat][i])
                        if catmap[cat][i] else 0, ic)
                    for ic, cat in enumerate(categories))
            result.append(maxcat)
        # Return the max-coverage cateogry for each label.
        return numpy.array(result)


    def segmentation_data(self, category, i, c=0, full=False, out=None):
        '''
        Returns a 2-d numpy matrix with segmentation data for the ith image,
        restricted to the given category.  By default, maps all label numbers
        to the category-specific dense mapping described in the c_*.csv
        listing; but can be asked to expose the fully unique indexing by
        using full=True.
        '''
        row = self.image[i]
        data_channels = row.get(category, ())
        if c >= len(data_channels):
            channel = 0  # Deal with unlabeled data in this category
        else:
            channel = data_channels[c]
        if out is None:
            out = numpy.empty((row['sh'], row['sw']), dtype=numpy.int16)
        if isinstance(channel, int):
            if not full:
                channel = self.category_map[category][channel]
            out[:,:] = channel  # Single-label for the whole image
            return out
        png = cv2.resize(cv2.imread(os.path.join(self.directory, 'images', channel)), (settings.SEGM_SIZE, settings.SEGM_SIZE))
        png[:,:,0] = 0
        png[:,:,2] = np.where(png[:,:,2]>0, 235, 0)
        png[:,:,1] = np.where(png[:,:,2]>0, 1, 0)
        if full:
            # Full case: just combine png channels.
            out[...] = png[:,:,0] + png[:,:,1] * 256
        else:
            # Dense case: combine png channels and apply the category map.
            catmap = self.category_map[category]
            out[...] = catmap[png[:,:,0] + png[:,:,1] * 256]
        return out

    def full_segmentation_data(self, i,
            categories=None, max_depth=None, out=None):
        '''
        Returns a 3-d numpy tensor with segmentation data for the ith image,
        with multiple layers represnting multiple lables for each pixel.
        The depth is variable depending on available data but can be
        limited to max_depth.
        '''
        row = self.image[i]
        if categories:
            groups = [d for cat, d in row.items() if cat in categories and d]
        else:
            groups = [d for cat, d in row.items() if d and (
                cat not in self.meta_categories)]
        depth = sum(len(c) for c in groups)
        if max_depth is not None:
            depth = min(depth, max_depth)
        # Allocate an array if not already allocated.
        if out is None:
            out = numpy.empty((depth, row['sh'], row['sw']), dtype=numpy.int16)
        i = 0
        # Stack up the result segmentation one channel at a time
        for group in groups:
            for channel in group:
                if isinstance(channel, int):
                    out[i] = channel
                else:
                    png = cv2.resize(cv2.imread(os.path.join(self.directory, 'images', channel)), (settings.SEGM_SIZE, settings.SEGM_SIZE))
                    png[:,:,0] = 0
                    png[:,:,2] = np.where(png[:,:,2]>0, 235, 0)
                    png[:,:,1] = np.where(png[:,:,2]>0, 1, 0)
                    out[i] = png[:,:,0] + png[:,:,1] * 256
                i += 1
                if i == depth:
                    return out
        # Return above when we get up to depth
        assert False

    def category_index_map(self, category):
        return numpy.array(self.category_map[category])

def build_dense_label_array(label_data, key='number', allow_none=False):
    '''
    Input: set of rows with 'number' fields (or another field name key).
    Output: array such that a[number] = the row with the given number.
    '''
    result = [None] * (max([d[key] for d in label_data]) + 1)
    for d in label_data:
        result[d[key]] = d
    # Fill in none
    if not allow_none:
        example = label_data[0]
        def make_empty(k):
            return dict((c, k if c is key else type(v)())
                    for c, v in example.items())
        for i, d in enumerate(result):
            if d is None:
                result[i] = dict(make_empty(i))
    return result

def build_numpy_category_map(map_data, key1='code', key2='number'):
    '''
    Input: set of rows with 'number' fields (or another field name key).
    Output: array such that a[number] = the row with the given number.
    '''
    results = list(numpy.zeros((max([d[key] for d in map_data]) + 1),
            dtype=numpy.int16) for key in (key1, key2))
    for d in map_data:
        results[0][d[key1]] = d[key2]
        results[1][d[key2]] = d[key1]
    return results

def decode_label_dict(row):
    result = {}
    for key, val in row.items():
        if key == 'category':
            result[key] = dict((c, int(n))
                for c, n in [re.match('^([^(]*)\(([^)]*)\)$', f).groups()
                    for f in val.split(';')])
        elif key == 'name':
            result[key] = val
        elif key == 'syns':
            result[key] = val.split(';')
        elif re.match('^\d+$', val):
            result[key] = int(val)
        elif re.match('^\d+\.\d*$', val):
            result[key] = float(val)
        else:
            result[key] = val
    return result

def decode_index_dict(row):
    result = {}
    for key, val in row.items():
        if key in ['image', 'split']:
            result[key] = val
        elif key in ['sw', 'sh', 'iw', 'ih']:
            result[key] = int(val)
        else:
            item = [s for s in val.split(';') if s]
            for i, v in enumerate(item):
                if re.match('^\d+$', v):
                    item[i] = int(v)
            result[key] = item
    return result

def index_has_any_data(row, categories):
    for c in categories:
        for data in row[c]:
            if data: return True
    return False

def index_has_all_data(row, categories):
    for c in categories:
        cat_has = False
        for data in row[c]:
            if data:
                cat_has = True
                break
        if not cat_has:
            return False
    return True

class SegmentationPrefetcher:
    '''
    SegmentationPrefetcher will prefetch a bunch of segmentation
    images using a multiprocessing pool, so you do not have to wait
    around while the files get opened and decoded.  Just request
    batches of images and segmentations calling fetch_batch().
    '''
    def __init__(self, segmentation, split=None, randomize=False,
            segmentation_shape=None, categories=None, once=False,
            start=None, end=None, batch_size=4, ahead=4, thread=False):
        '''
        Constructor arguments:
        segmentation: The AbstractSegmentation to load.
        split: None for no filtering, or 'train' or 'val' etc.
        randomize: True to randomly shuffle order, or a random seed.
        categories: a list of categories to include in each batch.
        batch_size: number of data items for each batch.
        ahead: the number of data items to prefetch ahead.
        '''
        self.segmentation = segmentation
        self.segm_to_label  = None
        with open(settings.SEGM_TO_LABEL_PKL, 'rb') as f:
            self.segm_to_label = pickle.load(f)
        self.split = split
        self.randomize = randomize
        self.random = random.Random()
        if randomize is not True:
            self.random.seed(randomize)
        self.categories = categories
        self.once = once
        self.batch_size = batch_size
        self.ahead = ahead
        # Initialize the multiprocessing pool
        n_procs = cpu_count()
        if thread:
            self.pool = ThreadPool(processes=n_procs)
        else:
            original_sigint_handler = setup_sigint()
            self.pool = Pool(processes=n_procs, initializer=setup_sigint)
            restore_sigint(original_sigint_handler)
        # Prefilter the image indexes of interest
        if start is None:
            start = 0
        if end is None:
            end = segmentation.size()
        self.indexes = range(start, end)
        if split:
            self.indexes = [i for i in self.indexes
                    if segmentation.split(i) == split]
        if self.randomize:
            self.random.shuffle(self.indexes)
        self.index = 0
        self.result_queue = []
        self.segmentation_shape = segmentation_shape
        # Get dense catmaps
        self.catmaps = [
                segmentation.category_index_map(cat) if cat != 'image' else None
                for cat in categories]

    def next_job(self):
        if self.index < 0:
            return None
        j = self.indexes[self.index]
        result = (j,
                self.segmentation.__class__,
                self.segmentation.metadata(j),
                self.segmentation.filename(j),
                self.categories,
                self.segm_to_label,
                self.segmentation_shape)
        self.index += 1
        if self.index >= len(self.indexes):
            if self.once:
                self.index = -1
            else:
                self.index = 0
                if self.randomize:
                    # Reshuffle every time through
                    self.random.shuffle(self.indexes)
        return result

    def batches(self):
        '''Iterator for all batches'''
        while True:
            batch = self.fetch_batch()
            if batch is None:
                break
            else:
                yield batch
    # def batches(self):
    #     '''Iterator for all batches'''
    #     while True:
    #         batch = self.fetch_batch()
    #         if batch is None:
    #             raise StopIteration
    #         yield batch
    def fetch_batch(self):
        '''Returns a single batch as an array of dictionaries.'''
        try:
            self.refill_tasks()
            if len(self.result_queue) == 0:
                return None
            result = self.result_queue.pop(0)
            return result.get(31536000)
        except KeyboardInterrupt:
            print("Caught KeyboardInterrupt, terminating workers")
            self.pool.terminate()
            raise

    def fetch_tensor_batch(self, bgr_mean=None, global_labels=False):
        '''Iterator for batches as arrays of tensors.'''
        batch = self.fetch_batch()
        return self.form_caffe_tensors(batch, bgr_mean, global_labels)

    def tensor_batches(self, bgr_mean=None, global_labels=False):
        '''Returns a single batch as an array of tensors, one per category.'''
        while True:
            batch = self.fetch_tensor_batch(
                    bgr_mean=bgr_mean, global_labels=global_labels)
            if batch is None:
                break
            else:
                yield batch

    def form_caffe_tensors(self, batch, bgr_mean=None, global_labels=False):
        # Assemble a batch in [{'cat': data,..},..] format into
        # an array of batch tensors, the first for the image, and the
        # remaining for each category in self.categories, in order.
        # This also applies a random flip if needed
        if batch is None:
            return None
        batches = [[] for c in self.categories]
        for record in batch:
            default_shape = (1, record['sh'], record['sw'])
            for c, cat in enumerate(self.categories):
                if cat == 'image':
                    # Normalize image with right RGB order and mean
                    batches[c].append(normalize_image(
                        record[cat], bgr_mean))
                elif global_labels:
                    batches[c].append(normalize_label(
                        record[cat], default_shape, flatten=True))
                else:
                    catmap = self.catmaps[c]
                    batches[c].append(catmap[normalize_label(
                        record[cat], default_shape, flatten=True)])
        return [numpy.concatenate(tuple(m[numpy.newaxis] for m in b))
                for b in batches]

    def refill_tasks(self):
        # It will call the sequencer to ask for a sequence
        # of batch_size jobs (indexes with categories)
        # Then it will call pool.map_async
        while len(self.result_queue) < self.ahead:
            data = []
            while len(data) < self.batch_size:
                job = self.next_job()
                if job is None:
                    break
                data.append(job)
            if len(data) == 0:
                return
            self.result_queue.append(self.pool.map_async(prefetch_worker, data))

    def close(self):
        while len(self.result_queue):
            result = self.result_queue.pop(0)
            if result is not None:
                result.wait(0.001)
        self.pool.close()
        self.poool.cancel_join_thread()

def prefetch_worker(d):
    if d is None:
        return None
    j, typ, m, fn, categories, segm_to_label, segmentation_shape = d
    segs, shape = typ.resolve_segmentation(m, categories=categories, segm_to_label=segm_to_label)
    if segmentation_shape is not None:
        for k, v in segs.items():
            print("k: ", k)
            segs[k] = scale_segmentation(v, segmentation_shape)
        shape = segmentation_shape
    # Some additional metadata to provide
    segs['sh'], segs['sw'] = shape
    segs['i'] = j
    segs['fn'] = fn
    if categories is None or 'image' in categories:
        segs['image'] = np.asarray(Image.open(fn).convert('L').resize((settings.IMG_SIZE, settings.IMG_SIZE)))
    return segs
# def convertRGBToGray(rgbImg):
#     return np.dot(rgbImg[...,:3], [0.114, 0.587, 0.299])
def scale_segmentation(segmentation, dims, crop=False):
    '''
    Zooms a 2d or 3d segmentation to the given dims, using nearest neighbor.
    '''
    shape = numpy.shape(segmentation)
    if len(shape) < 2 or shape[-2:] == dims:
        return segmentation
    peel = (len(shape) == 2)
    if peel:
        segmentation = segmentation[numpy.newaxis]
    levels = segmentation.shape[0]
    result = numpy.zeros((levels, ) + dims,
            dtype=segmentation.dtype)
    ratio = (1,) + tuple(res / float(orig)
            for res, orig in zip(result.shape[1:], segmentation.shape[1:]))
    if not crop:
        safezoom(segmentation, ratio, output=result, order=0)
    else:
        ratio = max(ratio[1:])
        height = int(round(dims[0] / ratio))
        hmargin = (segmentation.shape[0] - height) // 2
        width = int(round(dims[1] / ratio))
        wmargin = (segmentation.shape[1] - height) // 2
        safezoom(segmentation[:, hmargin:hmargin+height,
            wmargin:wmargin+width],
            (1, ratio, ratio), output=result, order=0)
    if peel:
        result = result[0]
    return result

def safezoom(array, ratio, output=None, order=0):
    '''Like numpy.zoom, but does not crash when the first dimension
    of the array is of size 1, as happens often with segmentations'''
    dtype = array.dtype
    if array.dtype == numpy.float16:
        array = array.astype(numpy.float32)
    if array.shape[0] == 1:
        if output is not None:
            output = output[0,...]
        result = zoom(array[0,...], ratio[1:],
                output=output, order=order)
        if output is None:
            output = result[numpy.newaxis]
    else:
        result = zoom(array, ratio, output=output, order=order)
        if output is None:
            output = result
    return output.astype(dtype)

def setup_sigint():
    import threading
    if not isinstance(threading.current_thread(), threading._MainThread):
        return None
    return signal.signal(signal.SIGINT, signal.SIG_IGN)

def restore_sigint(original):
    import threading
    if not isinstance(threading.current_thread(), threading._MainThread):
        return
    if original is None:
        original = signal.SIG_DFL
    signal.signal(signal.SIGINT, original)

def wants(what, option):
    if option is None:
        return True
    return what in option

def normalize_image(rgb_image, bgr_mean):
    """
    Load input image and preprocess for Caffe:
    - cast to float
    - switch channels RGB -> BGR
    - subtract mean
    - transpose to channel x height x width order
    """
    # img = numpy.array(rgb_image, dtype=numpy.float32)
    # if (img.ndim == 2):
    #     img = numpy.repeat(img[:,:,None], 3, axis = 2)
    # # img = img[:,:,::-1]
    # if bgr_mean is not None:
    #     img -= bgr_mean
    rgb_image = np.expand_dims(rgb_image, axis=0)
    # rgb_image = rgb_image.transpose((2,0,1))
    # print("rgb_image shape: ", rgb_image.shape)
    return rgb_image

### Original code
# def normalize_image(rgb_image, bgr_mean):
#     """
#     Load input image and preprocess for Caffe:
#     - cast to float
#     - switch channels RGB -> BGR
#     - subtract mean
#     - transpose to channel x height x width order
#     """
#     img = numpy.array(rgb_image, dtype=numpy.float32)
#     if (img.ndim == 2):
#         img = numpy.repeat(img[:,:,None], 3, axis = 2)
#     img = img[:,:,::-1]
#     if bgr_mean is not None:
#         img -= bgr_mean
#     img = img.transpose((2,0,1))
#     return img

def normalize_label(label_data, shape, flatten=False):
    """
    Given a 0, 1, 2, or 3-dimensional label_data and a default
    shape of the form (1, y, x), returns a 3d tensor by
    """
    dims = len(numpy.shape(label_data))
    if dims <= 2:
        # Scalar data on this channel: fill shape
        if dims == 1:
            if flatten:
                label_data = label_data[0] if len(label_data) else 0
            else:
                return (numpy.ones(shape, dtype=numpy.int16) *
                        numpy.asarray(label_data, dtype=numpy.int16)
                            [:, numpy.newaxis, numpy.newaxis])
        return numpy.full(shape, label_data, dtype=numpy.int16)
    else:
        if dims == 3:
            if flatten:
                label_data = label_data[0]
            else:
                return label_data
        return label_data[numpy.newaxis]

if __name__ == '__main__':
    data = SegmentationData('broden1_227')
    pd = SegmentationPrefetcher(data,categories=data.category_names()+['image'],once=True)
    bs = pd.batches().next()