unet.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Apr 23 10:01:27 2018

@author: shawn
"""

import tensorflow as tf
import tensorlayer as tl
from tensorlayer.layers import *
import numpy as np
import BatchDatasetReader as BDR
import read_Data_list as RDL
import sys
import time
from sklearn.metrics import roc_auc_score
from sklearn.metrics import auc
from tfrecords_util import read_and_decode

#path variable
logs_dir = 'logs/'
#data_dir = 'Data/'

#basic constant variable
IMG_SIZE = 512
num_of_classes = 2
print_freq = 10

#training constant variable
MAX_EPOCH = int(3000+1)
batch_size = 3
test_batchsize = 3
train_nbr = 54
test_nbr = 27
step_every_epoch = int(train_nbr/batch_size)
test_every_epoch = int(test_nbr/test_batchsize)
learning_rate = tf.Variable(1e-4, dtype=tf.float32)

#the parameters of aupr
range_threshold = [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1]

#flags parameters
FLAGS = tf.flags.FLAGS
tf.flags.DEFINE_string('mode', "train", "Mode train/ test/ visualize")

data_dir = "Data/"
class Unet:
    def __init__(self, img_rows = IMG_SIZE, img_cols = IMG_SIZE):
        self.img_rows = img_rows
        self.img_cols = img_cols
#    def load_data_util(self):
#        image_options = {'resize': False, 'resize_size': IMG_SIZE} #resize all your images
#        train_records, valid_records = RDL.read_dataset(data_dir) #get read lists
#        train_dataset_reader = BDR.BatchDatset(train_records, image_options)
#        validation_dataset_reader = BDR.BatchDatset(valid_records, image_options)
#        return train_dataset_reader,validation_dataset_reader
    def model(self, image, is_train=True, reuse=False):
        with tf.variable_scope("model", reuse=reuse):
            tl.layers.set_name_reuse(reuse)
            W_init = tf.contrib.layers.xavier_initializer()
            input_image = tl.layers.InputLayer(image, name='input_layer') #input image
        
            conv2d_1 = tl.layers.Conv2d(input_image, 64, (3, 3), (1, 1),
                                        padding='SAME', W_init=W_init, name='conv_1')
            BN1 = tl.layers.BatchNormLayer(conv2d_1, act=tf.nn.relu, is_train=is_train, name='BN1')
            conv2d_2 = tl.layers.Conv2d(BN1, 64, (3, 3), (1, 1),
                                        padding='SAME', W_init=W_init, name='conv_2')
            BN2 = tl.layers.BatchNormLayer(conv2d_2, act=tf.nn.relu, is_train=is_train, name='BN2')
            pool_1 = tl.layers.MaxPool2d(BN2, (2, 2), (2,2), 
                                         padding='SAME', name='maxpool_1')
            conv2d_3 = tl.layers.Conv2d(pool_1, 128, (3,3), (1,1), 
                                        padding='SAME', W_init=W_init, name='conv_3')
            BN3 = tl.layers.BatchNormLayer(conv2d_3, act=tf.nn.relu, is_train=is_train, name='BN3')
            conv2d_4 = tl.layers.Conv2d(BN3, 128, (3,3), (1,1),
                                        padding='SAME', W_init=W_init, name='conv_4')
            BN4 = tl.layers.BatchNormLayer(conv2d_4, act=tf.nn.relu, is_train=is_train, name='BN4')
            pool_2 = tl.layers.MaxPool2d(BN4, (2,2), (2,2),
                                         padding='SAME', name='maxpool_2')
            conv2d_5 = tl.layers.Conv2d(pool_2, 256, (3,3), (1,1), 
                                        padding='SAME', W_init=W_init, name='conv_5')
            BN5 = tl.layers.BatchNormLayer(conv2d_5, act=tf.nn.relu, is_train=is_train, name='BN5')
            conv2d_6 = tl.layers.Conv2d(BN5, 256, (3,3), (1,1), 
                                        padding='SAME', W_init=W_init, name='conv_6')
            BN6 = tl.layers.BatchNormLayer(conv2d_6, act=tf.nn.relu, is_train=is_train, name='BN6')
            pool_3 = tl.layers.MaxPool2d(BN6, (2,2), (2,2),
                                         padding='SAME', name='maxpool_3')
            conv2d_7 = tl.layers.Conv2d(pool_3, 512, (3,3), (1,1), 
                                        padding='SAME', W_init=W_init, name='conv_7')
            BN7 = tl.layers.BatchNormLayer(conv2d_7, act=tf.nn.relu, is_train=is_train, name='BN7')
            conv2d_8 = tl.layers.Conv2d(BN7, 512, (3,3), (1,1), 
                                        padding='SAME', W_init=W_init, name='conv_8')
            BN8 = tl.layers.BatchNormLayer(conv2d_8, act=tf.nn.relu, is_train=is_train, name='BN8')
            #dropout_1 = tl.layers.DropoutLayer(conv2d_8, keep= 0.5, is_fix=True, is_train=is_train, name = 'drop_1')
            pool_4 = tl.layers.MaxPool2d(BN8, (2,2), (2,2),
                                         padding='SAME', name='maxpool_4')
            conv2d_9 = tl.layers.Conv2d(pool_4, 1024, (3,3), (1,1), 
                                        padding='SAME', W_init=W_init, name='conv_9')
            BN9 = tl.layers.BatchNormLayer(conv2d_9, act=tf.nn.relu, is_train=is_train, name='BN9')
            conv2d_10 = tl.layers.Conv2d(BN9, 1024, (3,3), (1,1),
                                         padding='SAME', W_init=W_init, name='conv_10')
            BN10 = tl.layers.BatchNormLayer(conv2d_10, act=tf.nn.relu, is_train=is_train, name='BN10')
            #dropout_2 = tl.layers.DropoutLayer(conv2d_10, keep= 0.5, is_fix=True, is_train=is_train, name = 'drop_2')
            upsampling_1 = tl.layers.UpSampling2dLayer(BN10, (2,2), name='upsample2d_1')
            conv2d_11 = tl.layers.Conv2d(upsampling_1, 512, (3,3), (1,1),
                                         padding='SAME', W_init=W_init, name='conv_11')
            BN11 = tl.layers.BatchNormLayer(conv2d_11, act=tf.nn.relu, is_train=is_train, name='BN11')
            concat_1 = tl.layers.ConcatLayer([BN8, BN11], 3, name ='concat_1')
            conv2d_12 = tl.layers.Conv2d(concat_1, 512, (3,3), (1,1), 
                                         padding='SAME', W_init=W_init, name='conv_12')
            BN12 = tl.layers.BatchNormLayer(conv2d_12, act=tf.nn.relu, is_train=is_train, name='BN12')
            conv2d_13 = tl.layers.Conv2d(BN12, 512, (3,3), (1,1),
                                     padding='SAME', W_init=W_init, name='conv_13')
            BN13 = tl.layers.BatchNormLayer(conv2d_13, act=tf.nn.relu, is_train=is_train, name='BN13')
            upsampling_2 = tl.layers.UpSampling2dLayer(BN13, (2,2), name='upsample2d_2')
            conv2d_14 = tl.layers.Conv2d(upsampling_2, 256, (3,3), (1,1),
                                         padding='SAME', W_init=W_init, name='conv_14')
            BN14 = tl.layers.BatchNormLayer(conv2d_14, act=tf.nn.relu, is_train=is_train, name='B14')
            concat_2 = tl.layers.ConcatLayer([BN14,BN6], 3, name='concat_2')
            conv2d_15 = tl.layers.Conv2d(concat_2, 256, (3,3), (1,1), 
                                         padding='SAME', W_init=W_init, name='conv_15')
            BN15 = tl.layers.BatchNormLayer(conv2d_15, act=tf.nn.relu, is_train=is_train, name='BN15')
            conv2d_16 = tl.layers.Conv2d(BN15, 256, (3,3), (1,1),
                                         padding='SAME', W_init=W_init, name='conv_16')
            BN16 = tl.layers.BatchNormLayer(conv2d_16, act=tf.nn.relu, is_train=is_train, name='BN16')
            upsampling_3 = tl.layers.UpSampling2dLayer(BN16, (2,2), name='upsample2d_3')
            conv2d_17 = tl.layers.Conv2d(upsampling_3, 128, (3,3), (1,1),
                                         padding='SAME', W_init=W_init, name='conv_17')
            BN17 = tl.layers.BatchNormLayer(conv2d_17, act=tf.nn.relu, is_train=is_train, name='BN17')
            concat_3 = tl.layers.ConcatLayer([BN17,BN4], 3, name='concat_3')
            conv2d_18 = tl.layers.Conv2d(concat_3, 128, (3,3), (1,1), 
                                         padding='SAME', W_init=W_init, name='conv_18')
            BN18 = tl.layers.BatchNormLayer(conv2d_18, act=tf.nn.relu, is_train=is_train, name='BN18')
            conv2d_19 = tl.layers.Conv2d(BN18, 128, (3,3), (1,1),
                                         padding='SAME', W_init=W_init, name='conv_19')
            BN19 = tl.layers.BatchNormLayer(conv2d_19, act=tf.nn.relu, is_train=is_train, name='BN19')
            upsampling_4 = tl.layers.UpSampling2dLayer(BN19, (2,2), name='upsample2d_4')
            conv2d_20 = tl.layers.Conv2d(upsampling_4, 64, (3,3), (1,1),
                                         padding='SAME', W_init=W_init, name='conv_20')
            BN20 = tl.layers.BatchNormLayer(conv2d_20, act=tf.nn.relu, is_train=is_train, name='BN20')
            concat_4 = tl.layers.ConcatLayer([BN20,BN2], 3, name='concat_4')
            conv2d_21 = tl.layers.Conv2d(concat_4, 64, (3,3), (1,1),
                                         padding='SAME', W_init=W_init, name='conv_21')
            BN21 = tl.layers.BatchNormLayer(conv2d_21, act=tf.nn.relu, is_train=is_train, name='BN21')
            conv2d_22 = tl.layers.Conv2d(BN21, 32, (3,3), (1,1),
                                         padding='SAME', W_init=W_init, name='conv_22')
            BN22 = tl.layers.BatchNormLayer(conv2d_22, act=tf.nn.relu, is_train=is_train, name='BN22')
            conv2d_23 = tl.layers.Conv2d(BN22, num_of_classes, (3,3), (1,1),
                                         padding='SAME', W_init=W_init, name='conv_23')
            #maybe conv2d_23 should not be activation!
            y = conv2d_23.outputs #transfer tl object to logits tensor
            pred = tf.argmax(y, dimension=3, name="prediction")

        return pred,y,conv2d_23
    def loss(self, logits, annotation):
        loss = tf.reduce_mean((tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits,
                                                                          labels=tf.squeeze(annotation, squeeze_dims=[3]),
                                                                        name="entropy")))
        
        L2 = 0
        for p in tl.layers.get_variables_with_name('relu/W', True, True):
            L2 += tf.contrib.layers.l2_regularizer(0.004)(p)
        cost = loss + L2
        return cost
    def train(self, loss):
        #If use tf.nn.sparse_softmax_cross_entropy_with_logits ,
        #maybe loss will be NAN,because without clip
        #annotation = tf.cast(annotation,dtype = tf.float32)
        #prob = tf.nn.softmax(logits)
        #loss = -tf.reduce_mean(annotation*tf.log(tf.clip_by_value(prob,1e-11,1.0)))
        optimizer = tf.train.AdamOptimizer(learning_rate)
        var_list = tf.trainable_variables()
        grads = optimizer.compute_gradients(loss, var_list=var_list)
        train_op = optimizer.apply_gradients(grads)
        return train_op
    
#AUPR score
def computeConfMatElements(thresholded_proba_map, ground_truth):
    P = np.count_nonzero(ground_truth)
    TP = np.count_nonzero(thresholded_proba_map*ground_truth)
    FP = np.count_nonzero(thresholded_proba_map - (thresholded_proba_map*ground_truth))    
                                 
    return P,TP,FP
    
def computeAUPR(proba_map, ground_truth, threshold_list):
    proba_map = proba_map.astype(np.float32)
    precision_list_treshold = []
    recall_list_treshold = []
    #loop over thresholds
    for threshold in threshold_list:
        precision_list_proba_map = []
        recall_list_proba_map = []
        #loop over proba map
        for i in range(len(proba_map)): #batch_size
            #threshold the proba map
            thresholded_proba_map = np.zeros(np.shape(proba_map[i]))
            thresholded_proba_map[proba_map[i] >= threshold] = 1
            #print(np.shape(thresholded_proba_map)) #(400,640)
                   
            #compute P, TP, and FP for this threshold and this proba map
            P,TP,FP = computeConfMatElements(thresholded_proba_map, ground_truth[i])       
            
            #check that ground truth contains at least one positive
            if (P > 0 and (TP+FP) > 0) :
                precision_list_proba_map.append(TP*1./(TP+FP))
                recall_list_proba_map.append(TP*1./P)
            else:
                precision_list_proba_map.append(0)
                recall_list_proba_map.append(0)
            
        
        #average sensitivity and FP over the proba map, for a given threshold
        precision_list_treshold.append(np.mean(precision_list_proba_map))
        recall_list_treshold.append(np.mean(recall_list_proba_map))    
    
    return auc(recall_list_treshold, precision_list_treshold)

def main(argv=None):    
    myUnet = Unet()
    x_train_, y_train_ = read_and_decode("train.tfrecords")
    x_test_, y_test_ = read_and_decode("test.tfrecords")

    x_train_batch, y_train_batch = tf.train.shuffle_batch(
        [x_train_, y_train_], batch_size=batch_size, capacity=200, min_after_dequeue=100, num_threads=4
    )  
    # set the number of threads here
    # for testing, uses batch instead of shuffle_batch
    x_test_batch, y_test_batch = tf.train.batch(
        [x_test_, y_test_], batch_size=batch_size, capacity=100, num_threads=4
    )

    # define inferences
    train_pred, train_logits, train_tlnetwork = myUnet.model(x_train_batch, is_train=True, reuse=False)
    train_positive_prob = tf.nn.softmax(train_logits)[:, :, :, 1]
    train_loss_op = myUnet.loss(train_logits, y_train_batch)
    train_op = myUnet.train(train_loss_op)
    
    test_pred, test_logits, test_tlnetwork = myUnet.model(x_test_batch, is_train=False, reuse=True)
    test_positive_prob = tf.nn.softmax(test_logits)[:, :, :, 1]
    test_loss_op= myUnet.loss(test_logits, y_test_batch)
    
    lr_assign_op = tf.assign(learning_rate, learning_rate / 10) #learning_rate decay
    
    #only visualize the test images
    #first lighten the annotation images
    visual_annotation = tf.where(tf.equal(y_test_batch,1), y_test_batch+254, y_test_batch)
    visual_pred = tf.expand_dims(tf.where(tf.equal(test_pred,1), test_pred+254, test_pred), dim=3)
    tf.summary.image("input_image", x_test_batch, max_outputs=2)
    tf.summary.image("ground_truth", tf.cast(visual_annotation, tf.uint8), max_outputs=2)
    tf.summary.image("pred_annotation", tf.cast(visual_pred, tf.uint8), max_outputs=2)
    
    print("Setting up summary op...")
    test_summary_op = tf.summary.merge_all()
    
#    if FLAGS.mode == 'train':
#        train_dataset_reader,validation_dataset_reader = myUnet.load_data_util()
    sess = tf.Session()
    
    print("Setting up Saver...")
    saver = tf.train.Saver(max_to_keep=2)
    summary_writer = tf.summary.FileWriter(logs_dir, sess.graph)
    
    tl.layers.initialize_global_variables(sess)
    sess.run(tf.global_variables_initializer())
    
    ckpt = tf.train.get_checkpoint_state(logs_dir)#if model has been trained,restore it
    if ckpt and ckpt.model_checkpoint_path:
        saver.restore(sess, ckpt.model_checkpoint_path)   
        print("Model restored...")
    
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)
    
    for epo in range(MAX_EPOCH):
        start_time = time.time()
        train_loss, test_loss, train_aupr, test_aupr, train_auc, test_auc= 0, 0, 0, 0, 0, 0
        
        for s in range(step_every_epoch):
#            train_images, train_annotations = train_dataset_reader.next_batch(batch_size,is_train=True)
#            feed_dict = {image: train_images, annotation: train_annotations}
            tra_positive_prob, train_err, _ = sess.run([train_positive_prob, train_loss_op,train_op])
            
            #compute auc score
            label_train_batch = sess.run(y_train_batch) #get array from tensor
            temp_train_annotations = (np.reshape(label_train_batch,-1)).astype(np.int32)
            temp_tra_positive_prob = np.reshape(tra_positive_prob,-1)
            train_sauc = roc_auc_score(temp_train_annotations, temp_tra_positive_prob)
            #compute aupr
            train_saupr = computeAUPR(tra_positive_prob ,np.squeeze(label_train_batch, axis=3), range_threshold)
            
            
            train_loss += train_err
            train_auc += train_sauc
            train_aupr += train_saupr
            
        if epo + 1 == 1 or (epo + 1) % print_freq == 0:
            train_loss = train_loss/step_every_epoch
            train_auc = train_auc/step_every_epoch
            train_aupr = train_aupr/step_every_epoch
            #visualize the training loss
            print("%d epoches %d took %fs" % (print_freq, epo, time.time() - start_time))
            print("   train loss: %f" % train_loss)
            print("   train auc: %f" % train_auc)
            print("   train aupr: %f" % train_aupr)
            
            train_summary = tf.Summary(value=[
                    tf.Summary.Value(tag="train_loss", simple_value=train_loss), 
                    tf.Summary.Value(tag="train_auc", simple_value=train_auc),
                    tf.Summary.Value(tag="train_aupr", simple_value=train_aupr)
                ])
            summary_writer.add_summary(train_summary, epo)
            
            for test_s in range(test_every_epoch):
                #get validation data
                #valid_images, valid_annotations = validation_dataset_reader.next_batch(test_batchsize, is_train=False)
                #visualize the validation loss
                #feed_dict= {image:valid_images,annotation:valid_annotations}
                valid_positive_prob, validation_err = sess.run([test_positive_prob, test_loss_op])
                #compute auc score
                label_test_batch = sess.run(y_test_batch)
                temp_valid_annotations = (np.reshape(label_test_batch,-1)).astype(np.int32)
                temp_valid_positive_prob = np.reshape(valid_positive_prob,-1)
                test_sauc = roc_auc_score(temp_valid_annotations, temp_valid_positive_prob)
                #compute test aupr
                test_saupr = computeAUPR(valid_positive_prob ,np.squeeze(label_test_batch, axis=3), range_threshold)
                
                test_loss += validation_err
                test_auc += test_sauc
                test_aupr += test_saupr
            test_loss = test_loss/test_every_epoch
            test_auc = test_auc/test_every_epoch
            test_aupr = test_aupr/test_every_epoch
            test_summary = tf.Summary(value=[
                    tf.Summary.Value(tag="test_loss", simple_value=test_loss), 
                    tf.Summary.Value(tag="test_auc", simple_value=test_auc),
                    tf.Summary.Value(tag="test_aupr", simple_value=test_aupr)
                ])
            summary_writer.add_summary(test_summary, epo)
            
            #visualize the test result(only visualize the last batchsize of this epoch)
            #feed_dict= {image:valid_images,annotation:valid_annotations}
            summary_str = sess.run(test_summary_op)
            summary_writer.add_summary(summary_str, epo)
            
            #summary_writer.add_summary(summary_str, epo)
            #tensorboard flush
            summary_writer.flush()
            sys.stdout.flush()
#        if epo == 1500 or epo == 2000:   
#            sess.run(lr_assign_op)
        if epo % 1500 == 0:
            saver.save(sess, logs_dir + "model.ckpt", epo)
            print('the %d epoch , the model has been saved successfully' %epo)
            sys.stdout.flush()
            
    coord.request_stop()
    coord.join(threads)
    sess.close()       
if __name__ == '__main__':
    tf.app.run()