diff UNetCoreograph.py @ 0:99308601eaa6 draft

"planemo upload for repository https://github.com/ohsu-comp-bio/UNetCoreograph commit fb90660a1805b3f68fcff80d525b5459c3f7dfd6-dirty"
author perssond
date Wed, 19 May 2021 21:34:38 +0000
parents
children 57f1260ca94e
line wrap: on
line diff
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/UNetCoreograph.py	Wed May 19 21:34:38 2021 +0000
@@ -0,0 +1,802 @@
+import numpy as np
+from scipy import misc as sm
+import shutil
+import scipy.io as sio
+import os
+import skimage.exposure as sk
+import cv2
+import argparse
+import pytiff
+import tifffile
+import tensorflow as tf
+from skimage.morphology import *
+from skimage.exposure import rescale_intensity
+from skimage.segmentation import chan_vese, find_boundaries, morphological_chan_vese
+from skimage.measure import regionprops,label, find_contours
+from skimage.transform import resize
+from skimage.filters import gaussian
+from skimage.feature import peak_local_max,blob_log
+from skimage.color import label2rgb
+import skimage.io as skio
+from skimage import img_as_bool
+from skimage.draw import circle_perimeter
+from scipy.ndimage.filters import uniform_filter
+from scipy.ndimage import gaussian_laplace
+from os.path import *
+from os import listdir, makedirs, remove
+
+
+
+import sys
+from typing import Any
+
+#sys.path.insert(0, 'C:\\Users\\Public\\Documents\\ImageScience')
+from toolbox.imtools import *
+from toolbox.ftools import *
+from toolbox.PartitionOfImage import PI2D
+
+
+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):
+			with tf.name_scope('ld%d' % index):
+				ldXWeights1 = tf.Variable(tf.truncated_normal([UNet2D.hp['ks'], UNet2D.hp['ks'], nOutX[index], nOutX[index+1]], stddev=stdDev0),name='kernel1')
+				ldXWeightsExtra = []
+				for i in range(nExtraConvs):
+					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.relu(c00), ldXWeightsExtra[i], strides=[1, 1, 1, 1], padding='SAME')
+
+				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.layers.batch_normalization(tf.nn.relu(c00+shortcut), 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.name_scope('lb'):
+			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):
+				return tf.nn.relu(tf.nn.conv2d(hidden, lbWeights1, strides=[1, 1, 1, 1], padding='SAME'),name='conv')
+
+		# --------------------------------------------------
+		# 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.name_scope('lu%d' % index):
+				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.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.relu(tf.nn.conv2d(cc, luXWeights2, strides=[1, 1, 1, 1], padding='SAME'),name='conv2')
+				for i in range(nExtraConvs):
+					cv = tf.nn.relu(tf.nn.conv2d(cv, luXWeightsExtra[i], strides=[1, 1, 1, 1], padding='SAME'),name='conv2Extra%d' % i)
+				return cv
+
+		# --------------------------------------------------
+		# final (top) layer
+		# --------------------------------------------------
+
+		with tf.name_scope('lt'):
+			ltWeights1    = tf.Variable(tf.truncated_normal([1, 1, nOutX[1], nClasses], stddev=stdDev0),name='kernel')
+			def lt(hidden):
+				return tf.nn.conv2d(hidden, ltWeights1, strides=[1, 1, 1, 1], padding='SAME',name='conv')
+
+
+		# --------------------------------------------------
+		# 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,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))
+
+		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
+			datasetStDev = 0
+			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+nValid+nTest)
+		np.random.shuffle(perm)
+
+		for iSample in range(0, nTrain):
+			path = '%s/I%05d_Img.tif' % (imPath,perm[iSample])
+			im = im2double(tifread(path))
+			Train[iSample,:,:,0] = (im-datasetMean)/datasetStDev
+			path = '%s/I%05d_Ant.tif' % (imPath,perm[iSample])
+			im = tifread(path)
+			for i in range(nClasses):
+				LTrain[iSample,:,:,i] = (im == i+1)
+
+		for iSample in range(0, nValid):
+			path = '%s/I%05d_Img.tif' % (imPath,perm[nTrain+iSample])
+			im = im2double(tifread(path))
+			Valid[iSample,:,:,0] = (im-datasetMean)/datasetStDev
+			path = '%s/I%05d_Ant.tif' % (imPath,perm[nTrain+iSample])
+			im = tifread(path)
+			for i in range(nClasses):
+				LValid[iSample,:,:,i] = (im == i+1)
+
+		for iSample in range(0, nTest):
+			path = '%s/I%05d_Img.tif' % (imPath,perm[nTrain+nValid+iSample])
+			im = im2double(tifread(path))
+			Test[iSample,:,:,0] = (im-datasetMean)/datasetStDev
+			path = '%s/I%05d_Ant.tif' % (imPath,perm[nTrain+nValid+iSample])
+			im = tifread(path)
+			for i in range(nClasses):
+				LTest[iSample,:,:,i] = (im == i+1)
+
+		# --------------------------------------------------
+		# optimization
+		# --------------------------------------------------
+
+		tfLabels = tf.placeholder("float", shape=[None,imSize,imSize,nClasses],name='labels')
+
+		globalStep = tf.Variable(0,trainable=False)
+		learningRate0 = 0.01
+		decaySteps = 1000
+		decayRate = 0.95
+		learningRate = tf.train.exponential_decay(learningRate0,globalStep,decaySteps,decayRate,staircase=True)
+
+		with tf.name_scope('optim'):
+			loss = tf.reduce_mean(-tf.reduce_sum(tf.multiply(tfLabels,tf.log(UNet2D.nn)),3))
+			updateOps = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
+			# optimizer = tf.train.MomentumOptimizer(1e-3,0.9)
+			optimizer = tf.train.MomentumOptimizer(learningRate,0.9)
+			# optimizer = tf.train.GradientDescentOptimizer(learningRate)
+			with tf.control_dependencies(updateOps):
+				optOp = optimizer.minimize(loss,global_step=globalStep)
+
+		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,0],[-1,-1,-1,1])
+			split1 = tf.slice(UNet2D.nn,[0,0,0,1],[-1,-1,-1,1])
+			if nClasses > 2:
+				split2 = tf.slice(UNet2D.nn,[0,0,0,2],[-1,-1,-1,1])
+			tf.summary.image('pm0',split0)
+			tf.summary.image('pm1',split1)
+			if nClasses > 2:
+				tf.summary.image('pm2',split2)
+		merged = tf.summary.merge_all()
+
+
+		# --------------------------------------------------
+		# session
+		# --------------------------------------------------
+
+		saver = tf.train.Saver()
+		sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) # 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))
+		for i in range(nSteps):
+			# train
+
+			perm = np.arange(nTrain)
+			np.random.shuffle(perm)
+
+			for j in range(batchSize):
+				batchData[j,:,:,:] = Train[perm[j],:,:,:]
+				batchLabels[j,:,:,:] = LTrain[perm[j],:,:,:]
+
+			summary,_ = sess.run([merged,optOp],feed_dict={UNet2D.tfData: batchData, tfLabels: batchLabels, UNet2D.tfTraining: 1})
+			trainWriter.add_summary(summary, i)
+
+			# validation
+
+			perm = np.arange(nValid)
+			np.random.shuffle(perm)
+
+			for j in range(batchSize):
+				batchData[j,:,:,:] = Valid[perm[j],:,:,:]
+				batchLabels[j,:,:,:] = LValid[perm[j],:,:,:]
+
+			summary, es = sess.run([merged, errors],feed_dict={UNet2D.tfData: batchData, tfLabels: batchLabels, UNet2D.tfTraining: 0})
+			validWriter.add_summary(summary, i)
+
+			e = np.mean(es)
+			print('step %05d, e: %f' % (i,e))
+
+			if i == 0:
+				if restoreVariables:
+					lowestError = e
+				else:
+					lowestError = np.inf
+
+			if np.mod(i,100) == 0 and e < lowestError:
+				lowestError = e
+				print("Model saved in file: %s" % saver.save(sess, outModelPath))
+
+
+		# --------------------------------------------------
+		# test
+		# --------------------------------------------------
+
+		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, tfLabels: batchLabels, UNet2D.tfTraining: 0})
+
+				for k in range(j+1):
+					pm = output[k,:,:,testPMIndex]
+					gt = batchLabels[k,:,:,testPMIndex]
+					im = np.sqrt(normalize(batchData[k,:,:,0]))
+					imwrite(np.uint8(255*np.concatenate((im,np.concatenate((pm,gt),axis=1)),axis=1)),'%s/I%05d.png' % (outPMPath,i-j+k+1))
+
+
+		# --------------------------------------------------
+		# save hyper-parameters, clean-up
+		# --------------------------------------------------
+
+		saveData(UNet2D.hp,pathjoin(modelPath,'hp.data'))
+
+		trainWriter.close()
+		validWriter.close()
+		sess.close()
+
+	def deploy(imPath,nImages,modelPath,pmPath,gpuIndex,pmIndex):
+		os.environ['CUDA_VISIBLE_DEVICES']= '%d' % gpuIndex
+		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(tifread(path))
+			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):
+		os.environ['CUDA_VISIBLE_DEVICES']= '%d' % gpuIndex
+		variablesPath = pathjoin(modelPath,'model.ckpt')
+		hp = loadData(pathjoin(modelPath,'hp.data'))
+		UNet2D.setupWithHP(hp)
+
+		UNet2D.DatasetMean =loadData(pathjoin(modelPath,'datasetMean.data'))
+		UNet2D.DatasetStDev =  loadData(pathjoin(modelPath,'datasetStDev.data'))
+		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
+		#UNet2D.Session = tf.Session(config=tf.ConfigProto(device_count={'GPU': 0}))
+		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(nChannels)
+
+		batchData = np.zeros((batchSize,imSize,imSize,nChannels))
+		for i in range(PI2D.NumPatches):
+			j = np.mod(i,batchSize)
+			batchData[j,:,:,0] = (PI2D.getPatch(i)-UNet2D.DatasetMean)/UNet2D.DatasetStDev
+			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()
+
+
+def identifyNumChan(path):
+   tiff = tifffile.TiffFile(path)
+   shape = tiff.pages[0].shape
+   numChan=None
+   for i, page in enumerate(tiff.pages):
+      if page.shape != shape:
+         numChan = i
+         return numChan
+         break
+#      else:
+#         raise Exception("Did not find any pyramid subresolutions") 
+
+   if not numChan:
+      numChan = len(tiff.pages)
+      return numChan
+
+def getProbMaps(I,dsFactor,modelPath):
+   hsize = int((float(I.shape[0]) * float(0.5)))
+   vsize = int((float(I.shape[1]) * float(0.5)))
+   imagesub = cv2.resize(I,(vsize,hsize),cv2.INTER_NEAREST)
+
+   UNet2D.singleImageInferenceSetup(modelPath, 1)
+
+   for iSize in range(dsFactor):
+	   hsize = int((float(I.shape[0]) * float(0.5)))
+	   vsize = int((float(I.shape[1]) * float(0.5)))
+	   I = cv2.resize(I,(vsize,hsize),cv2.INTER_NEAREST)
+   I = im2double(I)
+   I = im2double(sk.rescale_intensity(I, in_range=(np.min(I), np.max(I)), out_range=(0, 0.983)))
+   probMaps = UNet2D.singleImageInference(I,'accumulate',1)
+   UNet2D.singleImageInferenceCleanup()
+   return probMaps 
+
+def coreSegmenterOutput(I,probMap,initialmask,preBlur,findCenter):
+	hsize = int((float(I.shape[0]) * float(0.1)))
+	vsize = int((float(I.shape[1]) * float(0.1)))
+	nucGF = cv2.resize(I,(vsize,hsize),cv2.INTER_CUBIC)
+#	Irs = cv2.resize(I,(vsize,hsize),cv2.INTER_CUBIC)
+#	I=I.astype(np.float)
+#	r,c = I.shape
+#	I+=np.random.rand(r,c)*1e-6
+#	c1 = uniform_filter(I, 3, mode='reflect')
+#	c2 = uniform_filter(I*I, 3, mode='reflect')
+#	nucGF = np.sqrt(c2 - c1*c1)*np.sqrt(9./8)
+#	nucGF[np.isnan(nucGF)]=0
+	#active contours
+	hsize = int(float(nucGF.shape[0]))
+	vsize = int(float(nucGF.shape[1]))
+	initialmask = cv2.resize(initialmask,(vsize,hsize),cv2.INTER_NEAREST)
+	initialmask = dilation(initialmask,disk(15)) >0
+		
+#	init=np.argwhere(eroded>0)
+	nucGF = gaussian(nucGF,0.7)
+	nucGF=nucGF/np.amax(nucGF)
+	
+   
+#	initialmask = nucGF>0
+	nuclearMask = morphological_chan_vese(nucGF, 100, init_level_set=initialmask, smoothing=10,lambda1=1.001, lambda2=1)
+	
+#	nuclearMask = chan_vese(nucGF, mu=1.5, lambda1=6, lambda2=1, tol=0.0005, max_iter=2000, dt=15, init_level_set=initialmask, extended_output=True)	
+#	nuclearMask = nuclearMask[0]
+  
+	
+	TMAmask = nuclearMask
+#	nMaskDist =distance_transform_edt(nuclearMask)
+#	fgm = peak_local_max(h_maxima(nMaskDist, 2*preBlur),indices =False)
+#	markers= np.logical_or(erosion(1-nuclearMask,disk(3)),fgm)
+#	TMAmask=watershed(-nMaskDist,label(markers),watershed_line=True)
+#	TMAmask = nuclearMask*(TMAmask>0)
+	TMAmask = remove_small_objects(TMAmask>0,round(TMAmask.shape[0])*round(TMAmask.shape[1])*0.005)
+	TMAlabel = label(TMAmask)
+# find object closest to center
+	if findCenter==True:
+		
+		stats= regionprops(TMAlabel)
+		counter=1
+		minDistance =-1
+		index =[]
+		for props in stats:
+			centroid = props.centroid
+			distanceFromCenter = np.sqrt((centroid[0]-nucGF.shape[0]/2)**2+(centroid[1]-nucGF.shape[1]/2)**2)
+	#		if distanceFromCenter<0.6/2*np.sqrt(TMAlabel.shape[0]*TMAlabel.shape[1]):
+			if distanceFromCenter<minDistance or minDistance==-1 :
+				minDistance =distanceFromCenter
+				index = counter
+			counter=counter+1
+	#		dist = 0.6/2*np.sqrt(TMAlabel.shape[0]*TMAlabel.shape[1])
+		TMAmask = morphology.binary_closing(TMAlabel==index,disk(3))
+
+	return TMAmask
+
+def overlayOutline(outline,img):
+	img2 = img.copy()
+	stacked_img = np.stack((img2,)*3, axis=-1)
+	stacked_img[outline > 0] = [1, 0, 0]
+	imshowpair(img2,stacked_img)
+
+def imshowpair(A,B):
+	plt.imshow(A,cmap='Purples')
+	plt.imshow(B,cmap='Greens',alpha=0.5)
+	plt.show()
+
+
+if __name__ == '__main__':
+	parser=argparse.ArgumentParser()
+	parser.add_argument("--imagePath")
+	parser.add_argument("--outputPath")
+	parser.add_argument("--maskPath")
+	parser.add_argument("--downsampleFactor",type = int, default = 5)
+	parser.add_argument("--channel",type = int, default = 0)
+	parser.add_argument("--buffer",type = float, default = 2)
+	parser.add_argument("--outputChan", type=int, nargs = '+', default=[-1])
+	parser.add_argument("--sensitivity",type = float, default=0.3)
+	parser.add_argument("--useGrid",action='store_true')
+	parser.add_argument("--cluster",action='store_true')
+	args = parser.parse_args()
+
+	outputPath = args.outputPath
+	imagePath = args.imagePath
+	sensitivity = args.sensitivity
+	#scriptPath = os.path.dirname(os.path.realpath(__file__))
+	#modelPath = os.path.join(scriptPath, 'TFModel - 3class 16 kernels 5ks 2 layers')
+	#modelPath = 'D:\\LSP\\Coreograph\\model-4layersMaskAug20'
+	scriptPath = os.path.dirname(os.path.realpath(__file__))
+	modelPath = os.path.join(scriptPath, 'model')
+#	outputPath = 'D:\\LSP\\cycif\\testsets\\exemplar-002\\dearrayPython' ############
+	maskOutputPath = os.path.join(outputPath, 'masks')
+#	imagePath = 'D:\\LSP\\cycif\\testsets\\exemplar-002\\registration\\exemplar-002.ome.tif'###########
+#	imagePath = 'Y:\\sorger\\data\\RareCyte\\Connor\\TMAs\\CAJ_TMA11_13\\original_data\\TMA11\\registration\\TMA11.ome.tif'
+#	imagePath = 'Y:\\sorger\\data\\RareCyte\\Connor\\TMAs\\Z124_TMA20_22\\TMA22\\registration\\TMA22.ome.tif'
+#	classProbsPath = 'D:\\unetcoreograph.tif'
+#	imagePath = 'Y:\\sorger\\data\\RareCyte\\Connor\\Z155_PTCL\\TMA_552\\registration\\TMA_552.ome.tif'
+#	classProbsPath = 'Y:\\sorger\\data\\RareCyte\\Connor\\Z155_PTCL\\TMA_552\\probMapCore\\TMA_552_CorePM_1.tif'
+#	imagePath = 'Y:\\sorger\\data\\RareCyte\\Zoltan\\Z112_TMA17_19\\190403_ashlar\\TMA17_1092.ome.tif'
+#	classProbsPath = 'Z:\\IDAC\\Clarence\\LSP\\CyCIF\\TMA\\probMapCore\\1new_CorePM_1.tif'
+#	imagePath = 'Y:\\sorger\\data\\RareCyte\\ANNIINA\\Julia\\2018\\TMA6\\julia_tma6.ome.tif'
+#	classProbsPath = 'Z:\\IDAC\\Clarence\\LSP\\CyCIF\\TMA\\probMapCore\\3new_CorePM_1.tif'
+
+	
+#	if not os.path.exists(outputPath):
+#		os.makedirs(outputPath)
+#	else:
+#		shutil.rmtree(outputPath)
+	if not os.path.exists(maskOutputPath):
+		os.makedirs(maskOutputPath)
+
+
+	channel = args.channel 
+	dsFactor = 1/(2**args.downsampleFactor)
+#	I = tifffile.imread(imagePath, key=channel)
+	I = skio.imread(imagePath, img_num=channel)
+
+	imagesub = resize(I,(int((float(I.shape[0]) * dsFactor)),int((float(I.shape[1]) * dsFactor))))
+	numChan = identifyNumChan(imagePath)
+	
+	outputChan = args.outputChan
+	if len(outputChan)==1:
+		if outputChan[0]==-1:
+			outputChan = [0, numChan-1]
+		else:
+			outputChan.append(outputChan[0])
+	
+	classProbs = getProbMaps(I,args.downsampleFactor,modelPath)
+#	classProbs = tifffile.imread(classProbsPath,key=0)
+	preMask = gaussian(np.uint8(classProbs*255),1)>0.8
+	
+	P = regionprops(label(preMask),cache=False)
+	area = [ele.area for ele in P]
+	print(str(len(P)) + ' cores detected!')
+	if len(P) <3:
+		medArea = np.median(area)
+		maxArea = np.percentile(area,99)
+	else:
+		count=0
+		labelpreMask = np.zeros(preMask.shape,dtype=np.uint32)
+		for props in P:
+				count += 1
+				yi = props.coords[:, 0]
+				xi = props.coords[:, 1]
+				labelpreMask[yi, xi] = count              
+				P=regionprops(labelpreMask)
+				area = [ele.area for ele in P]
+		medArea =  np.median(area)
+		maxArea = np.percentile(area,99)
+	preMask = remove_small_objects(preMask,0.2*medArea)
+	coreRad = round(np.sqrt(medArea/np.pi))
+	estCoreDiam = round(np.sqrt(maxArea/np.pi)*1.2*args.buffer)
+
+#preprocessing
+	fgFiltered = blob_log(preMask,coreRad*0.6,threshold=sensitivity)
+	Imax = np.zeros(preMask.shape,dtype=np.uint8)
+	for iSpot in range(fgFiltered.shape[0]):
+		yi = np.uint32(round(fgFiltered[iSpot, 0]))
+		xi = np.uint32(round(fgFiltered[iSpot, 1]))
+		Imax[yi, xi] = 1
+	Imax = Imax*preMask
+	Idist = distance_transform_edt(1-Imax)
+	markers = label(Imax)
+	coreLabel  = watershed(Idist,markers,watershed_line=True,mask = preMask)
+	P = regionprops(coreLabel)
+	centroids = np.array([ele.centroid for ele in P])/dsFactor
+	numCores = len(centroids)
+	estCoreDiamX = np.ones(numCores)*estCoreDiam/dsFactor
+	estCoreDiamY = np.ones(numCores)*estCoreDiam/dsFactor
+
+	if numCores ==0 & args.cluster:
+		print('No cores detected. Try adjusting the downsample factor')
+		sys.exit(255)
+
+	singleMaskTMA = np.zeros(imagesub.shape)
+	maskTMA = np.zeros(imagesub.shape)
+	bbox = [None] * numCores
+
+ 
+	x=np.zeros(numCores)
+	xLim=np.zeros(numCores)
+	y=np.zeros(numCores)
+	yLim=np.zeros(numCores)
+	
+# segmenting each core   	
+	#######################
+	for iCore in range(numCores):
+		x[iCore] = centroids[iCore,1] - estCoreDiamX[iCore]/2
+		xLim[iCore] = x[iCore]+estCoreDiamX[iCore]
+		if xLim[iCore] > I.shape[1]:
+			xLim[iCore] = I.shape[1]
+		if x[iCore]<1:
+			x[iCore]=1
+
+		y[iCore] = centroids[iCore,0] - estCoreDiamY[iCore]/2
+		yLim[iCore] = y[iCore] + estCoreDiamY[iCore]
+		if yLim[iCore] > I.shape[0]:
+			yLim[iCore] = I.shape[0]
+		if y[iCore]<1:
+			y[iCore]=1
+
+		bbox[iCore] = [round(x[iCore]), round(y[iCore]), round(xLim[iCore]), round(yLim[iCore])]
+		
+		for iChan in range(outputChan[0],outputChan[1]+1):
+			with pytiff.Tiff(imagePath, "r", encoding='utf-8') as handle:
+				handle.set_page(iChan)
+				coreStack= handle[np.uint32(bbox[iCore][1]):np.uint32(bbox[iCore][3]-1), np.uint32(bbox[iCore][0]):np.uint32(bbox[iCore][2]-1)]
+			skio.imsave(outputPath + os.path.sep + str(iCore+1)  + '.tif',coreStack,append=True)	
+
+		with pytiff.Tiff(imagePath, "r", encoding='utf-8') as handle:
+			handle.set_page(args.channel)
+			coreSlice= handle[np.uint32(bbox[iCore][1]):np.uint32(bbox[iCore][3]-1), np.uint32(bbox[iCore][0]):np.uint32(bbox[iCore][2]-1)]
+
+		core = (coreLabel ==(iCore+1))
+		initialmask = core[np.uint32(y[iCore]*dsFactor):np.uint32(yLim[iCore]*dsFactor),np.uint32(x[iCore]*dsFactor):np.uint32(xLim[iCore]*dsFactor)]
+		initialmask = resize(initialmask,size(coreSlice),cv2.INTER_NEAREST)
+
+		singleProbMap = classProbs[np.uint32(y[iCore]*dsFactor):np.uint32(yLim[iCore]*dsFactor),np.uint32(x[iCore]*dsFactor):np.uint32(xLim[iCore]*dsFactor)]
+		singleProbMap = resize(np.uint8(255*singleProbMap),size(coreSlice),cv2.INTER_NEAREST)
+		TMAmask = coreSegmenterOutput(coreSlice,singleProbMap,initialmask,coreRad/20,False) 
+		if np.sum(TMAmask)==0:
+			TMAmask = np.ones(TMAmask.shape)
+		vsize = int(float(coreSlice.shape[0]))
+		hsize = int(float(coreSlice.shape[1]))
+		masksub = resize(resize(TMAmask,(vsize,hsize),cv2.INTER_NEAREST),(int((float(coreSlice.shape[0])*dsFactor)),int((float(coreSlice.shape[1])*dsFactor))),cv2.INTER_NEAREST)
+		singleMaskTMA[int(y[iCore]*dsFactor):int(y[iCore]*dsFactor)+masksub.shape[0],int(x[iCore]*dsFactor):int(x[iCore]*dsFactor)+masksub.shape[1]]=masksub
+		maskTMA = maskTMA + resize(singleMaskTMA,maskTMA.shape,cv2.INTER_NEAREST)
+		cv2.putText(imagesub, str(iCore+1), (int(P[iCore].centroid[1]),int(P[iCore].centroid[0])), 0, 0.5, (np.amax(imagesub), np.amax(imagesub), np.amax(imagesub)), 1, cv2.LINE_AA)
+		
+		skio.imsave(maskOutputPath + os.path.sep + str(iCore+1)  + '_mask.tif',np.uint8(TMAmask))
+		print('Segmented core ' + str(iCore+1))	
+		
+	boundaries = find_boundaries(maskTMA)
+	imagesub = imagesub/np.percentile(imagesub,99.9)
+	imagesub[boundaries==1] = 1
+	skio.imsave(outputPath + os.path.sep + 'TMA_MAP.tif' ,np.uint8(imagesub*255))
+	print('Segmented all cores!')
+	
+
+#restore GPU to 0
+	#image load using tifffile