from __future__ import absolute_import, division, print_function
import argparse
from datetime import datetime
from os import environ
import keras.backend as K
from keras.datasets import cifar10
import keras
import numpy as np
import pandas as pd
import tensorflow as tf
from random_eraser import get_random_eraser
from skimage import io, color, transform, exposure
from keras.applications import MobileNet, ResNet50
from keras.applications import DenseNet169, InceptionResNetV2,DenseNet201
from keras.applications.vgg19 import VGG19
from keras.callbacks import (EarlyStopping, ModelCheckpoint, ReduceLROnPlateau, TensorBoard)
from keras.layers import Dense, GlobalAveragePooling2D, GlobalMaxPooling2D
from keras.metrics import binary_accuracy, binary_crossentropy, kappa_error, kullback_leibler_divergence
from keras.models import Model
from keras.optimizers import SGD, Adam
from keras.preprocessing.image import ImageDataGenerator
from sklearn.utils import class_weight
from keras.utils.training_utils import multi_gpu_model
from custom_layers import *
environ['TF_CPP_MIN_LOG_LEVEL'] = '3' # Shut up tensorflow!
print("tf : {}".format(tf.__version__))
print("keras : {}".format(keras.__version__))
print("numpy : {}".format(np.__version__))
print("pandas : {}".format(pd.__version__))
parser = argparse.ArgumentParser(description='Hyperparameters')
parser.add_argument('--classes', default=1, type=int)
parser.add_argument('--workers', default=4, type=int)
parser.add_argument('--epochs', default=30, type=int)
parser.add_argument('-b', '--batch-size', default=8, type=int, help='mini-batch size')
parser.add_argument('--lr', '--learning-rate', default=1e-4, type=float)
parser.add_argument('--lr-wait', default=10, type=int, help='how long to wait on plateu')
parser.add_argument('--decay', default=1e-4, type=float)
parser.add_argument('--momentum', default=0.9, type=float)
parser.add_argument('--resume', default='', type=str, help='path to latest checkpoint')
parser.add_argument('--fullretrain', dest='fullretrain', action='store_true', help='retrain all layers of the model')
parser.add_argument('--seed', default=1953, type=int, help='random seed')
parser.add_argument('--img_channels', default=3, type=int)
parser.add_argument('--img_size', default=499, type=int)
parser.add_argument('--early_stop', default=20, type=int)
#def preprocess_img():
# def preprocess_img(img):
# Histogram normalization in v channel
# args = parser.parse_args()
# hsv = color.rgb2hsv(img)
# hsv[:, :, 2] = exposure.equalize_hist(hsv[:, :, 2])
# img = color.hsv2rgb(hsv)
# central square crop
# min_side = min(img.shape[:-1])
# centre = img.shape[0] // 2, img.shape[1] // 2
# img = img[centre[0] - min_side // 2:centre[0] + min_side // 2,
# centre[1] - min_side // 2:centre[1] + min_side // 2,
# :]
# rescale to standard size
# img = transform.resize(img, (args.img_size, args.img_size))
# roll color axis to axis 0
# img = np.rollaxis(img, -1)
# img = img.transpose([2,0,1])
# img = img.transpose([2,0,1])
# return img
# return preprocess_img
def train(args=None):
args = parser.parse_args()
img_shape = ( args.img_size, args.img_size, args.img_channels) # blame theano
now_iso = datetime.now().strftime('%Y-%m-%dT%H:%M:%S%z')
#(x_train, y_train), (x_test,y_test) = cifar10.load_data()
# We then scale the variable-sized images to 224x224
# We augment .. by applying random lateral inversions and rotations.
train_datagen = ImageDataGenerator(
rescale=1. / 255,
rotation_range=30,
# contrast_stretching=True,
# adaptive_equalization=True,
histogram_equalization=True,
# featurewise_center=True,
# samplewise_center=True,
# featurewise_std_normalization=True,
# samplewise_std_normalization=True,
# channel_shift_range=0.2,
# brightness_range=[-0.3, 0.3],
width_shift_range=0.2,
height_shift_range=0.2,
zoom_range=0.3,
horizontal_flip=True,
# preprocessing_function= get_random_eraser(v_l=0, v_h=1, pixel_level=True)
# preprocessing_function=preprocess_img()
)
#train_datagen.fit(x_train)
train_generator = train_datagen.flow_from_directory(
'data/train/',
shuffle=True,
target_size=(args.img_size, args.img_size),
# save_to_dir='data/AUG_ELBOW_HIST',
class_mode='binary',
# color_mode='grayscale',
# interpolation='bicubic',
batch_size=args.batch_size, )
val_datagen = ImageDataGenerator(rescale=1. / 255,
# contrast_stretching=True
# ,adaptive_equalization=True
# histogram_equalization=True
)
val_generator = val_datagen.flow_from_directory(
'data/val/',
shuffle=True, # otherwise we get distorted batch-wise metrics
class_mode='binary',
# color_mode='grayscale',
target_size=(args.img_size, args.img_size),
batch_size=args.batch_size, )
# val_datagen.fit(x_train)
classes = len(train_generator.class_indices)
assert classes > 0
assert classes is len(val_generator.class_indices)
n_of_train_samples = train_generator.samples
n_of_val_samples = val_generator.samples
# Architectures
base_model = InceptionResNetV2(input_shape=img_shape, weights='imagenet', include_top=False)
x = base_model.output # Recast classification layer
#x = Flatten()(x) # Uncomment for Resnet based model
# x = GlobalAveragePooling2D(name='predictions_avg_pool')(x) # comment for RESNET models
x = WildcatPool2d()(x)
# n_classes; softmax for multi-class, sigmoid for binary
x = Dense(args.classes, activation='sigmoid', name='predictions')(x)
model = Model(inputs=base_model.input, outputs=x)
# model = multi_gpu_model(model, gpus=2)
# checkpoints
#
checkpoint = ModelCheckpoint(filepath='./models/InceptionResNetV2_499_NEW_HIST_WC_1.hdf5', verbose=1, save_best_only=True)
early_stop = EarlyStopping(patience=args.early_stop)
tensorboard = TensorBoard(log_dir='./logs/InceptionResNetV2_499_NEW_HIST_WC_1/{}/'.format(now_iso))
reduce_lr = ReduceLROnPlateau(factor=0.03, cooldown=0, patience=args.lr_wait, min_lr=0.1e-6)
callbacks = [checkpoint, tensorboard, reduce_lr]
# Calculate class weights
weights = class_weight.compute_class_weight('balanced', np.unique(train_generator.classes), train_generator.classes)
weights = {0: weights[0], 1: weights[1]}
# for layer in base_model.layers:
# layer.set_trainable = False
#print(model.summary())
#for i, layer in enumerate(base_model.layers):
# print(i, layer.name)
if args.resume:
model.load_weights(args.resume)
for layer in model.layers:
layer.set_trainable = True
#if TRAIN_FULL:
# print("=> retrain all layers of network")
# for layer in model.layers:
# set_trainable = True
#else:
# print("=> retraining only bottleneck and fc layers")
# import pdb
# pdb.set_trace()
# set_trainable = False
# for layer in base_model.layers:
# if "block12" in layer.name: # what block do we want to start unfreezing
# set_trainable = True
# if set_trainable:
# layer.trainable = True
# else:
# layer.trainable = False
# The network is trained end-to-end using Adam with default parameters
model.compile(
optimizer=Adam(lr=args.lr, decay=args.decay),
# optimizer=SGD(lr=args.lr, decay=args.decay,momentum=args.momentum, nesterov=True),
loss=binary_crossentropy,
# loss=kappa_error,
metrics=['accuracy', binary_accuracy], )
model_out = model.fit_generator(
train_generator,
steps_per_epoch=n_of_train_samples // args.batch_size,
epochs=args.epochs,
validation_data=val_generator,
validation_steps=n_of_val_samples // args.batch_size,
class_weight=weights,
workers=args.workers,
use_multiprocessing=True,
callbacks=callbacks)
if __name__ == '__main__':
train()