ifind_data.py

"""
Code for reading image in a folder structure and sorting them into labels based on the folder names on the lowest
hirarchy level.

The data then gets stored as 2 *.npz arrays for easier access later on.
"""

import os
import numpy as np
import cv2

SIZE_TOPBOTTOM = 64
SIZE_LEFTRIGHT = 79  #int((float(im.shape[1])/float(im.shape[0]))*64)
NUM_CHANNELS = 1  # i.e. number of colour channels (gray images have 1)
IFIND_ROOT = '/vol/medic01/users/cbaumgar/data/iFind1/iFind1_1200/simple/images/bylabels_manual_one_other/'
DATA_POSTFIX = 'rectangular'

def read_npz(filename):
    """
    Helper function for reading out all arrays in an npz archive
    """
    with np.load(filename) as f:
        values = [f['arr_%d' % i] for i in range(len(f.files))]
        return values

def read_labels(filename):
    """
    Helper function for reading the label data
    """
    values = read_npz(filename)
    y_train = values[0]
    y_test = values[1]
    label_names = values[2].item()
    label_numbers = values[3].item()
    return y_train, y_test, label_names, label_numbers

def read_data(filename):
    """
    Helper function for reading the image data
    """
    values = read_npz(filename)
    X_train = values[0]
    X_test = values[1]
    return X_train, X_test

def load_ifind_data(path):
    """
    This function tries to load the preprocessed training and testing data
    located in `path`
    If it cannot find the data it preprocesses the data if finds in
    the folder IFIND_ROOT
    It returns training and testing data plus two dictionaries for going from label numbers to label names
    and back.
    """
    expected_label_path = os.path.join(path, 'ifind1_scanplane_labels_' + DATA_POSTFIX + '.npz')
    expected_data_path = os.path.join(path, 'ifind1_scanplane_data_' + DATA_POSTFIX + '.npz')

    if os.path.exists(expected_data_path) and os.path.exists(expected_label_path):

        y_train, y_test, label_names, label_numbers = read_labels(expected_label_path)
        X_train, X_test = read_data(expected_data_path)

    else:
        print "Dataset is being created from %s" % IFIND_ROOT
        X_train, y_train, X_test, y_test, label_names, label_numbers = _preprocess_dataset(path)

    return X_train, y_train, X_test, y_test, label_names, label_numbers

"""
 Helper functions for the data loader
"""

def _preprocess_dataset(path):
    """
    Starts the proprocessing
    """
    test_folder = os.path.join(IFIND_ROOT, 'test')
    train_folder = os.path.join(IFIND_ROOT, 'train')

    print "Doing test folder: "
    X_test, y_test, label_numbers = _process_folder(test_folder) # don't augment test dataset (we want to test on original data)
    print X_test.shape
    print y_test.shape
    print X_test.dtype
    print y_test.dtype
    print "Doing train folder: "
    X_train, y_train, dmy = _process_folder(train_folder, True, label_numbers, augment_dataset=True)
    print X_train.shape
    print y_train.shape
    print X_train.dtype
    print y_train.dtype

    label_names = _invert_dictionary(label_numbers)

    # Convert into a format suitable for tensorflow
    X_train = _whiten_images(X_train)
    X_test = _whiten_images(X_test)

    # MAY HAVE TO UNCOMMENT THIS FOR TENSORFLOW CODE!!!!!!!!!!!!!!!!!!
    #y_train = _dense_to_one_hot(y_train, len(label_names))
    #y_test = _dense_to_one_hot(y_test, len(label_names))

    # saving
    labels_path = os.path.join(path, 'ifind1_scanplane_labels_' + DATA_POSTFIX + '.npz')
    print "Saving labels to %s..." % labels_path
    np.savez(labels_path, y_train, y_test, label_names, label_numbers)

    data_path = os.path.join(path, 'ifind1_scanplane_data_' + DATA_POSTFIX + '.npz')
    print "Saving data to %s..." % data_path
    np.savez(data_path, X_train, X_test)

    return X_train, y_train, X_test, y_test, label_names, label_numbers


def _process_folder(folder, use_existing_label_numbers=False, label_numbers=None, augment_dataset=False):
    """
    Goes into a specific folder and preprocess the data it finds there
    :param folder: The folder
    :param use_existing_label_numbers: Use a label dictionary found in a previous run of this function?
    :param label_numbers: If use_existing_label_numbers, give it the label numbers
    :param augment_dataset: Augment the dataset? I.e., from each image create multiple versions
    :return: the preprocessed data, the labels, and the label dictionary
    """

    if not use_existing_label_numbers:
        label_numbers = {}

    label_counter = -1
    labels = []
    data = []
    done_labels = []

    for directory in next(os.walk(folder))[1]:

        label_name = directory.split('/')[-1]
        directory_path = os.path.join(folder, directory)

        for root, dir_list, file_list in os.walk(directory_path):

            for file_name in file_list:

                image_path = os.path.join(root, file_name)

                # Sort out the labels
                if label_name not in done_labels:
                    label_counter += 1
                    if not use_existing_label_numbers:
                        label_number = label_counter
                        label_numbers[label_name] = label_number
                        print "added label %d which is %s" % (label_number, label_name)
                    else:
                        label_number = label_numbers[label_name]
                        print "Doing label %s which has label %d" % (label_name, label_number)

                    done_labels.append(label_name)
                else:
                    #print "! Going back to label %s" % label_name
                    label_number = label_numbers[label_name]

                # Load image
                im = cv2.imread(image_path)

                # figure out if video frame or still-frame image
                is_highres = True if im.shape[1] > 850 else False

                # the crop ranges are square regions from [center, left, right]
                if is_highres:
                    crop_range = (106, 106, 853, 713)
                else:
                    crop_range = (98, 80, 652, 530)

                im_array = _crop_and_preprocess(im, crop_range, is_highres)
                data.append(im_array)
                labels.append(label_number)

    data = np.asarray(data)
    labels = np.asarray(labels, dtype=np.int32)

    return data, labels, label_numbers

