models.py

import tensorflow as tf, numpy as np, data, os, math, pickle, matplotlib.pyplot as plt
from datetime import datetime
from tensorflow.keras import backend as K
from classes import modelSet, modelDictVal

huberDelta = .5

def smoothL1(y_true, y_pred): # custom loss function for unit vector output
	x = tf.keras.backend.abs(y_true - y_pred)
	x = tf.where(x < huberDelta, 0.5 * x ** 2, huberDelta * (x - 0.5 * huberDelta))
	return	tf.keras.backend.sum(x)

def coordsOutPut(x): # add coordinate output layer
	coords = tf.keras.layers.Conv2D(18, (1,1), name = 'coordsOut', kernel_initializer = tf.keras.initializers.GlorotUniform(seed=0), padding = 'same')(x)
	#coords = tf.keras.layers.BatchNormalization(name = 'batchCoords')(coords)
	#coords = tf.keras.layers.Activation('relu')(coords)
	return coords
	
def classOutput(x): # add class output layer
	classPred = tf.keras.layers.Conv2D(1, (1,1), name = 'classConv', kernel_initializer = tf.keras.initializers.GlorotUniform(seed=0), padding = 'same')(x)
	classPred = tf.keras.layers.BatchNormalization(name = 'classBatch')(classPred)
	classPred = tf.keras.layers.Activation('relu', name = "classOut")(classPred)
	
	return classPred

def convLayer(x, numFilters, kernelSize, strides = 1, dilation = 1):
	x = tf.keras.layers.Conv2D(numFilters, kernelSize, strides = strides, kernel_initializer = tf.keras.initializers.GlorotUniform(seed=0), padding = 'same', dilation_rate = dilation)(x)
	x = tf.keras.layers.BatchNormalization()(x)
	x = tf.keras.layers.Activation('relu')(x)
	
	return x
	
def stvNet(inputShape = (480, 640, 3), outVectors = True, outClasses = True, modelName = "stvNet"):
	
	xIn = tf.keras.Input(inputShape, dtype = np.dtype('uint8'))
	
	x = tf.keras.layers.Lambda(lambda x: x / 255) (xIn)
	
	x = tf.keras.layers.Conv2D(64, 7, input_shape = inputShape, kernel_initializer = tf.keras.initializers.GlorotUniform(seed=0), padding = 'same')(x)
	x = tf.keras.layers.BatchNormalization()(x)
	x = tf.keras.layers.Activation('relu')(x)
	
	res1 = x
	
	x = tf.keras.layers.MaxPool2D(pool_size = 3, strides = 2, padding = 'same')(x)
	
	skip = x
	
	x = convLayer(x, 64, 3)
	x = convLayer(x, 64, 3)
	
	x = tf.keras.layers.Add()([x, skip])
	skip = x
	
	x = convLayer(x, 64, 3)
	x = convLayer(x, 64, 3)
	
	x = tf.keras.layers.Add()([x, skip])
	skip = tf.keras.layers.MaxPool2D(pool_size = 2, padding = 'same')(x)
	skip = tf.pad(skip, [[0, 0], [0, 0], [0, 0], [32, 32]]) # linear projection
	res2 = x
	
	x = convLayer(x, 128, 3, 2)
	x = convLayer(x, 128, 3)
	
	x = tf.keras.layers.Add()([x, skip])
	skip = x
	
	x = convLayer(x, 128, 3)
	x = convLayer(x, 128, 3)
	
	x = tf.keras.layers.Add()([x, skip])
	skip = tf.keras.layers.MaxPool2D(pool_size = 2, padding = 'same')(x)
	skip = tf.pad(skip, [[0, 0], [0, 0], [0, 0], [64, 64]]) # linear projection
	res3 = x
	
	x = convLayer(x, 256, 3, 2)
	x = convLayer(x, 256, 3)
	
	x = tf.keras.layers.Add()([x, skip])
	skip = x
	
	x = convLayer(x, 256, 3)
	x = convLayer(x, 256, 3)
	
	x = tf.keras.layers.Add()([x, skip])
	skip = tf.pad(x, [[0, 0], [0, 0], [0, 0], [128, 128]])
	res4 = x
	
	x = convLayer(x, 512, 3, dilation = 2)
	x = convLayer(x, 512, 3, dilation = 2)
	
	x = tf.keras.layers.Add()([x, skip])
	skip = x
	
	x = convLayer(x, 512, 3, dilation = 2)
	x = convLayer(x, 512, 3, dilation = 2)
	
	x = tf.keras.layers.Add()([x, skip])
	
	x = convLayer(x, 256, 3)
	x = convLayer(x, 256, 3)
	
	x = tf.keras.layers.Add()([x, res4])
	x = tf.keras.layers.UpSampling2D()(x)
	
	x = convLayer(x, 128, 3)
	
	x = tf.keras.layers.Add()([x, res3])
	x = tf.keras.layers.UpSampling2D()(x)
	
	x = convLayer(x, 64, 3)
	
	x = tf.keras.layers.Add()([x, res2])
	x = tf.keras.layers.UpSampling2D()(x)
	
	x = convLayer(x, 32, 3)
	
	outputs = []
	
	if outVectors:
		outputs.append(coordsOutPut(x))
	if outClasses:
		outputs.append(classOutput(x))
	
	return tf.keras.Model(inputs = xIn, outputs = outputs, name = modelName)

