UnMicst2.py

import numpy as np
from scipy import misc
import os
import logging

import tensorflow.compat.v1 as tf
from tensorflow import keras
from tensorflow.keras import layers
import shutil
import scipy.io as sio
import fnmatch, glob
import skimage.exposure as sk
import skimage.io
import argparse
import czifile
from nd2reader import ND2Reader
import tifffile
import sys
tf.disable_v2_behavior()
# sys.path.insert(0, 'C:\\Users\\Public\\Documents\\ImageScience')
from toolbox.imtools import *
from toolbox.ftools import *
from toolbox.PartitionOfImage import PI2D
from toolbox import GPUselect

def concat3(lst):
    return tf.concat(lst, 3)


class UNet2D:
    hp = None  # hyper-parameters
    nn = None  # network
    tfTraining = None  # if training or not (to handle batch norm)
    tfData = None  # data placeholder
    Session = None
    DatasetMean = 0
    DatasetStDev = 0

    def setupWithHP(hp):
        UNet2D.setup(hp['imSize'],
                     hp['nChannels'],
                     hp['nClasses'],
                     hp['nOut0'],
                     hp['featMapsFact'],
                     hp['downSampFact'],
                     hp['ks'],
                     hp['nExtraConvs'],
                     hp['stdDev0'],
                     hp['nLayers'],
                     hp['batchSize'])

    def setup(imSize, nChannels, nClasses, nOut0, featMapsFact, downSampFact, kernelSize, nExtraConvs, stdDev0,
              nDownSampLayers, batchSize):
        UNet2D.hp = {'imSize': imSize,
                     'nClasses': nClasses,
                     'nChannels': nChannels,
                     'nExtraConvs': nExtraConvs,
                     'nLayers': nDownSampLayers,
                     'featMapsFact': featMapsFact,
                     'downSampFact': downSampFact,
                     'ks': kernelSize,
                     'nOut0': nOut0,
                     'stdDev0': stdDev0,
                     'batchSize': batchSize}

        nOutX = [UNet2D.hp['nChannels'], UNet2D.hp['nOut0']]
        dsfX = []
        for i in range(UNet2D.hp['nLayers']):
            nOutX.append(nOutX[-1] * UNet2D.hp['featMapsFact'])
            dsfX.append(UNet2D.hp['downSampFact'])

        # --------------------------------------------------
        # downsampling layer
        # --------------------------------------------------

        with tf.name_scope('placeholders'):
            UNet2D.tfTraining = tf.placeholder(tf.bool, name='training')
            UNet2D.tfData = tf.placeholder("float", shape=[None, UNet2D.hp['imSize'], UNet2D.hp['imSize'],
                                                           UNet2D.hp['nChannels']], name='data')

        def down_samp_layer(data, index):
            regularizer = tf.keras.regularizers.l2(0.01)
            with tf.variable_scope('ld%d' % index):
                ldXWeights1 = tf.Variable(
                    tf.truncated_normal([UNet2D.hp['ks'], UNet2D.hp['ks'], nOutX[index], nOutX[index + 1]],
                                        stddev=stdDev0), name='kernelD%d' % index)
                # ldXWeights1 = tf.get_variable(
                # 	initializer=tf.contrib.layers.variance_scaling_initializer(mode='FAN_IN'),
                # 	shape=[UNet2D.hp['ks'], UNet2D.hp['ks'], nOutX[index], nOutX[index + 1]], name='kernelD%d' % index, regularizer=regularizer)
                ldXWeightsExtra = []
                for i in range(nExtraConvs):
                    ldXWeightsExtra.append(
                        tf.get_variable(initializer=tf.compat.v1.keras.initializers.VarianceScaling(mode='fan_in'),
                                        shape=[UNet2D.hp['ks'], UNet2D.hp['ks'], nOutX[index + 1], nOutX[index + 1]],
                                        name='kernelExtra%d' % i))
                    # ldXWeightsExtra.append(tf.Variable(
                    # 	tf.truncated_normal([UNet2D.hp['ks'], UNet2D.hp['ks'], nOutX[index + 1], nOutX[index + 1]],
                    # 						stddev=stdDev0), name='kernelExtra%d' % i))

                c00 = tf.nn.conv2d(data, ldXWeights1, strides=[1, 1, 1, 1], padding='SAME')
                for i in range(nExtraConvs):
                    c00 = tf.nn.conv2d(tf.nn.leaky_relu(c00), ldXWeightsExtra[i], strides=[1, 1, 1, 1], padding='SAME')

                ldXWeightsShortcut = tf.get_variable(initializer=tf.compat.v1.keras.initializers.VarianceScaling(mode='fan_in'),
                                                     shape=[UNet2D.hp['ks'], UNet2D.hp['ks'], nOutX[index],
                                                            nOutX[index + 1]],
                                                     name='shortcutWeights', regularizer=regularizer)
                # ldXWeightsShortcut = tf.Variable(
                # 	tf.truncated_normal([1, 1, nOutX[index], nOutX[index + 1]], stddev=stdDev0), name='shortcutWeights')
                shortcut = tf.nn.conv2d(data, ldXWeightsShortcut, strides=[1, 1, 1, 1], padding='SAME')

                bn = tf.nn.leaky_relu(tf.layers.batch_normalization(c00+shortcut, training=UNet2D.tfTraining))
                # bn = tf.layers.batch_normalization(tf.nn.leaky_relu(c00 + shortcut), training=UNet2D.tfTraining)
                bn = tf.layers.dropout(bn, 0.05 * index, training=UNet2D.tfTraining)
                return tf.nn.max_pool(bn, ksize=[1, dsfX[index], dsfX[index], 1],
                                      strides=[1, dsfX[index], dsfX[index], 1], padding='SAME', name='maxpool')

        # --------------------------------------------------
        # bottom layer
        # --------------------------------------------------

        with tf.variable_scope('lb'):
            regularizer = tf.keras.regularizers.l2(0.01)
            lbWeights1 = tf.get_variable(initializer=tf.compat.v1.keras.initializers.VarianceScaling(mode='fan_in'),
                                         shape=[UNet2D.hp['ks'], UNet2D.hp['ks'], nOutX[UNet2D.hp['nLayers']],
                                                nOutX[UNet2D.hp['nLayers'] + 1]],
                                         name='kernel1', regularizer=regularizer)
            # lbWeights1 = tf.Variable(tf.truncated_normal(
            # 	[UNet2D.hp['ks'], UNet2D.hp['ks'], nOutX[UNet2D.hp['nLayers']], nOutX[UNet2D.hp['nLayers'] + 1]],
            # 	stddev=stdDev0), name='kernel1')
            def lb(hidden):
                # lbn= tf.nn.leaky_relu(
                # 	tf.nn.conv2d(hidden, lbWeights1, strides=[1, 1, 1, 1], padding='SAME'), name='conv')
                lbn = tf.nn.leaky_relu(tf.layers.batch_normalization(
                    tf.nn.conv2d(hidden, lbWeights1, strides=[1, 1, 1, 1], padding='SAME'),
                    name='conv',training=UNet2D.tfTraining))
                return tf.layers.dropout(lbn, 0.3, training=UNet2D.tfTraining)

        # --------------------------------------------------
        # downsampling
        # --------------------------------------------------

        with tf.name_scope('downsampling'):
            dsX = []
            dsX.append(UNet2D.tfData)

            for i in range(UNet2D.hp['nLayers']):
                dsX.append(down_samp_layer(dsX[i], i))

            b = lb(dsX[UNet2D.hp['nLayers']])

        # --------------------------------------------------
        # upsampling layer
        # --------------------------------------------------

        def up_samp_layer(data, index):
            with tf.variable_scope('lu%d' % index):
                regularizer = tf.keras.regularizers.l2(0.005)
                luXWeights1 = tf.get_variable(
                    initializer=tf.compat.v1.keras.initializers.VarianceScaling(mode='fan_in'),
                    shape=[UNet2D.hp['ks'], UNet2D.hp['ks'], nOutX[index + 1], nOutX[index + 2]],
                    name='kernelU%d' % index,regularizer=regularizer)
                luXWeights2 = tf.get_variable(
                    initializer=tf.compat.v1.keras.initializers.VarianceScaling(mode='fan_in'),
                    shape=[UNet2D.hp['ks'], UNet2D.hp['ks'], nOutX[index] + nOutX[index + 1], nOutX[index + 1]],
                    name='kernel2',regularizer=regularizer)
                luXWeightsExtra = []
                for i in range(nExtraConvs):
                    luXWeightsExtra.append(tf.get_variable(
                        initializer=tf.compat.v1.keras.initializers.VarianceScaling(mode='fan_in'),
                        shape=[UNet2D.hp['ks'], UNet2D.hp['ks'], nOutX[index + 1], nOutX[index + 1]],
                        name='kernel2Extra%d' % i))
                # luXWeights1 = tf.Variable(
                # 	tf.truncated_normal([UNet2D.hp['ks'], UNet2D.hp['ks'], nOutX[index + 1], nOutX[index + 2]],
                # 						stddev=stdDev0), name='kernel1')
                # luXWeights2 = tf.Variable(tf.truncated_normal(
                # 	[UNet2D.hp['ks'], UNet2D.hp['ks'], nOutX[index] + nOutX[index + 1], nOutX[index + 1]],
                # 	stddev=stdDev0), name='kernel2')
                # luXWeightsExtra = []
                # for i in range(nExtraConvs):
                # 	luXWeightsExtra.append(tf.Variable(
                # 		tf.truncated_normal([UNet2D.hp['ks'], UNet2D.hp['ks'], nOutX[index + 1], nOutX[index + 1]],
                # 							stddev=stdDev0), name='kernel2Extra%d' % i))

                outSize = UNet2D.hp['imSize']
                for i in range(index):
                    outSize /= dsfX[i]
                outSize = int(outSize)

                outputShape = [UNet2D.hp['batchSize'], outSize, outSize, nOutX[index + 1]]
                us = tf.nn.leaky_relu(
                    tf.nn.conv2d_transpose(data, luXWeights1, outputShape, strides=[1, dsfX[index], dsfX[index], 1],
                                           padding='SAME'), name='conv1')
                cc = concat3([dsX[index], us])
                # cv = tf.nn.leaky_relu(tf.nn.conv2d(cc, luXWeights2, strides=[1, 1, 1, 1], padding='SAME'), name='conv2')
                cv = tf.nn.leaky_relu(
                    tf.layers.batch_normalization(tf.nn.conv2d(cc, luXWeights2, strides=[1, 1, 1, 1], padding='SAME'),
                                                  name='conv2',
                                                  training=UNet2D.tfTraining))
                for i in range(nExtraConvs):
                    cv = tf.nn.leaky_relu(tf.nn.conv2d(cv, luXWeightsExtra[i], strides=[1, 1, 1, 1], padding='SAME'),
                                          name='conv2Extra%d' % i)
                cv = tf.layers.dropout(cv, 0.25 - 0.05 * index, training=UNet2D.tfTraining)
                return cv

        # --------------------------------------------------
        # final (top) layer
        # --------------------------------------------------

        with tf.variable_scope('lt'):
            regularizer = tf.keras.regularizers.l2(0.005)
            # ltWeights1 = tf.Variable(tf.truncated_normal([1, 1, nOutX[1], nClasses], stddev=stdDev0), name='kernel')
            ltWeights1 = tf.get_variable(initializer=tf.compat.v1.keras.initializers.VarianceScaling(mode='fan_in'),
                                         shape=[1, 1, nOutX[1], nClasses],
                                         name='kernel',regularizer=regularizer)
            def lt(hidden):
                # return tf.nn.conv2d(hidden, ltWeights1, strides=[1, 1, 1, 1], padding='SAME', name='conv')
                return tf.layers.batch_normalization(
                    tf.nn.conv2d(hidden, ltWeights1, strides=[1, 1, 1, 1], padding='SAME', name='conv'),
                    training=UNet2D.tfTraining)
        # --------------------------------------------------
        # upsampling
        # --------------------------------------------------

        with tf.name_scope('upsampling'):
            usX = []
            usX.append(b)

            for i in range(UNet2D.hp['nLayers']):
                usX.append(up_samp_layer(usX[i], UNet2D.hp['nLayers'] - 1 - i))

            t = lt(usX[UNet2D.hp['nLayers']])

        sm = tf.nn.softmax(t, -1)
        UNet2D.nn = sm

    def train(imPath, validPath, testPath, logPath, modelPath, pmPath, nTrain, nValid, nTest, restoreVariables, nSteps, gpuIndex,
              testPMIndex):
        os.environ['CUDA_VISIBLE_DEVICES'] = '%d' % gpuIndex

        outLogPath = logPath
        trainWriterPath = pathjoin(logPath, 'Train')
        validWriterPath = pathjoin(logPath, 'Valid')
        outModelPath = pathjoin(modelPath, 'model.ckpt')
        outPMPath = pmPath

        batchSize = UNet2D.hp['batchSize']
        imSize = UNet2D.hp['imSize']
        nChannels = UNet2D.hp['nChannels']
        nClasses = UNet2D.hp['nClasses']

        # --------------------------------------------------
        # data
        # --------------------------------------------------

        Train = np.zeros((nTrain, imSize, imSize, nChannels))
        Valid = np.zeros((nValid, imSize, imSize, nChannels))
        Test = np.zeros((nTest, imSize, imSize, nChannels))
        LTrain = np.zeros((nTrain, imSize, imSize, nClasses))
        LValid = np.zeros((nValid, imSize, imSize, nClasses))
        LTest = np.zeros((nTest, imSize, imSize, nClasses))
        WTrain = np.zeros((nTrain, imSize, imSize, nClasses))
        WValid = np.zeros((nValid, imSize, imSize, nClasses))
        WTest = np.zeros((nTest, imSize, imSize, nClasses))

        print('loading data, computing mean / st dev')
        if not os.path.exists(modelPath):
            os.makedirs(modelPath)
        # if restoreVariables:
        # 	datasetMean = loadData(pathjoin(modelPath,'datasetMean.data'))
        # 	datasetStDev = loadData(pathjoin(modelPath,'datasetStDev.data'))
        # else:
        datasetMean = 0.19
        datasetStDev = 0.17
        bgWeight = 1
        contourWeight = 2
        nucleiWeight = 5
        intersectWeight = 10

        # for iSample in range(nTrain + nValid + nTest):
        #     I = im2double(tifread('%s/I%05d_Img.tif' % (imPath, iSample)))
        #     datasetMean += np.mean(I)
        #     datasetStDev += np.std(I)
        # datasetMean /= (nTrain + nValid + nTest)
        # datasetStDev /= (nTrain + nValid + nTest)
        saveData(datasetMean, pathjoin(modelPath, 'datasetMean.data'))
        saveData(datasetStDev, pathjoin(modelPath, 'datasetStDev.data'))

        perm = np.arange(nTrain)
        np.random.shuffle(perm)

        for iSample in range(0, nTrain):
            path = '%s/I%05d_Img.tif' % (imPath, perm[iSample])

            # for iChan in range(nChannels):
            im = im2double(skio.imread(path, img_num=0))
            Train[iSample, :, :, 0] = (im - datasetMean) / datasetStDev
            path = '%s/I%05d_Ant.tif' % (imPath, perm[iSample])
            im = tifread(path)
            path = '%s/I%05d_wt.tif' % (imPath, perm[iSample])
            W = tifread(path)
            for i in range(nClasses):
                LTrain[iSample, :, :, i] = (im == i + 1)
                if i == 1:
                    WTrain[iSample, :, :, i] = (W * intersectWeight) + contourWeight
                elif i == 2:
                    WTrain[iSample, :, :, i] = (W * 0) + nucleiWeight
                else:
                    WTrain[iSample, :, :, i] = (W * 0) + bgWeight

        permV = np.arange(nValid)
        np.random.shuffle(permV)
        for iSample in range(0, nValid):
            path = '%s/I%05d_Img.tif' % (validPath, permV[iSample])
            # im = im2double(tifread(path))
            # for iChan in range(nChannels):
            im = im2double(skio.imread(path, img_num=0))
            Valid[iSample, :, :, 0] = (im - datasetMean) / datasetStDev
            path = '%s/I%05d_Ant.tif' % (validPath, permV[iSample])
            im = tifread(path)
            path = '%s/I%05d_wt.tif' % (validPath,  permV[iSample])
            W = tifread(path)
            for i in range(nClasses):
                LValid[iSample, :, :, i] = (im == i + 1)
                if i == 1:
                    WValid[iSample, :, :, i] = (W * intersectWeight) + contourWeight
                elif i == 2:
                    WValid[iSample, :, :, i] = (W * 0) + nucleiWeight
                else:
                    WValid[iSample, :, :, i] = (W * 0) + bgWeight

        for iSample in range(0, nTest):
            path = '%s/I%05d_Img.tif' % (testPath, iSample)
            # for iChan in range(nChannels):
            im = im2double(skio.imread(path, img_num=0))
            Test[iSample, :, :, 0] = (im - datasetMean) / datasetStDev
            path = '%s/I%05d_Ant.tif' % (testPath, iSample)
            im = tifread(path)
            path = '%s/I%05d_wt.tif' % (testPath, iSample)
            W = tifread(path)
            for i in range(nClasses):
                LTest[iSample, :, :, i] = (im == i + 1)
                if i == 1:
                    WTest[iSample, :, :, i] = (W * intersectWeight) + contourWeight
                elif i == 2:
                    WTest[iSample, :, :, i] = (W * 0) + nucleiWeight
                else:
                    WTest[iSample, :, :, i] = (W * 0) + bgWeight

        # --------------------------------------------------
        # optimization
        # --------------------------------------------------

        tfLabels = tf.placeholder("float", shape=[None, imSize, imSize, nClasses], name='labels')
        tfWeights = tf.placeholder("float", shape=[None, imSize, imSize, nClasses], name='weights')
        globalStep = tf.Variable(0, trainable=False)
        learningRate0 = 0.00006
        decaySteps = 4000
        decayRate = 0.99
        learningRate = tf.train.exponential_decay(learningRate0, globalStep, decaySteps, decayRate, staircase=True)

        with tf.name_scope('optim'):
            l2_loss = tf.compat.v1.losses.get_regularization_loss()
            loss = tf.reduce_mean(-tf.reduce_sum(tf.multiply(tf.cast(tfWeights, tf.float32),
                                                             tf.multiply(tf.cast(tfLabels, tf.float32),
                                                                         tf.log(UNet2D.nn))), 3)) + l2_loss

            updateOps = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
            # optimizer = tf.train.MomentumOptimizer(1e-3,0.9)
            # optimizer = tf.train.MomentumOptimizer(learningRate,0.99)
            optimizer = tf.train.AdamOptimizer(learning_rate=learningRate)
            with tf.control_dependencies(updateOps):
                optOp = optimizer.minimize(loss, global_step=globalStep)

        # for g, v in gradients:
        # 	tf.summary.histogram(v.name, v)
        # 	tf.summary.histogram(v.name + '_grad', g)

        with tf.name_scope('eval'):
            error = []
            for iClass in range(nClasses):
                labels0 = tf.reshape(tf.to_int32(tf.slice(tfLabels, [0, 0, 0, iClass], [-1, -1, -1, 1])),
                                     [batchSize, imSize, imSize])
                predict0 = tf.reshape(tf.to_int32(tf.equal(tf.argmax(UNet2D.nn, 3), iClass)),
                                      [batchSize, imSize, imSize])
                correct = tf.multiply(labels0, predict0)
                nCorrect0 = tf.reduce_sum(correct)
                nLabels0 = tf.reduce_sum(labels0)
                error.append(1 - tf.to_float(nCorrect0) / tf.to_float(nLabels0))
            errors = tf.tuple(error)

        # --------------------------------------------------
        # inspection
        # --------------------------------------------------

        with tf.name_scope('scalars'):
            tf.summary.scalar('avg_cross_entropy', loss)
            for iClass in range(nClasses):
                tf.summary.scalar('avg_pixel_error_%d' % iClass, error[iClass])
            tf.summary.scalar('learning_rate', learningRate)
        with tf.name_scope('images'):
            split0 = tf.slice(UNet2D.nn, [0, 0, 0, 1], [-1, -1, -1, 1])
            split1 = tf.slice(UNet2D.tfData, [0, 0, 0, 0], [-1, -1, -1, 1])

            planeImN = tf.div(tf.subtract(split1, tf.reduce_min(split1, axis=(1, 2), keep_dims=True)),
                              tf.subtract(tf.reduce_max(split1, axis=(1, 2), keep_dims=True),
                                          tf.reduce_min(split1, axis=(1, 2), keep_dims=True)))

            plane = tf.concat([planeImN, split0], 2)
            split2 = tf.slice(UNet2D.nn, [0, 0, 0, 2], [-1, -1, -1, 1])

            # planeImN2 = tf.div(tf.subtract(split3, tf.reduce_min(split3, axis=(1, 2), keep_dims=True)),
            #                   tf.subtract(tf.reduce_max(split3, axis=(1, 2), keep_dims=True),
            #                               tf.reduce_min(split3, axis=(1, 2), keep_dims=True)))
            plane = tf.concat([plane, split2], 2)
            tf.summary.image('impm', plane, max_outputs=4)
        merged = tf.summary.merge_all()

        # --------------------------------------------------
        # session
        # --------------------------------------------------

        saver = tf.train.Saver()
        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True
        config.allow_soft_placement = True
        sess = tf.Session(config=config)  # config parameter needed to save variables when using GPU

        if os.path.exists(outLogPath):
            shutil.rmtree(outLogPath)
        trainWriter = tf.summary.FileWriter(trainWriterPath, sess.graph)
        validWriter = tf.summary.FileWriter(validWriterPath, sess.graph)

        if restoreVariables:
            saver.restore(sess, outModelPath)
            print("Model restored.")
        else:
            sess.run(tf.global_variables_initializer())

        # --------------------------------------------------
        # train
        # --------------------------------------------------

        batchData = np.zeros((batchSize, imSize, imSize, nChannels))
        batchLabels = np.zeros((batchSize, imSize, imSize, nClasses))
        batchWeights = np.zeros((batchSize, imSize, imSize, nClasses))
        permT = np.arange(nTrain)
        np.random.shuffle(permT)

        permV = np.arange(nValid)
        np.random.shuffle(permV)

        maxBrig = 1 * datasetStDev
        maxCont = 0.1 * datasetStDev
        jT = 0
        jV = 0
        epochCounter = 1
        for i in range(nSteps):
            # train

            for j in range(batchSize):
                fBrig = maxBrig * np.float_power(-1, np.random.rand() < 0.5) * np.random.rand()
                fCont = 1 + maxCont * np.float_power(-1, np.random.rand() < 0.5) * np.random.rand()
                image= Train[permT[jT + j], :, :, :] * fCont + fBrig
                # if np.random.rand() < 0.2:
                # 	image[:,:,1] = image[:,:,0]
                batchData[j, :, :, :] =  image

                batchLabels[j, :, :, :] = LTrain[permT[jT + j], :, :, :]
                batchWeights[j, :, :, :] = WTrain[permT[jT + j], :, :, :]
            summary, _ = sess.run([merged, optOp], feed_dict={UNet2D.tfData: batchData, tfLabels: batchLabels,
                                                              tfWeights: batchWeights, UNet2D.tfTraining: 1})
            jT = jT + batchSize
            if jT > (nTrain - batchSize - 1):
                jT = 0
                np.random.shuffle(permT)
                epochCounter = epochCounter + 1
            trainWriter.add_summary(summary, i)

            # validation
            for j in range(batchSize):
                batchData[j, :, :, :] = Valid[permV[jV + j], :, :, :]
                batchLabels[j, :, :, :] = LValid[permV[jV + j], :, :, :]
                batchWeights[j, :, :, :] = WValid[permV[jV + j], :, :, :]
            summary, es = sess.run([merged, errors], feed_dict={UNet2D.tfData: batchData, tfLabels: batchLabels,
                                                                tfWeights: batchWeights, UNet2D.tfTraining: 0})
            jV = jV + batchSize
            if jV > (nValid - batchSize - 1):
                jV = 0
                np.random.shuffle(permV)
            validWriter.add_summary(summary, i)

            e = np.mean(es)
            print('step %05d, e: %f' % (i, e) + ', epoch: ' + str(epochCounter))

            if i == 0:
                if restoreVariables:
                    lowestError = e
                else:
                    lowestError = np.inf

            if np.mod(i, 10) == 0 and e < lowestError:
                lowestError = e
                print("Model saved in file: %s" % saver.save(sess, outModelPath))

        # --------------------------------------------------
        # save hyper-parameters, clean-up
        # --------------------------------------------------

        saveData(UNet2D.hp, pathjoin(modelPath, 'hp.data'))

        trainWriter.close()
        validWriter.close()
        sess.close()


        # --------------------------------------------------
        # test
        # --------------------------------------------------
        os.environ['CUDA_VISIBLE_DEVICES'] = '%d' % gpuIndex
        tf.reset_default_graph()
        variablesPath = pathjoin(modelPath, 'model.ckpt')
        outPMPath = pmPath

        hp = loadData(pathjoin(modelPath, 'hp.data'))
        UNet2D.setupWithHP(hp)
        saver = tf.train.Saver()
        sess = tf.Session(config=tf.ConfigProto(
            allow_soft_placement=True))  # config parameter needed to save variables when using GPU
        saver.restore(sess, variablesPath)
        print("Model restored.")

        if not os.path.exists(outPMPath):
            os.makedirs(outPMPath)

        for i in range(nTest):
            j = np.mod(i, batchSize)

            batchData[j, :, :, :] = Test[i, :, :, :]
            batchLabels[j, :, :, :] = LTest[i, :, :, :]

            if j == batchSize - 1 or i == nTest - 1:

                output = sess.run(UNet2D.nn,
                                  feed_dict={UNet2D.tfData: batchData, UNet2D.tfTraining: 0})

                for k in range(j + 1):
                    pm = output[k, :, :, 2]
                    gt = batchLabels[k, :, :, 2]
                    im = np.sqrt(normalize(batchData[k, :, :, 0]))
                    imwrite(np.uint8(255 * np.concatenate((im, np.concatenate((pm, gt), axis=1)), axis=1)),
                            '%s/I%05dNuc.png' % (outPMPath, i - j + k + 1))

                for k in range(j + 1):
                    pm = output[k, :, :, 1]
                    gt = batchLabels[k, :, :, 1]
                    im = np.sqrt(normalize(batchData[k, :, :, 0]))
                    imwrite(np.uint8(255 * np.concatenate((im, np.concatenate((pm, gt), axis=1)), axis=1)),
                            '%s/I%05dCon.png' % (outPMPath, i - j + k + 1))



    def deploy(imPath, nImages, modelPath, pmPath, gpuIndex, pmIndex):
        os.environ['CUDA_VISIBLE_DEVICES'] = '%d' % gpuIndex
        tf.reset_default_graph()
        variablesPath = pathjoin(modelPath, 'model.ckpt')
        outPMPath = pmPath

        hp = loadData(pathjoin(modelPath, 'hp.data'))
        UNet2D.setupWithHP(hp)

        batchSize = UNet2D.hp['batchSize']
        imSize = UNet2D.hp['imSize']
        nChannels = UNet2D.hp['nChannels']
        nClasses = UNet2D.hp['nClasses']

        # --------------------------------------------------
        # data
        # --------------------------------------------------

        Data = np.zeros((nImages, imSize, imSize, nChannels))

        datasetMean = loadData(pathjoin(modelPath, 'datasetMean.data'))
        datasetStDev = loadData(pathjoin(modelPath, 'datasetStDev.data'))

        for iSample in range(0, nImages):
            path = '%s/I%05d_Img.tif' % (imPath, iSample)
            im = im2double(skio.imread(path, img_num=0))
            Data[iSample, :, :, 0] = (im - datasetMean) / datasetStDev

        # --------------------------------------------------
        # session
        # --------------------------------------------------

        saver = tf.train.Saver()
        sess = tf.Session(config=tf.ConfigProto(
            allow_soft_placement=True))  # config parameter needed to save variables when using GPU

        saver.restore(sess, variablesPath)
        print("Model restored.")

        # --------------------------------------------------
        # deploy
        # --------------------------------------------------

        batchData = np.zeros((batchSize, imSize, imSize, nChannels))

        if not os.path.exists(outPMPath):
            os.makedirs(outPMPath)

        for i in range(nImages):
            print(i, nImages)

            j = np.mod(i, batchSize)

            batchData[j, :, :, :] = Data[i, :, :, :]

            if j == batchSize - 1 or i == nImages - 1:

                output = sess.run(UNet2D.nn, feed_dict={UNet2D.tfData: batchData, UNet2D.tfTraining: 0})

                for k in range(j + 1):
                    pm = output[k, :, :, pmIndex]
                    im = np.sqrt(normalize(batchData[k, :, :, 0]))
                    # imwrite(np.uint8(255*np.concatenate((im,pm),axis=1)),'%s/I%05d.png' % (outPMPath,i-j+k+1))
                    imwrite(np.uint8(255 * im), '%s/I%05d_Im.png' % (outPMPath, i - j + k + 1))
                    imwrite(np.uint8(255 * pm), '%s/I%05d_PM.png' % (outPMPath, i - j + k + 1))

        # --------------------------------------------------
        # clean-up
        # --------------------------------------------------

        sess.close()

    def singleImageInferenceSetup(modelPath, gpuIndex,mean,std):
        variablesPath = pathjoin(modelPath, 'model.ckpt')

        hp = loadData(pathjoin(modelPath, 'hp.data'))
        UNet2D.setupWithHP(hp)
        if mean == -1:
            UNet2D.DatasetMean = loadData(pathjoin(modelPath, 'datasetMean.data'))
        else:
            UNet2D.DatasetMean = mean

        if std == -1:
            UNet2D.DatasetStDev = loadData(pathjoin(modelPath, 'datasetStDev.data'))
        else:
            UNet2D.DatasetStDev = std
        print(UNet2D.DatasetMean)
        print(UNet2D.DatasetStDev)

        # --------------------------------------------------
        # session
        # --------------------------------------------------

        saver = tf.train.Saver()
        UNet2D.Session = tf.Session(config=tf.ConfigProto(
            allow_soft_placement=True))  # config parameter needed to save variables when using GPU

        saver.restore(UNet2D.Session, variablesPath)
        print("Model restored.")

    def singleImageInferenceCleanup():
        UNet2D.Session.close()

    def singleImageInference(image, mode, pmIndex):
        print('Inference...')

        batchSize = UNet2D.hp['batchSize']
        imSize = UNet2D.hp['imSize']
        nChannels = UNet2D.hp['nChannels']

        PI2D.setup(image, imSize, int(imSize / 8), mode)
        PI2D.createOutput(1)

        batchData = np.zeros((batchSize, imSize, imSize, nChannels))
        for i in range(PI2D.NumPatches):
            j = np.mod(i, batchSize)
            P = (PI2D.getPatch(i) - UNet2D.DatasetMean) / UNet2D.DatasetStDev
            for iChan in range(nChannels):
                batchData[j, :, :, iChan] = P[iChan, :, :]
            if j == batchSize - 1 or i == PI2D.NumPatches - 1:
                output = UNet2D.Session.run(UNet2D.nn, feed_dict={UNet2D.tfData: batchData, UNet2D.tfTraining: 0})
                for k in range(j + 1):
                    pm = output[k, :, :, pmIndex]
                    PI2D.patchOutput(i - j + k, pm)
        # PI2D.patchOutput(i-j+k,normalize(imgradmag(PI2D.getPatch(i-j+k),1)))

        return PI2D.getValidOutput()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("imagePath", help="path to the .tif file")
    parser.add_argument("--model",  help="type of model. For example, nuclei vs cytoplasm",default = 'nucleiDAPILAMIN')
    parser.add_argument("--outputPath", help="output path of probability map")
    parser.add_argument("--channel", help="channel to perform inference on",  nargs = '+', default=[0])
    # parser.add_argument("--channel2", help="channel2 to perform inference on", type=int, default=-1)
    parser.add_argument("--classOrder", help="background, contours, foreground", type = int, nargs = '+', default=-1)
    parser.add_argument("--mean", help="mean intensity of input image. Use -1 to use model", type=float, default=-1)
    parser.add_argument("--std", help="mean standard deviation of input image. Use -1 to use model", type=float, default=-1)
    parser.add_argument("--scalingFactor", help="factor by which to increase/decrease image size by", type=float,
                        default=1)
    parser.add_argument("--stackOutput", help="save probability maps as separate files", action='store_true')
    parser.add_argument("--GPU", help="explicitly select GPU", type=int, default = -1)
    parser.add_argument("--outlier", help="map percentile intensity to max when rescaling intensity values. Max intensity as default", type=float, default=-1)
    parser.add_argument("--verbose", help="display error messages for debugging", action='store_true')
    args = parser.parse_args()

    if args.verbose:
        os.environ['TF_CPP_MIN_LOG_LEVEL'] = '0'
        logging.getLogger('tensorflow').setLevel(logging.DEBUG)
    else:
        os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
        logging.getLogger('tensorflow').setLevel(logging.FATAL)


    logPath = ''
    scriptPath = os.path.dirname(os.path.realpath(__file__))
    modelPath = os.path.join(scriptPath, 'models', args.model)
    # modelPath = os.path.join(scriptPath, 'models/cytoplasmINcell')
    # modelPath = os.path.join(scriptPath, 'cytoplasmZeissNikon')
    pmPath = ''
    if os.system('nvidia-smi') == 0:
        if args.GPU == -1:
            print("automatically choosing GPU")
            GPU = GPUselect.pick_gpu_lowest_memory()
        else:
            GPU = args.GPU
        print('Using GPU ' + str(GPU))

    else:
        if sys.platform == 'win32':  # only 1 gpu on windows
            if args.GPU==-1:
                GPU = 0
                print('using default GPU')
            else:
                GPU = args.GPU
            print('Using GPU ' + str(GPU))
        else:
            GPU=0
            print('Using CPU')
    os.environ['CUDA_VISIBLE_DEVICES'] = '%d' % GPU
    UNet2D.singleImageInferenceSetup(modelPath, GPU,args.mean,args.std)
    nClass = UNet2D.hp['nClasses']
    imagePath = args.imagePath
    dapiChannel = args.channel[0]
    if len(args.channel)==1:
        channel = [dapiChannel, dapiChannel]
    else:
        channel = args.channel
    print ('Using channels ' + str(int(channel[0])+1) + ' and ' + str(int(channel[1])+1))
    dsFactor = args.scalingFactor
    parentFolder = os.path.dirname(os.path.dirname(imagePath))
    fileName = os.path.basename(imagePath)
    fileNamePrefix = fileName.split(os.extsep, 1)
    # print(fileName)
    fileType = fileNamePrefix[1]

    for iChan in range(len(channel)):
        if fileType=='ome.tif' or fileType=='ome.tiff' or fileType == 'btf' :
            I = skio.imread(imagePath, img_num=int(channel[iChan]),plugin='tifffile')
        elif fileType == 'tif' :
            I = tifffile.imread(imagePath, key=int(channel[iChan]))
        elif fileType == 'czi':
            with czifile.CziFile(imagePath) as czi:
                image = czi.asarray()
                I = image[0, 0, int(channel[iChan]), 0, 0, :, :, 0]
        elif fileType == 'nd2':
            with ND2Reader(imagePath) as fullStack:
                I = fullStack[int(channel[iChan])]
        if I.dtype == 'float32':
            I = np.uint16(I)
        rawVert = I.shape[0]
        rawHorz = I.shape[1]
        rawI = I

        hsize = int((float(I.shape[0]) * float(dsFactor)))
        vsize = int((float(I.shape[1]) * float(dsFactor)))
        I = resize(I, (hsize, vsize))
        if iChan ==0:
            cells = np.zeros((len(channel), I.shape[0], I.shape[1]))
        if args.outlier == -1:
            maxLimit = np.max(I)
        else:
            maxLimit = np.percentile(I,args.outlier)
        I = im2double(sk.rescale_intensity(I, in_range=(np.min(I), maxLimit), out_range=(0, 0.983)))
        cells[iChan, :, :] = I

    if args.classOrder == -1:
        args.classOrder = range(nClass)
    rawI = im2double(rawI) / np.max(im2double(rawI))
    if not args.outputPath:
        args.outputPath = parentFolder + '//probability_maps'

    if not os.path.exists(args.outputPath):
        os.makedirs(args.outputPath)
    os.makedirs(args.outputPath + '//qc', exist_ok=True)

    append_kwargs = {
        'bigtiff': True,
        'metadata': None,
        'append': True,
    }
    save_kwargs = {
        'bigtiff': True,
        'metadata': None,
        'append': False,
    }
    if args.stackOutput:
        slice=0
        for iClass in args.classOrder[::-1]:
            PM = np.uint8(255*UNet2D.singleImageInference(cells, 'accumulate', iClass)) # backwards in order to align with ilastik...
            PM = resize(PM, (rawVert, rawHorz))
            if slice==0:
                skimage.io.imsave(args.outputPath + '//' + fileNamePrefix[0] + '_Probabilities_' + str(int(dapiChannel)+1) + '.tif', np.uint8(255 * PM),**save_kwargs)
            else:
                skimage.io.imsave(args.outputPath + '//' + fileNamePrefix[0] + '_Probabilities_' + str(int(dapiChannel)+1) + '.tif',np.uint8(255 * PM),**append_kwargs)
            if slice==1:
                save_kwargs['append'] = False
                skimage.io.imsave(args.outputPath +  '//qc//' + fileNamePrefix[0] + '_Preview_' + str(int(dapiChannel)+1) + '.tif',	np.uint8(255 * PM), **save_kwargs)
                skimage.io.imsave(args.outputPath +  '//qc//' + fileNamePrefix[0] + '_Preview_' + str(int(dapiChannel)+1) + '.tif', np.uint8(255 * rawI), **append_kwargs)
            slice = slice + 1

    else:
        contours = np.uint8(255*UNet2D.singleImageInference(cells, 'accumulate', args.classOrder[1]))
        contours = resize(contours, (rawVert, rawHorz))
        skimage.io.imsave(args.outputPath + '//' + fileNamePrefix[0] + '_ContoursPM_' + str(int(dapiChannel)+1) + '.tif',np.uint8(255 * contours),**save_kwargs)
        skimage.io.imsave(args.outputPath + '//' + fileNamePrefix[0] + '_ContoursPM_' + str(int(dapiChannel)+1) + '.tif',np.uint8(255 * rawI), **append_kwargs)
        del contours
        nuclei = np.uint8(255*UNet2D.singleImageInference(cells, 'accumulate', args.classOrder[2]))
        nuclei = resize(nuclei, (rawVert, rawHorz))
        skimage.io.imsave(args.outputPath + '//' + fileNamePrefix[0] + '_NucleiPM_' + str(int(dapiChannel)+1) + '.tif',np.uint8(255 * nuclei), **save_kwargs)
        del nuclei
    UNet2D.singleImageInferenceCleanup()