def _crop_and_preprocess(im, crop_range, is_highres):
    """
    Helper for cropping images. I determined the constants for contrast, brightness, and the
    Gaussian manually to match the appearance of the video frames
    (using find_image_to_video_mapping.py)
    """
    # crop image
    im_cropped = im.copy()[crop_range[1]:crop_range[3], crop_range[0]:crop_range[2], :]
    # change to gray levels
    im_cropped = cv2.cvtColor(im_cropped, cv2.COLOR_BGR2GRAY)

    # cv2.imshow('debug2', im_cropped)
    # cv2.waitKey(0)

    if is_highres:
        # contrast
        im_adjusted = _adjust_contrast(im_cropped, 0.92)
        # brightness
        im_adjusted = _adjust_brightness(im_adjusted, 13)
        # gaussian blur
        im_adjusted = _blur_gaussian(im_adjusted, 7)
        # resize
    else:
        # if the frame comes from a video then there is no need for adjustment
        im_adjusted = im_cropped

    im_resized = cv2.resize(im_adjusted, (SIZE_LEFTRIGHT, SIZE_TOPBOTTOM))

    # debugging
    # cv2.imshow('debug', im_resized)
    # cv2.waitKey(0)

    # change type
    im_array = im_resized.astype(np.float32)
    # add a dummy dimensions where usually the color channels would be
    im_array = np.reshape(im_array, (1, SIZE_TOPBOTTOM, SIZE_LEFTRIGHT))

    # debugging
    # plt.imshow(np.squeeze(im_array[0,:,:]))
    # plt.show()

    return im_array

def _invert_dictionary(dictionary):
    """
    Helper for inverting a dictionary
    """
    return {v: k for k, v in dictionary.items()}


def _whiten_images(X):
    """
    Helper for making the images zero mean and unit standard deviation i.e. `white`
    """

    X_white = np.zeros(X.shape, dtype=np.float32)

    for ii in xrange(X.shape[0]):

        Xc = X[ii,:,:,:]
        mc = Xc.mean()
        sc = Xc.std()

        Xc_white = np.divide((Xc - mc), sc)

        X_white[ii,:,:,:] = Xc_white

    return X_white


def _dense_to_one_hot(labels_dense, num_classes):
        """
        Convert class labels from scalars to one-hot vectors.
        This means if there are 10 possible labels
        1 -> [1,0,0,0,0,0,0,0,0,0]
        4 -> [0,0,0,1,0,0,0,0,0,0]
         etc...
        """
        num_labels = labels_dense.shape[0]
        index_offset = np.arange(num_labels) * num_classes
        labels_one_hot = np.zeros((num_labels, num_classes))
        labels_one_hot.flat[index_offset + labels_dense.ravel()] = 1
        return labels_one_hot

def _adjust_contrast(image, contrast_factor):
    """
    Helper function for adjusting the contrast of an image using simple multiplication and rebinning into [0,255]
    """
    max_value = image.max()
    min_value = image.min()
    new_image = contrast_factor * image.copy().astype(np.float32)
    new_image[new_image>max_value] = max_value
    new_image[new_image<min_value] = min_value
    return new_image.astype(np.uint8)

def _adjust_brightness(image, brightness_factor):
    """
    Helper function for adjusting the brightness of an image by adding a constant and rebinning into [0,255]
    """
    max_value = image.max()
    min_value = image.min()
    new_image = image.copy().astype(np.uint16) + brightness_factor
    new_image[new_image > max_value] = max_value
    new_image[new_image < min_value] = min_value
    return new_image.astype(np.uint8)

def _blur_average(image, blur_kernel_size):
    """
    Helper function for blurring a square average filter (sliding window)
    """
    blur_kernel = (blur_kernel_size, blur_kernel_size)
    kernel = np.ones(blur_kernel,np.float32)/(blur_kernel_size**2)
    new_image = cv2.filter2D(image.copy().astype(np.float32),-1,kernel)
    return new_image.astype(np.uint8)

def _blur_gaussian(image, blur_kernel_size):
    """
    Helper function for blurring an image with a Gaussian
    """
    blur_kernel = (blur_kernel_size, blur_kernel_size)
    new_image = cv2.GaussianBlur(image.copy().astype(np.float32),blur_kernel,0)
    return new_image.astype(np.uint8)

def _blur_billateral(image, param1, param2):
    """
    Helper function for applying the billateral smoothing filter to an image
    """
    new_image = cv2.bilateralFilter(image.copy().astype(np.float32),param1,param2,param2)
    return new_image.astype(np.uint8)