This module takes a dataset in which the images are in labeled folders and trains a Keras neural network for image recognition.
TrainModel
TrainModel (data_dir)
Initialize self. See help(type(self)) for accurate signature.
TrainModel.load_data
TrainModel.load_data (img_height, img_width, batch_size)
Loads data from the directory provided in data_dir
TrainModel.train
TrainModel.train (img_height, img_width, epochs, optim_choice, batch_size, model_path)
Trains a new neural network model.
Args: img_height (int): image pixel height img_width (int): image pixel width epochs (int): Number of epochs to train optim_choice (string): Loss function to be used
Returns: keras_model, statistics
TrainModel.prediction
TrainModel.prediction (model, class_names:list)
Predicts on the image provided in the path. Args: model (tflite model): tflite model to be used in the prediction
Returns: img: image predicted, result: formatted string for the result
TrainModel.plot_statistics
TrainModel.plot_statistics (history, epochs_range)
Plot model training statistics.
Args: history (tuple?): tuple containing loss and accuracy values over training epochs_range (int): amount of epochs used to train over
Returns: BytesIO buffer: Matplotlib figure containing graphs about the training process
TrainModel.continue_training
TrainModel.continue_training (img_height, img_width, epochs, batch_size, model_path)
Takes an excisting model and further trains it with more images.
Args: img_height (int): image pixel height img_width (int): image pixel width epochs (int): Number of epochs to train optim_choice (string): Loss function to be used batch_size (int): Batch size to be used model_path (string): Location for fetching the model
Returns: keras_model, statistics