def stvNetNew(inputShape = (480, 640, 3), outVectors = True, outClasses = True, modelName = "stvNetNew"):
	
	xIn = tf.keras.Input(inputShape, dtype = np.dtype('uint8'))
	
	x = tf.keras.layers.Lambda(lambda x: x / 255) (xIn)
	
	x = tf.keras.layers.Conv2D(64, 7, input_shape = inputShape, kernel_initializer = tf.keras.initializers.GlorotUniform(seed=0), padding = 'same')(x)
	x = tf.keras.layers.BatchNormalization()(x)
	x = tf.keras.layers.Activation('relu')(x)
	
	res1 = x
	
	x = tf.keras.layers.MaxPool2D(pool_size = 3, strides = 2, padding = 'same')(x)
	
	skip = x
	
	x = convLayer(x, 64, 3)
	x = convLayer(x, 64, 3)
	
	x = tf.keras.layers.Add()([x, skip])
	skip = x
	
	x = convLayer(x, 64, 3)
	x = convLayer(x, 64, 3)
	
	x = tf.keras.layers.Add()([x, skip])
	skip = tf.keras.layers.MaxPool2D(pool_size = 2, padding = 'same')(x)
	skip = tf.pad(skip, [[0, 0], [0, 0], [0, 0], [32, 32]]) # linear projection
	res2 = x
	
	x = convLayer(x, 128, 3, 2)
	x = convLayer(x, 128, 3)
	
	x = tf.keras.layers.Add()([x, skip])
	skip = x
	
	x = convLayer(x, 128, 3)
	x = convLayer(x, 128, 3)
	
	x = tf.keras.layers.Add()([x, skip])
	skip = tf.keras.layers.MaxPool2D(pool_size = 2, padding = 'same')(x)
	skip = tf.pad(skip, [[0, 0], [0, 0], [0, 0], [64, 64]]) # linear projection
	res3 = x
	
	x = convLayer(x, 256, 3, 2)
	x = convLayer(x, 256, 3)
	
	x = tf.keras.layers.Add()([x, skip])
	skip = x
	
	x = convLayer(x, 256, 3)
	x = convLayer(x, 256, 3)
	
	x = tf.keras.layers.Add()([x, skip])
	skip = tf.pad(x, [[0, 0], [0, 0], [0, 0], [128, 128]])
	res4 = x
	
	x = convLayer(x, 512, 3, dilation = 2)
	x = convLayer(x, 512, 3, dilation = 2)
	
	x = tf.keras.layers.Add()([x, skip])
	skip = x
	
	x = convLayer(x, 512, 3, dilation = 2)
	x = convLayer(x, 512, 3, dilation = 2)
	
	x = tf.keras.layers.Add()([x, skip])
	
	x = convLayer(x, 256, 3)
	x = convLayer(x, 256, 3)
	
	x = tf.keras.layers.Add()([x, res4])
	x = tf.keras.layers.UpSampling2D()(x)
	
	x = convLayer(x, 128, 3)
	
	x = tf.keras.layers.Add()([x, res3])
	x = tf.keras.layers.UpSampling2D()(x)
	
	x = convLayer(x, 64, 3)
	
	x = tf.keras.layers.Add()([x, res2])
	x = tf.keras.layers.UpSampling2D()(x)
	
	x = tf.keras.layers.Add()([x, res1])
	
	x = convLayer(x, 32, 3)
	
	outputs = []
	
	if outVectors:
		outputs.append(coordsOutPut(x))
	if outClasses:
		outputs.append(classOutput(x))
	
	return tf.keras.Model(inputs = xIn, outputs = outputs, name = modelName)
	
