model.py

# Reduce terminal logging verbosity
import os
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
print("Starting up...")

# Import libraries
print("\nImporting libraries...")
import tensorflow as tf
from tensorflow.python import keras, data
from tensorflow.python.keras import layers, losses
print("Tensorflow imported")
from keras import utils
from keras.layers import Rescaling
from keras.callbacks import History
print("Keras imported")
import pathlib
import matplotlib.pyplot as plt
print("Other dependencies imported")
print("All libraries imported\n")


# Environment Variables
dataPath = "../Generated Images"
testPath = "../Equalised Images"
modelPath = ".model"
print("Environment variables set\n")


# Class container for model
class CNN(tf.Module):
	def __init__(self, imgDim: tuple = (512, 512, 3)):
		self.model = keras.Sequential(
			[
				layers.Conv2D(8, (4, 4), padding = 'same', input_shape = imgDim, activation = 'relu'),
				Rescaling(1. / 255),
				layers.MaxPooling2D(pool_size = (2, 2)),
				layers.Conv2D(16, (4, 4), padding = 'same', activation = 'relu'),
				layers.MaxPooling2D(pool_size = (2, 2)),
				layers.Conv2D(8, (4, 4), padding = 'same', activation = 'relu'),
				layers.MaxPooling2D(pool_size = (2, 2)),
				layers.Dropout(0.2),
				layers.Flatten(),
				layers.Dense(64, activation = "relu"),
				layers.Dense(2)
			]
		)
		self.model.compile(
			optimizer = "adam",
			loss = losses.SparseCategoricalCrossentropy(from_logits = True)
		)

		self.trainData = None
		self.testData = None
		self.history = None
	
	def train(self, seed: int = 1542, epochs: int = 10):
		# Configuring datasets for better loading performance
		tuner = data.AUTOTUNE
		self.trainData = self.trainData.cache().shuffle(seed).prefetch(buffer_size = tuner)
		self.testData = self.testData.cache().prefetch(buffer_size = tuner)

		# Training CNN Model
		return self.model.fit(self.trainData,
			validation_data = self.testData,
			epochs = epochs)

	def trainSummary(self):
		loss = self.history.history['loss']
		val_loss = self.history.history['val_loss']

		epochs_range = self.history.epoch

		plt.figure(figsize=(8, 8))
		plt.plot(epochs_range, loss, label='Training Loss')
		plt.plot(epochs_range, val_loss, label='Validation Loss')
		plt.legend(loc='upper right')
		plt.title('Training and Validation Loss')
		plt.show()

	@tf.function
	def loadData(self, trainTestSplit: float = 0.35, seed: int = 1574):
		dataDir = pathlib.Path(dataPath)
		batchSize = 32

		# Creating dataset loaders
		trainData = utils.image_dataset_from_directory(dataDir,
						      validation_split = trainTestSplit,
							  subset = "training",
							  seed = seed,
							  image_size = (512, 512),
							  batch_size = batchSize)

		testData = utils.image_dataset_from_directory(dataDir,
						      validation_split = trainTestSplit,
							  subset = "validation",
							  seed = seed,
							  image_size = (512, 512),
							  batch_size = batchSize)

		return trainData, testData

	def rawPredict(self, testPath: str = "test"):
		testDir = pathlib.Path(testPath)
		testData = utils.image_dataset_from_directory(testDir,
							  validation_split = 0,
							  seed = 0,
							  image_size = (512, 512))

		return testData, self.model.predict(testData)

	def saveModel(self, modelFile: str = ".model"):
		self.model.save(modelFile)

	def loadModel(self, modelFile: str = ".model"):
		self.model.load_weights(modelFile)

	def debugInfo(self):
		print(self.model)
		for layer in self.model.layers:
			print(layer)
		
		self.model.summary()
		print(self.model.get_weights())


# Getting the OS currently running the model
from platform import system
currentOs = system()

# Helper functions
def clrscr():
	if currentOs == 'Linux':
		os.system("clear")
	else:
		os.system("cls")

def interpretPred(data, pred):
	[]


# Wrapper functions
def manualTrain():
	[]

def rawPredict():
	[]

def editEnvVars():
	[]


# Main menu function
def mainMenu() -> str:
	clrscr()	
	print("Main Menu\n")
	print("T -> Train the model")
	print("P -> Make a prediction")
	print("E -> Edit Environment Variables")
	print("Q -> Quit")

	ch = input("\nEnter your choice: ")
	ch = ch[0].lower()
	return ch


# Operation mode specified by launch arguments
# Available options:
# Refresh mode: Overwrite existing model files to start a fresh model
# Continue mode: Use existing model files to continue from previously saved model
# Model is automatically saved at the end of each program run
import argparse

parser = argparse.ArgumentParser(prog = "cnn-model", description = "CNN Windows EXE Classification Program")
parser.add_argument('-m', '--model-path', type = str, help = "Model files destination", dest = "modelPath", default = ".model")
parser.add_argument('-d', '--data-path', type = str, help = "Image Dataset path", dest = "dataPath", default = "../Generated Images")
parser.add_argument('-t', '--test-path', type = str, help = "Image Dataset for Testing path", dest = "testPath", default = "../Equalised Images")

# Mode Selection
progMode = parser.add_mutually_exclusive_group(required = True)
progMode.add_argument('-r', '--refresh', action = "store_true", help = "Run program in refresh mode", dest = "ref")
progMode.add_argument('-c', '--continue', action = "store_true", help = "Run program in continue mod", dest = "con")

args = parser.parse_args()

if args.ref:
	print("Initialising new Model...")
	model = CNN()
	print("Saving model (overwrites existing model files in path)...")
	model.saveModel(modelFile = args.modelPath)
	print("Model successfully saved!")
elif args.con:
	print("Creating Model...")
	model = CNN()
	print("Loading model values")
	model.loadModel(modelFile = args.modelPath)
	print("Model successfully loaded!")

input("\nPress enter to continue to main menu...")


# Main menu loop
ex = 'a'
while ex != 'q':
	ex = mainMenu()

model.saveModel(modelFile = args.modelPath)