# part of this script was taken from https://github.com/jocicmarko/ultrasound-nerve-segmentation
import argparse
from glob import glob
import numpy as np
from PIL import Image
from keras import backend as K
from keras import losses
from keras.callbacks import ModelCheckpoint, EarlyStopping, ReduceLROnPlateau
from keras.layers import Input, MaxPooling2D
from keras.layers import concatenate, Conv2D, Conv2DTranspose, Dropout, ReLU
from keras.models import Model
from keras.optimizers import Adam
from numpy import random
import tensorflow as tf
from aug_utils import random_augmentation
from random import randint
batch_size = 16
input_shape = (64, 64)
def custom_activation(x):
return K.relu(x, alpha=0.0, max_value=1)
def focal_loss(gamma=2., alpha=.25):
def focal_loss_fixed(y_true, y_pred):
pt_1 = tf.where(tf.equal(y_true, 1), y_pred, tf.ones_like(y_pred))
pt_0 = tf.where(tf.equal(y_true, 0), y_pred, tf.zeros_like(y_pred))
return -K.sum(alpha * K.pow(1. - pt_1, gamma) * K.log(pt_1))-K.sum((1-alpha) * K.pow( pt_0, gamma) * K.log(1. - pt_0))
return focal_loss_fixed
smooth = 1.
def get_unet(do=0, activation=ReLU):
inputs = Input((None, None, 3))
conv1 = Dropout(do)(activation()(Conv2D(32, (3, 3), padding='same')(inputs)))
conv1 = Dropout(do)(activation()(Conv2D(32, (3, 3), padding='same')(conv1)))
pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)
conv2 = Dropout(do)(activation()(Conv2D(64, (3, 3), padding='same')(pool1)))
conv2 = Dropout(do)(activation()(Conv2D(64, (3, 3), padding='same')(conv2)))
pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)
conv3 = Dropout(do)(activation()(Conv2D(128, (3, 3), padding='same')(pool2)))
conv3 = Dropout(do)(activation()(Conv2D(128, (3, 3), padding='same')(conv3)))
pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)
conv4 = Dropout(do)(activation()(Conv2D(256, (3, 3), padding='same')(pool3)))
conv4 = Dropout(do)(activation()(Conv2D(256, (3, 3), padding='same')(conv4)))
pool4 = MaxPooling2D(pool_size=(2, 2))(conv4)
conv5 = Dropout(do)(activation()(Conv2D(512, (3, 3), padding='same')(pool4)))
conv5 = Dropout(do)(activation()(Conv2D(512, (3, 3), padding='same')(conv5)))
up6 = concatenate([Conv2DTranspose(256, (2, 2), strides=(2, 2), padding='same')(conv5), conv4], axis=3)
conv6 = Dropout(do)(activation()(Conv2D(256, (3, 3), padding='same')(up6)))
conv6 = Dropout(do)(activation()(Conv2D(256, (3, 3), padding='same')(conv6)))
up7 = concatenate([Conv2DTranspose(128, (2, 2), strides=(2, 2), padding='same')(conv6), conv3], axis=3)
conv7 = Dropout(do)(activation()(Conv2D(128, (3, 3), padding='same')(up7)))
conv7 = Dropout(do)(activation()(Conv2D(128, (3, 3), padding='same')(conv7)))
up8 = concatenate([Conv2DTranspose(64, (2, 2), strides=(2, 2), padding='same')(conv7), conv2], axis=3)
conv8 = Dropout(do)(activation()(Conv2D(64, (3, 3), padding='same')(up8)))
conv8 = Dropout(do)(activation()(Conv2D(64, (3, 3), padding='same')(conv8)))
up9 = concatenate([Conv2DTranspose(32, (2, 2), strides=(2, 2), padding='same')(conv8), conv1], axis=3)
conv9 = Dropout(do)(activation()(Conv2D(32, (3, 3), padding='same')(up9)))
conv9 = Dropout(do)(activation()(Conv2D(32, (3, 3), padding='same')(conv9)))
conv10 = Conv2D(1, (1, 1), activation='sigmoid')(conv9)
model = Model(inputs=[inputs], outputs=[conv10])
model.compile(optimizer=Adam(lr=1e-3), loss=losses.binary_crossentropy, metrics=['accuracy'])
return model
def read_input(path):
x = np.array(Image.open(path))/255.
return x
def read_gt(path):
x = np.array(Image.open(path))/255.
return x[..., np.newaxis]
def random_crop(img, mask, crop_size=input_shape[0]):
imgheight= img.shape[0]
imgwidth = img.shape[1]
i = randint(0, imgheight-crop_size)
j = randint(0, imgwidth-crop_size)
return img[i:(i+crop_size), j:(j+crop_size), :], mask[i:(i+crop_size), j:(j+crop_size)]
def gen(data, au=False):
while True:
repeat = 4
index= random.choice(list(range(len(data))), batch_size//repeat)
index = list(map(int, index))
list_images_base = [read_input(data[i][0]) for i in index]
list_gt_base = [read_gt(data[i][1]) for i in index]
list_images = []
list_gt = []
for image, gt in zip(list_images_base, list_gt_base):
for _ in range(repeat):
image_, gt_ = random_crop(image.copy(), gt.copy())
list_images.append(image_)
list_gt.append(gt_)
list_images_aug = []
list_gt_aug = []
for image, gt in zip(list_images, list_gt):
if au:
image, gt = random_augmentation(image, gt)
list_images_aug.append(image)
list_gt_aug.append(gt)
yield np.array(list_images_aug), np.array(list_gt_aug)
if __name__ == '__main__':
ap = argparse.ArgumentParser()
ap.add_argument("-d", "--dropout", required=False,
help="dropout", type=float, default=0.1)
ap.add_argument("-a", "--activation", required=False,
help="activation", default="ReLU")
args = vars(ap.parse_args())
activation = globals()[args['activation']]
model_name = "baseline_unet_do_%s_activation_%s_"%(args['dropout'], args['activation'])
print("Model : %s"%model_name)
train_data = list(zip(sorted(glob('../input/DRIVE/training/images/*.tif')),
sorted(glob('../input/DRIVE/training/1st_manual/*.gif'))))
model = get_unet(do=args['dropout'], activation=activation)
file_path = model_name + "weights.best.hdf5"
try:
model.load_weights(file_path, by_name=True)
except:
pass
history = model.fit_generator(gen(train_data, au=True), epochs=100, verbose=2,
steps_per_epoch= 100*len(train_data)//batch_size,
use_multiprocessing=True, workers=16)
model.save_weights(file_path)
