sbb_pixelwise_segmentation/train.py

import os
import sys
import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
import keras , warnings
from keras.optimizers import *
from sacred import Experiment
from models import *
from utils import *
from metrics import *
from keras.models import load_model

def configuration():
    keras.backend.clear_session()
    tf.reset_default_graph()
    warnings.filterwarnings('ignore')
    
    os.environ['CUDA_DEVICE_ORDER']='PCI_BUS_ID'
    config = tf.ConfigProto(log_device_placement=False, allow_soft_placement=True)
    
    
    config.gpu_options.allow_growth = True
    config.gpu_options.per_process_gpu_memory_fraction=0.95#0.95
    config.gpu_options.visible_device_list="0"
    set_session(tf.Session(config=config))

def get_dirs_or_files(input_data):
    if os.path.isdir(input_data):
        image_input, labels_input = os.path.join(input_data, 'images/'), os.path.join(input_data, 'labels/')
        # Check if training dir exists
        assert os.path.isdir(image_input), "{} is not a directory".format(image_input)
        assert os.path.isdir(labels_input), "{} is not a directory".format(labels_input)
    return image_input, labels_input

ex = Experiment()

@ex.config
def config_params():
    n_classes=None # Number of classes. If your case study is binary case the set it to 2 and otherwise give your number of cases.
    n_epochs=1
    input_height=224*1
    input_width=224*1 
    weight_decay=1e-6 # Weight decay of l2 regularization of model layers.
    n_batch=1 # Number of batches at each iteration.
    learning_rate=1e-4
    patches=False # Make patches of image in order to use all information of image. In the case of page
    # extraction this should be set to false since model should see all image.
    augmentation=False
    flip_aug=False # Flip image (augmentation).
    blur_aug=False # Blur patches of image (augmentation). 
    scaling=False # Scaling of patches (augmentation) will be imposed if this set to true.
    binarization=False # Otsu thresholding. Used for augmentation in the case of binary case like textline prediction. For multicases should not be applied.
    dir_train=None # Directory of training dataset (sub-folders should be named images and labels).
    dir_eval=None # Directory of validation dataset (sub-folders should be named images and labels).
    dir_output=None # Directory of output where the model should be saved.
    pretraining=False # Set true to load pretrained weights of resnet50 encoder.
    scaling_bluring=False
    scaling_binarization=False
    scaling_flip=False
    thetha=[10,-10]
    blur_k=['blur','guass','median'] # Used in order to blur image. Used for augmentation.
    scales= [ 0.5, 2 ] # Scale patches with these scales. Used for augmentation.
    flip_index=[0,1,-1] # Flip image. Used for augmentation.


@ex.automain
def run(n_classes,n_epochs,input_height,
        input_width,weight_decay,
        n_batch,patches,augmentation,flip_aug
        ,blur_aug,scaling, binarization,
        blur_k,scales,dir_train,
        scaling_bluring,scaling_binarization,rotation,
        rotation_not_90,thetha,scaling_flip,
        flip_index,dir_eval ,dir_output,pretraining,learning_rate):
    
    data_is_provided = False
    
    if data_is_provided:
        dir_train_flowing=os.path.join(dir_output,'train')
        dir_eval_flowing=os.path.join(dir_output,'eval')
        
        dir_flow_train_imgs=os.path.join(dir_train_flowing,'images')
        dir_flow_train_labels=os.path.join(dir_train_flowing,'labels')
        
        dir_flow_eval_imgs=os.path.join(dir_eval_flowing,'images')
        dir_flow_eval_labels=os.path.join(dir_eval_flowing,'labels')
        
        configuration()
        
    else:
        dir_img,dir_seg=get_dirs_or_files(dir_train)
        dir_img_val,dir_seg_val=get_dirs_or_files(dir_eval)
        
        # make first a directory in output for both training and evaluations in order to flow data from these directories.
        dir_train_flowing=os.path.join(dir_output,'train')
        dir_eval_flowing=os.path.join(dir_output,'eval')
        
        dir_flow_train_imgs=os.path.join(dir_train_flowing,'images/')
        dir_flow_train_labels=os.path.join(dir_train_flowing,'labels/')
        
        dir_flow_eval_imgs=os.path.join(dir_eval_flowing,'images/')
        dir_flow_eval_labels=os.path.join(dir_eval_flowing,'labels/')
        
        if os.path.isdir(dir_train_flowing):
            os.system('rm -rf '+dir_train_flowing)
            os.makedirs(dir_train_flowing)
        else:
            os.makedirs(dir_train_flowing)
            
        if os.path.isdir(dir_eval_flowing):
            os.system('rm -rf '+dir_eval_flowing)
            os.makedirs(dir_eval_flowing)
        else:
            os.makedirs(dir_eval_flowing)
            

        os.mkdir(dir_flow_train_imgs)
        os.mkdir(dir_flow_train_labels)
        
        os.mkdir(dir_flow_eval_imgs)
        os.mkdir(dir_flow_eval_labels)
        
        
        #set the gpu configuration
        configuration()


        #writing patches into a sub-folder in order to be flowed from directory. 
        provide_patches(dir_img,dir_seg,dir_flow_train_imgs,
                        dir_flow_train_labels,
                        input_height,input_width,blur_k,blur_aug,
                        flip_aug,binarization,scaling,scales,flip_index,
                        scaling_bluring,scaling_binarization,rotation,
                        rotation_not_90,thetha,scaling_flip,
                        augmentation=augmentation,patches=patches)
        
        provide_patches(dir_img_val,dir_seg_val,dir_flow_eval_imgs,
                        dir_flow_eval_labels,
                        input_height,input_width,blur_k,blur_aug,
                        flip_aug,binarization,scaling,scales,flip_index,
                        scaling_bluring,scaling_binarization,rotation,
                        rotation_not_90,thetha,scaling_flip,
                        augmentation=False,patches=patches)
        

    continue_train = False

    if continue_train:
        model_dir_start = '/home/vahid/Documents/struktur_full_data/output_multi/model_0.h5'
        model = load_model (model_dir_start, compile = True, custom_objects={'soft_dice_loss': soft_dice_loss})
        index_start = 1
    else:
        #get our model.
        index_start = 0
        model = resnet50_unet(n_classes,  input_height, input_width,weight_decay,pretraining)
    
    #if you want to see the model structure just uncomment model summary.
    #model.summary()
    

    
    #model.compile(loss='categorical_crossentropy',
                        #optimizer = Adam(lr=learning_rate),metrics=['accuracy'])
                        
    model.compile(loss=soft_dice_loss,
                        optimizer = Adam(lr=learning_rate),metrics=['accuracy'])
    
    #generating train and evaluation data
    train_gen = data_gen(dir_flow_train_imgs,dir_flow_train_labels, batch_size =  n_batch,
                         input_height=input_height, input_width=input_width,n_classes=n_classes  )
    val_gen = data_gen(dir_flow_eval_imgs,dir_flow_eval_labels, batch_size =  n_batch,
                         input_height=input_height, input_width=input_width,n_classes=n_classes  )
    
    for i in range(index_start, n_epochs+index_start):
        model.fit_generator(
            train_gen,
            steps_per_epoch=int(len(os.listdir(dir_flow_train_imgs))/n_batch)-1,
            validation_data=val_gen,
            validation_steps=1,
            epochs=1)
        model.save(dir_output+'/'+'model_'+str(i)+'.h5')
    

    os.system('rm -rf '+dir_train_flowing)
    os.system('rm -rf '+dir_eval_flowing)

    #model.save(dir_output+'/'+'model'+'.h5')