def uNet(inputShape = (480, 640, 3), outVectors = True, outClasses = True, modelName = "uNet"): # neural net structure used for image segmentation
	xIn = tf.keras.Input(inputShape, dtype = np.dtype('uint8'))
	
	x = tf.keras.layers.Lambda(lambda x: x / 255) (xIn)
	
	c1 = tf.keras.layers.Conv2D(16, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (x)
	c1 = tf.keras.layers.Dropout(0.1) (c1)
	c1 = tf.keras.layers.Conv2D(16, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c1)
	p1 = tf.keras.layers.MaxPool2D((2, 2)) (c1)

	c2 = tf.keras.layers.Conv2D(32, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (p1)
	c2 = tf.keras.layers.Dropout(0.1) (c2)
	c2 = tf.keras.layers.Conv2D(32, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c2)
	p2 = tf.keras.layers.MaxPool2D((2, 2)) (c2)

	c3 = tf.keras.layers.Conv2D(64, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (p2)
	c3 = tf.keras.layers.Dropout(0.2) (c3)
	c3 = tf.keras.layers.Conv2D(64, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c3)
	p3 = tf.keras.layers.MaxPool2D((2, 2)) (c3)

	c4 = tf.keras.layers.Conv2D(128, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (p3)
	c4 = tf.keras.layers.Dropout(0.2) (c4)
	c4 = tf.keras.layers.Conv2D(128, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c4)
	p4 = tf.keras.layers.MaxPool2D(pool_size=(2, 2)) (c4)

	c5 = tf.keras.layers.Conv2D(256, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (p4)
	c5 = tf.keras.layers.Dropout(0.3) (c5)
	c5 = tf.keras.layers.Conv2D(256, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c5)

	u6 = tf.keras.layers.Conv2DTranspose(128, (2, 2), strides=(2, 2), padding='same') (c5)
	u6 = tf.keras.layers.concatenate([u6, c4])
	c6 = tf.keras.layers.Conv2D(128, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (u6)
	c6 = tf.keras.layers.Dropout(0.2) (c6)
	c6 = tf.keras.layers.Conv2D(128, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c6)

	u7 = tf.keras.layers.Conv2DTranspose(64, (2, 2), strides=(2, 2), padding='same') (c6)
	u7 = tf.keras.layers.concatenate([u7, c3])
	c7 = tf.keras.layers.Conv2D(64, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (u7)
	c7 = tf.keras.layers.Dropout(0.2) (c7)
	c7 = tf.keras.layers.Conv2D(64, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c7)

	u8 = tf.keras.layers.Conv2DTranspose(32, (2, 2), strides=(2, 2), padding='same') (c7)
	u8 = tf.keras.layers.concatenate([u8, c2])
	c8 = tf.keras.layers.Conv2D(32, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (u8)
	c8 = tf.keras.layers.Dropout(0.1) (c8)
	c8 = tf.keras.layers.Conv2D(32, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c8)

	u9 = tf.keras.layers.Conv2DTranspose(16, (2, 2), strides=(2, 2), padding='same') (c8)
	u9 = tf.keras.layers.concatenate([u9, c1], axis=3)
	c9 = tf.keras.layers.Conv2D(16, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (u9)
	c9 = tf.keras.layers.Dropout(0.1) (c9)
	c9 = tf.keras.layers.Conv2D(16, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c9)

	outputs = []
	if outVectors:
		#outputs.append(tf.keras.layers.Conv2D(18, (1, 1), activation='sigmoid') (c9))
		outputs.append(tf.keras.layers.Conv2D(18, (1,1), kernel_initializer = tf.keras.initializers.GlorotUniform(seed=0), padding = 'same') (c9))
		#outputs.append(coordsOutPut(c9))
	if outClasses:
		outputs.append(tf.keras.layers.Conv2D(1, (1, 1), activation='sigmoid') (c9))
		#outputs.append(classOutput(c9))
		#outputs.append(tf.keras.layers.Conv2D(1, (1,1), name = 'classConv', kernel_initializer = tf.keras.initializers.GlorotUniform(seed=0), padding = 'same', activation = tf.keras.layers.Activation('relu')) (c9))
	
	return tf.keras.Model(inputs = [xIn], outputs = outputs, name = modelName)
	
def trainModel(modelStruct, modelGen, modelClass = 'cat', batchSize = 2, optimizer = tf.keras.optimizers.Adam, learning_rate = 0.01, losses = None, metrics = ['accuracy'], saveModel = True, modelName = 'stvNet_weights', epochs = 1, loss_weights = None, outVectors = False, outClasses = False, dataSplit = True, altLabels = True, augmentation = True): # train and save model weights
	if not (outVectors or outClasses):
		print("At least one of outVectors or outClasses must be set to True.")
		return
	model = modelStruct(outVectors = outVectors, outClasses = outClasses, modelName = modelName)
	#model.summary()
	model.compile(optimizer = optimizer(learning_rate = learning_rate), loss = losses, metrics = metrics, loss_weights = loss_weights)
	
	trainData, validData = None, None
	if dataSplit: # if using datasplit, otherwise all available data is used
		trainData, validData = data.getDataSplit(modelClass = modelClass)
	
	logger = tf.keras.callbacks.CSVLogger("models\\history\\" + modelName + "_" + modelClass + "_history.csv", append = True)
	#evalLogger = tf.keras.callbacks.CSVLogger("models\\history\\" + modelName + "_" + modelClass + "_eval_history.csv", append = True)
	
	history, valHistory = [], []
	
	if type(losses) is dict:
		outKeys = list(losses.keys())
		if len(outKeys) == 2: # combined output
			for i in range(epochs):
				print("Epoch {0} of {1}".format(i + 1, epochs))
				hist = model.fit(modelGen(modelClass, batchSize, masterList = trainData, out0 = outKeys[0], out1 = outKeys[1], altLabels = altLabels, augmentation = augmentation), steps_per_epoch = math.ceil(len(trainData) / batchSize), max_queue_size = 2, callbacks = [logger])
				history.append(hist.history)
				if dataSplit:
					print("Validation:")
					valHist = model.evaluate(modelGen(modelClass, batchSize, masterList = validData, out0 = outKeys[0], out1 = outKeys[1], altLabels = altLabels, augmentation = False), steps = math.ceil(len(validData) / batchSize), max_queue_size = 2)
					valHistory.append(valHist)
		else:
			raise Exception("Probably shouldn't be here ever..")
	else:
		for i in range(epochs):
			print("Epoch {0} of {1}".format(i + 1, epochs))
			hist = model.fit(modelGen(modelClass, batchSize, masterList = trainData, altLabels = altLabels, augmentation = augmentation), steps_per_epoch = math.ceil(len(trainData) / batchSize), max_queue_size = 2, callbacks = [logger])
			history.append(hist.history)
			if dataSplit:
				print("Validation:")
				valHist = model.evaluate(modelGen(modelClass, batchSize, masterList = validData, altLabels = altLabels, augmentation = False), steps = math.ceil(len(validData) / batchSize), max_queue_size = 2)
				valHistory.append(valHist)
		
	historyLog = {"struct": modelStruct.__name__,
		"class" : modelClass,
		"optimizer": optimizer,
		"lr" : learning_rate,
		"losses": losses,
		"name": modelName,
		"epochs": epochs,
		"history": history,
		"evalHistory": valHistory,
		"timestamp": datetime.now().strftime("%d/%m/%Y %H:%M:%S"),
	}
	
	if saveModel:
		model.save_weights(os.path.dirname(os.path.realpath(__file__)) + '\\models\\' + modelName + '_' + modelClass)
		chunkify(os.path.dirname(os.path.realpath(__file__)) + '\\models\\' + modelName + '_' + modelClass)
		#model.save(os.path.dirname(os.path.realpath(__file__)) + '\\models\\' + modelName + '_' + modelClass)
		if not os.path.exists("models\\history\\" + modelName + '_trainHistory'):
			with open("models\\history\\" + modelName + '_' + modelClass + '_trainHistory', 'wb') as f: # create model history
				pickle.dump([], f)
		with open("models\\history\\" + modelName + '_' + modelClass + '_trainHistory', 'rb') as f: # loading old history 
			histories = pickle.load(f)
		histories.append(historyLog)
		with open("models\\history\\" + modelName + '_' + modelClass + '_trainHistory', 'wb') as f: # saving the history of the model
			pickle.dump(histories, f)
		
	return model

def chunkify(input_file, chunk_size = 104856576 ):
	with open(input_file + '.data-00000-of-00001', 'rb') as f:
		chunk_num = 0
		while True:
			chunk = f.read(chunk_size)
			if not chunk:
				break
			with open(f'{input_file}.part-{chunk_num}', 'wb') as chunk_file:
				chunk_file.write(chunk)
			chunk_num += 1

	os.remove(input_file + '.data-00000-of-00001')

def joinFiles(input_file, model):
	with open(input_file + '.data-00000-of-00001', 'wb') as combined:
		chunk_num = 0
		while True:
			try:
				with open(f'{input_file}.part-{chunk_num}', 'rb') as f:
					combined.write(f.read())
					chunk_num += 1
			except FileNotFoundError:
				break
	
	model.load_weights(input_file)
	os.remove(input_file + '.data-00000-of-00001')

def trainModels(modelSets, shutDown = False):
	for modelSet in modelSets:
		print("Training {0}".format(modelSet.name))
		model = modelsDict[modelSet.name]
		trainModel(model.structure, model.generator, modelClass = modelSet.modelClass, epochs = model.epochs, losses = model.losses, modelName = modelSet.name, outClasses = model.outClasses, outVectors = model.outVectors, learning_rate = model.lr, metrics = model.metrics, altLabels = model.altLabels, augmentation = model.augmentation)
		
		K.clear_session()
		K.reset_uids()
	
	if shutDown:
		os.system('shutdown -s')

def evaluateModel(modelStruct, modelName, evalGen, modelClass = 'cat', outVectors = False, outClasses = False, batchSize = 2, optimizer = tf.keras.optimizers.Adam, learning_rate = 0.01, losses = None, metrics = ['accuracy'], samples = 100): # test existing model performance
	model = tf.keras.models.load_model(os.path.dirname(os.path.realpath(__file__)) + '\\models\\' + modelName + '_' + modelClass)
	model.evaluate(evalGen(modelClass, batchSize), steps = samples // batchSize)
	
def evaluateModels(modelSets, batchSize = 2, dataSplit = True):
	
	for modelSet in modelSets:
		validData = (data.getDataSplit(modelClass = modelSet.modelClass)[1] if dataSplit else None)
		modelEnt = modelsDict[modelSet.name]
		model = loadModelWeights(modelEnt.structure, modelSet.name, modelSet.modelClass, modelEnt.outVectors, modelEnt.outClasses, losses = modelEnt.losses, metrics = modelEnt.metrics)
		if type(model.losses) is dict:
			outKeys = list(model.losses.keys())
			if len(outKeys) == 2: # combined output
				model.evaluate(modelEnt.generator(modelSet.modelClass, batchSize = batchSize, masterList = validData, out0 = outKeys[0], out1 = outKeys[1], altLabels = modelEnt.altLabels, augmentation = False), steps = math.ceil(len(validData) / batchSize), max_queue_size = 2)
			else:
				raise Exception("Probably shouldn't be here ever..")
		else:
			model.evaluate(modelEnt.generator(modelSet.modelClass, batchSize = batchSize, masterList = validData, altLabels = modelEnt.altLabels, augmentation = False), steps = math.ceil(len(validData) / batchSize), max_queue_size = 2)

def loadModelWeights(modelStruct, modelName, modelClass = 'cat', outVectors = False, outClasses = False, optimizer = tf.keras.optimizers.Adam, learning_rate = 0.01, losses = None, metrics = ['accuracy']): # return compiled tf keras model
	if not (outVectors or outClasses):
		raise Exception("At least one of outVectors or outClasses must be set to True.")
	model = modelStruct(outVectors = outVectors, outClasses = outClasses, modelName = modelName)
	joinFiles(os.path.dirname(os.path.realpath(__file__)) + '\\models\\' + modelName + '_' + modelClass, model)
	model.compile(optimizer = optimizer(learning_rate = learning_rate), loss = losses, metrics = metrics)
	return model

def loadHistory(modelName, modelClass = 'cat'):
	with open("models\\history\\" + modelName + '_' + modelClass + '_trainHistory', 'rb') as f: # loading old history 
		histories = pickle.load(f)
		for hist in histories:
			print("Structure: {0}\nClass: {1}\nOptimizer: {2}\nLearningRate: {3}\nLosses: {4}\nName: {5}\nEpochs: {6}\nTimestamp: {7}\nTraining History:\n".format(hist['struct'], hist['class'], hist['optimizer'], hist['lr'], hist['losses'], hist['name'], hist['epochs'], hist['timestamp']))
			for i, epoch in enumerate(hist['history']):
				print("{0}: {1}".format(i, epoch))
			print("\nEvaluation History:\n")
			for i, epoch in enumerate(hist['evalHistory']):
				print("{0}: {1}".format(i, epoch))
			print("\n")

def loadHistories(modelSets):
	for modelSet in modelSets:
		print("Loading {0}".format(modelSet.name))
		loadHistory(modelSet.name, modelSet.modelClass)
		
def plotHistories(modelSets): # display loss values over epochs using pyplot
	plt.figure()
	maxLen = 0
	for modelSet in modelSets:
		with open("models\\history\\" + modelSet.name + '_' + modelSet.modelClass + '_trainHistory', 'rb') as f: # loading old history 
			histories = pickle.load(f)
		for hist in histories:
			if len(hist['history']) > maxLen:
				maxLen = len(hist['history'])
			plt.subplot(211)
			plt.plot([x['loss'] for x in hist['history']], label = hist['name'])
			plt.subplot(212)
			plt.plot([x[0] for x in hist['evalHistory']], label = hist['name'])
	plt.subplot(211)
	plt.ylabel("Training Loss")
	plt.xlabel("Epoch")
	plt.xticks(np.arange(0, maxLen, 1.0))
	plt.subplot(212)
	plt.ylabel("Validation Loss")
	plt.xlabel("Epoch")
	plt.xticks(np.arange(0, maxLen, 1.0))
	plt.legend()
	plt.show()
	plt.close()
			
modelsDict = {
	'uNet_classes' : modelDictVal(uNet, data.classTrainingGenerator, tf.keras.losses.BinaryCrossentropy(), False, True, epochs = 20, lr = 0.001, augmentation = False),
	'uNet_coords' : modelDictVal(uNet, data.coordsTrainingGenerator, tf.keras.losses.Huber(), True, False, epochs = 5, lr = 0.001, metrics = ['mae', 'mse']),
	'uNet_coords_smooth' : modelDictVal(uNet, data.coordsTrainingGenerator, smoothL1, True, False, epochs = 3, lr = 0.0001, metrics = ['mae', 'mse']),
	'stvNet' : modelDictVal(stvNet, data.combinedTrainingGenerator, {'coordsOut': tf.keras.losses.Huber(), 'classOut': tf.keras.losses.BinaryCrossentropy()}, True, True, epochs = 5, lr = 0.00005, metrics = {'coordsOut': ['mae', 'mse'], "classOut": ['accuracy']}),
	'stvNet_coords_slow_learner' : modelDictVal(stvNet, data.coordsTrainingGenerator, tf.keras.losses.Huber(), True, False, epochs = 40, lr = 0.00001, metrics = ['mae', 'mse'], outVecName = 'coordsOut'),
	'stvNetAltLabels' : modelDictVal(stvNet, data.combinedTrainingGenerator, {'coordsOut': tf.keras.losses.Huber(), 'classOut': tf.keras.losses.BinaryCrossentropy()}, True, True, epochs = 10, lr = 0.001, metrics = {'coordsOut': ['mae', 'mse'], "classOut": ['accuracy']}, altLabels = True, augmentation = True),
	'stvNetNormLabels' : modelDictVal(stvNet, data.combinedTrainingGenerator, {'coordsOut': tf.keras.losses.Huber(), 'classOut': tf.keras.losses.BinaryCrossentropy()}, True, True, epochs = 10, lr = 0.001, metrics = {'coordsOut': ['mae', 'mse'], "classOut": ['accuracy']}, altLabels = False, augmentation = True),
	'stvNet_coords' : modelDictVal(stvNet, data.coordsTrainingGenerator, tf.keras.losses.Huber(), True, False, epochs = 2, lr = 0.001, metrics = ['mae', 'mse'], altLabels = False, augmentation = True),
	'stvNet_coords_altLabels' : modelDictVal(stvNet, data.coordsTrainingGenerator, tf.keras.losses.Huber(), True, False, epochs = 20, lr = 0.001, metrics = ['mae', 'mse'], altLabels = True, augmentation = True),
	'stvNet_coords_altLabels_noAug' : modelDictVal(stvNet, data.coordsTrainingGenerator, tf.keras.losses.Huber(), True, False, epochs = 20, lr = 0.001, metrics = ['mae', 'mse'], altLabels = True, augmentation = False),
	'stvNet_coords_noAug' : modelDictVal(stvNet, data.coordsTrainingGenerator, tf.keras.losses.Huber(), True, False, epochs = 20, lr = 0.001, metrics = ['mae', 'mse'], altLabels = False, augmentation = False),
	'stvNet_classes' : modelDictVal(stvNet, data.classTrainingGenerator, tf.keras.losses.BinaryCrossentropy(), False, True, epochs = 10, lr = 0.001, altLabels = False, augmentation = True),
	'stvNet_classes_noAug' : modelDictVal(stvNet, data.classTrainingGenerator, tf.keras.losses.BinaryCrossentropy(), False, True, epochs = 10, lr = 0.001, altLabels = False, augmentation = False),
	'stvNet_new_coords_alt' : modelDictVal(stvNetNew, data.coordsTrainingGenerator, tf.keras.losses.Huber(), True, False, epochs = 20, lr = 0.001, metrics = ['mae', 'mse'], altLabels = True, augmentation = False),
	'stvNet_new_coords' : modelDictVal(stvNetNew, data.coordsTrainingGenerator, tf.keras.losses.Huber(), True, False, epochs = 10, lr = 0.001, metrics = ['mae', 'mse'], altLabels = False, augmentation = False),
	'stvNet_new_coords_alt_aug' : modelDictVal(stvNetNew, data.coordsTrainingGenerator, tf.keras.losses.Huber(), True, False, epochs = 20, lr = 0.001, metrics = ['mae', 'mse'], altLabels = True, augmentation = True),
	'stvNet_new_coords_aug' : modelDictVal(stvNetNew, data.coordsTrainingGenerator, tf.keras.losses.Huber(), True, False, epochs = 20, lr = 0.001, metrics = ['mae', 'mse'], altLabels = False, augmentation = True),
	'stvNet_new_classes' : modelDictVal(stvNetNew, data.classTrainingGenerator, tf.keras.losses.BinaryCrossentropy(), False, True, epochs = 20, lr = 0.001, augmentation = False),
	'stvNet_new_combined' : modelDictVal(stvNetNew, data.combinedTrainingGenerator, {'coordsOut': tf.keras.losses.Huber(), 'classOut': tf.keras.losses.BinaryCrossentropy()}, True, True, epochs = 20, lr = 0.001, metrics = {'coordsOut': ['mae', 'mse'], "classOut": ['accuracy']}, augmentation = False),
}
	
if __name__ == "__main__" :
	modelSets = [modelSet('stvNet_new_coords')]
	trainModels(modelSets)
	
	#evaluateModels(modelSets)
	#loadHistories(modelSets)
	#plotHistories(modelSets)