train_fcn

import numpy as np
from skimage.io import imshow, imread
import pandas as pd
import matplotlib.pyplot as plt

from time import time
import itertools

import sklearn.preprocessing


import keras.models
import keras.layers
import keras.optimizers
import keras.metrics
import keras.callbacks
import keras.utils
from keras.regularizers import l1, l2

import cv2
import scipy.ndimage
import imgaug
from imgaug import augmenters as iaa


NAME = '13_bin_multiple_images'
LR = 0.00003
BATCH_SIZE = 16
EPOCHS = 2048
NUM_CLASSES = 2
IMAGE_SHAPE = (256, 256)
TOTAL_FILES = 21


def load_train_image(index):
    assert(1 <= index <= TOTAL_FILES)
    return imread('regions/binary/g{}_g.tif'.format(index))

def load_val_image(index):
    assert(1 <= index <= TOTAL_FILES)
    if index == 6:
        return imread('regions/binary/g6_b.tif')
    return imread('regions/binary/g{}_y.tif'.format(index))

def load_mask(index):
    assert(1 <= index <= TOTAL_FILES)
    mask = imread('regions/binary/g{}_m.tif'.format(index))[:, :, 0]
    mask[mask > 0] = 1
    return mask

aug = iaa.Sequential([
        iaa.Fliplr(0.5),
        iaa.Flipud(0.5),
        # iaa.ContrastNormalization((0.75, 1.5)),
        iaa.Multiply((0.9, 1.1), per_channel=0.6),
        iaa.Sometimes(0.4,
            iaa.GaussianBlur(sigma=(0, 0.5))
        )
    ])


def get_rotated_subregion(image, shape, angle, shift=(0, 0), scale=1):
    radius = max(shape) / 2 ** 0.5
    shift = (min(max(shift[1], radius), image.shape[1] - radius),
             min(max(shift[0], radius), image.shape[0] - radius))
    cosa = np.cos(angle)
    sina = np.sin(angle)

    matShiftB = np.array([[1., 0., -shift[0]], [0., 1., -shift[1]], [0., 0., 1.]])
    matRot = np.array([[cosa, sina, 0.], [-sina, cosa, 0.], [0., 0., 1.]])
    matShiftF = np.array([[1., 0., shape[0] / 2.], [0., 1., shape[1] / 2.], [0., 0., 1.]])
    matScale = np.array([[scale, 0., 0.], [0., scale, 0.], [0., 0., 1.]])
    matTotal = matShiftF.dot(matRot.dot(matScale.dot(matShiftB)))

    return cv2.warpAffine(image, matTotal[:2, :], shape)


def image_generator(batch_size, load_image, load_mask):
    while True:
        image_index = np.random.randint(1, TOTAL_FILES + 1)
        big_image = load_image(image_index)
        mask = load_mask(image_index)

        angles = np.random.uniform(0, np.pi * 2, (batch_size,))
        angles = np.random.uniform(0, np.pi * 2, (batch_size,))
        shifts_x = np.random.uniform(0, mask.shape[0], (batch_size,))
        shifts_y = np.random.uniform(0, mask.shape[1], (batch_size,))

        images = np.array([get_rotated_subregion(
                big_image, IMAGE_SHAPE, angle, shift
            ) for angle, shift in zip(angles, zip(shifts_x, shifts_y))
        ])
        Ys = np.array([get_rotated_subregion(mask, IMAGE_SHAPE,
            angle, shift) for angle, shift in zip(angles,
                zip(shifts_x, shifts_y))])
        labels = Ys.reshape(batch_size, IMAGE_SHAPE[0], IMAGE_SHAPE[1], 1)
        # labels = keras.utils.to_categorical(Ys.flatten(),
        #         NUM_CLASSES).reshape(batch_size,
        #                              IMAGE_SHAPE[0],
        #                              IMAGE_SHAPE[1], NUM_CLASSES)
        yield images, labels


def get_model():
    CONVS = 2
    LAYERS = 6
    FILTERS = 8

    dataInput = keras.layers.Input(shape=(None, None, 3))
    x = dataInput
    # -------- Encoder --------
    lstMaxPools = []
    for cc in range(LAYERS):
        for ii in range(CONVS):
            x = keras.layers.Conv2D(filters=FILTERS * (2**cc),
                kernel_size=(3, 3),
                padding='same',
                activation='relu')(x)
        lstMaxPools.append(x)
        x = keras.layers.MaxPooling2D(pool_size=(2, 2))(x)
    # -------- Decoder --------
    for cc in range(LAYERS):
        for ii in range(CONVS):
            x = keras.layers.Conv2D(filters=FILTERS * (2 ** (LAYERS - 1 - cc)),
                kernel_size=(3, 3),
                padding='same',
                activation='relu')(x)
        x = keras.layers.UpSampling2D(size=(2, 2))(x)
        if cc + 2 < LAYERS:
            x = keras.layers.concatenate([x, lstMaxPools[-1 - cc]], axis=-1)

    # 1x1 Convolution: emulation of Dense layer
    if NUM_CLASSES == 2:
        x = keras.layers.Conv2D(filters=1, kernel_size=(1,1), padding='valid',
                                activation='sigmoid')(x)
    else:
        x = keras.layers.Conv2D(filters=NUM_CLASSES, kernel_size=(1, 1),
                                padding='valid')(x)
        x = keras.layers.Activation('softmax')(x)
    model = keras.models.Model(dataInput, x)

    model.compile(keras.optimizers.adam(lr=LR), 'binary_crossentropy',
                  ['binary_accuracy'])
    return model


def limit_gpu_usage():
    config = keras.backend.tf.ConfigProto(
        gpu_options=keras.backend.tf.GPUOptions(per_process_gpu_memory_fraction=0.2))
    keras.backend.tensorflow_backend.set_session(keras.backend.tf.Session(config=config))


def fit():
    steps = 192
    steps_val = steps

    model = get_model()

    hist = model.fit_generator(
        image_generator(BATCH_SIZE, load_train_image, load_mask),
        steps_per_epoch=steps,
        validation_data=image_generator(BATCH_SIZE, load_val_image, load_mask),
        validation_steps=steps_val,
        epochs=EPOCHS,
        callbacks=[
            keras.callbacks.TensorBoard(
                './logs/fcn/' + NAME,
                write_images=False,
                batch_size=BATCH_SIZE
            ),
            keras.callbacks.ModelCheckpoint(
                NAME + '.h5', verbose=False,
                save_best_only=True, monitor='val_loss'
            )
        ]
    )


if __name__ == '__main__':
    limit_gpu_usage()

    fit()