diff --git a/README.md b/README.md index 04b4ec9..1b86967 100644 --- a/README.md +++ b/README.md @@ -3,26 +3,30 @@ This project uses TensorFlow2.0 for image classification tasks. ## How to use ### Requirements -+ **Python 3.x** (My Python version is 3.6.8)
-+ **TensorFlow version: 2.0.0-beta1**
++ **Python 3.x** (My Python version is 3.8.0)
++ **TensorFlow version: 2.11.0
+ The file directory of the dataset should look like this:
``` ${dataset_root} |——train -| |——class_name_0 -| |——class_name_1 -| |——class_name_2 -| |——class_name_3 +| |——class_dir_0 +| | |——image_1.jpg +| | |——image_2.jpg +| | |——image_3.jpg +| | ... +| |——class_dir_1 +| |——class_dir_2 +| |——class_dir_3 |——valid -| |——class_name_0 -| |——class_name_1 -| |——class_name_2 -| |——class_name_3 +| |——class_dir_0 +| |——class_dir_1 +| |——class_dir_2 +| |——class_dir_3 |——test - |——class_name_0 - |——class_name_1 - |——class_name_2 - |——class_name_3 + |——class_dir_0 + |——class_dir_1 + |——class_dir_2 + |——class_dir_3 ``` ### Train @@ -40,3 +44,4 @@ The structure of the network is defined in `model_definition.py`, you can change ## References 1. AlexNet : http://www.cs.toronto.edu/~fritz/absps/imagenet.pdf 2. VGG : https://arxiv.org/abs/1409.1556 +3. Keras : https://keras.io/api/applications/ diff --git a/config.py b/config.py index b228054..5ebdd0c 100644 --- a/config.py +++ b/config.py @@ -1,12 +1,15 @@ # some training parameters -EPOCHS = 50 -BATCH_SIZE = 8 -NUM_CLASSES = 5 -image_height = 224 -image_width = 224 +EPOCHS = 100 +BATCH_SIZE = 128 +NUM_CLASSES = 6 +image_height = 128 +image_width = 128 channels = 3 -model_dir = "image_classification_model.h5" -train_dir = "dataset/train" -valid_dir = "dataset/valid" -test_dir = "dataset/test" -test_image_path = "" \ No newline at end of file + +model_save_name = "EfficientNetB2" +model_dir = "trained_models/salmon_crop_128/"+model_save_name+"/" # = save_dir + +train_dir = "/home/mirap/0_DATABASE/IMAS_Salmon/7_SalmonTest/train" +valid_dir = "/home/mirap/0_DATABASE/IMAS_Salmon/7_SalmonTest/valid" +test_dir = "/home/mirap/0_DATABASE/IMAS_Salmon/7_SalmonTest/test" +test_image_path = "/home/mirap/0_DATABASE/IMAS_Salmon/7_SalmonTest/test/5/5_108.jpg" diff --git a/evaluate.py b/evaluate.py index 221ba4b..d255ab8 100644 --- a/evaluate.py +++ b/evaluate.py @@ -1,14 +1,32 @@ import tensorflow as tf import config from prepare_data import get_datasets +from sklearn.metrics import classification_report +import numpy as np -train_generator, valid_generator, test_generator, \ -train_num, valid_num, test_num= get_datasets() +def eval_model(new_model): + # Load data + train_generator, valid_generator, test_generator, \ + train_num, valid_num, test_num= get_datasets() -# Load the model -new_model = tf.keras.models.load_model(config.model_dir) -# Get the accuracy on the test set -loss, acc = new_model.evaluate_generator(test_generator, - steps=test_num // config.BATCH_SIZE) -print("The accuracy on test set is: {:6.3f}%".format(acc*100)) \ No newline at end of file + # Get the accuracy on the test set + loss, acc, auc, precision, recall = new_model.evaluate(test_generator, + batch_size=config.BATCH_SIZE, + steps=test_num // config.BATCH_SIZE) + print("result of ",config.model_dir) + print("The accuracy on test set is: {:6.3f}%".format(acc*100)) + print("The auc on test set is: {:6.3f}%".format(auc*100)) + print("The precision on test set is: {:6.3f}%".format(precision*100)) + print("The recall on test set is: {:6.3f}%".format(recall*100)) + + # Evaluate per class + lables_array = test_generator.classes + predictions = new_model.predict(test_generator) + predictions = np.argmax(predictions, axis=1) + print(classification_report(lables_array, predictions)) + +if __name__ == '__main__': + # Load the model + new_model = tf.keras.models.load_model(config.model_dir+config.model_save_name+".h5") + eval_model(new_model) \ No newline at end of file diff --git a/log/README.md b/log/README.md deleted file mode 100644 index 2128441..0000000 --- a/log/README.md +++ /dev/null @@ -1 +0,0 @@ -This is the log dir. \ No newline at end of file diff --git a/models/pretrained_models.py b/models/pretrained_models.py new file mode 100755 index 0000000..b00c9fe --- /dev/null +++ b/models/pretrained_models.py @@ -0,0 +1,328 @@ + +import tensorflow as tf +import config +from models.testnet import TestNet + +if config.NUM_CLASSES == 2: + ACTIVATION = "sigmoid" +else: + ACTIVATION = "softmax" + +def pretrained_model(model_name, load_weight="imagenet"): + # TestNet + if model_name == "TestNet": + base_model = TestNet() + # MobileNetV2 + if model_name == "MobileNetV2": + base_model = tf.keras.applications.MobileNetV2( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION + ) + # Xception (2017) + if model_name == "Xception": + base_model = tf.keras.applications.Xception( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION + ) + + # EfficientNetB0~B7 (2019) + if model_name == "EfficientNetB0": + base_model = tf.keras.applications.EfficientNetB0( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION + ) + + if model_name == "EfficientNetB1": + base_model = tf.keras.applications.EfficientNetB1( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION + ) + + if model_name == "EfficientNetB2": + base_model = tf.keras.applications.EfficientNetB2( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION + ) + + if model_name == "EfficientNetB3": + base_model = tf.keras.applications.EfficientNetB3( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION + ) + + if model_name == "EfficientNetB4": + base_model = tf.keras.applications.EfficientNetB4( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION + ) + + if model_name == "EfficientNetB5": + base_model = tf.keras.applications.EfficientNetB5( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION + ) + + if model_name == "EfficientNetB6": + base_model = tf.keras.applications.EfficientNetB6( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION + ) + + if model_name == "EfficientNetB7": + base_model = tf.keras.applications.EfficientNetB7( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION + ) + + # EfficientNetV2 B0 to B3 and S, M, L (2021) + if model_name == "EfficientNetV2B0": + base_model = tf.keras.applications.EfficientNetV2B0( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION, + include_preprocessing=True, + ) + + if model_name == "EfficientNetV2B1": + base_model = tf.keras.applications.EfficientNetV2B1( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION, + include_preprocessing=True, + ) + + if model_name == "EfficientNetV2B2": + base_model = tf.keras.applications.EfficientNetV2B2( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION, + include_preprocessing=True, + ) + + if model_name == "EfficientNetV2B3": + base_model = tf.keras.applications.EfficientNetV2B3( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION, + include_preprocessing=True, + ) + + if model_name == "EfficientNetV2S": + base_model = tf.keras.applications.EfficientNetV2S( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION, + include_preprocessing=True, + ) + + if model_name == "EfficientNetV2M": + base_model = tf.keras.applications.EfficientNetV2M( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION, + include_preprocessing=True, + ) + + if model_name == "EfficientNetV2L": + base_model = tf.keras.applications.EfficientNetV2L( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION, + include_preprocessing=True, + ) + + # VGG series (2015) + if model_name == "VGG16": + base_model = tf.keras.applications.VGG16( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION, + ) + + if model_name == "VGG19": + base_model = tf.keras.applications.VGG19( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION, + ) + + # DenseNet Series (2017) + if model_name == "DenseNet121": + base_model = tf.keras.applications.DenseNet121( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION + ) + + if model_name == "DenseNet169": + base_model = tf.keras.applications.DenseNet169( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION + ) + + if model_name == "DenseNet201": + base_model = tf.keras.applications.DenseNet201( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION + ) + + # NasNet Series (2018) + if model_name == "NASNetLarge": + base_model = tf.keras.applications.NASNetLarge( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION + ) + + if model_name == "NASNetMobile": + base_model = tf.keras.applications.NASNetMobile( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION + ) + + # InceptionV3 (2016) + if model_name == "InceptionV3": + base_model = tf.keras.applications.InceptionV3( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION, + ) + + # InceptionResNetV2 (2016) + if model_name == "InceptionResNetV2": + base_model = tf.keras.applications.InceptionResNetV2( + include_top=True, + weights=load_weight, + input_tensor=None, + input_shape=(config.image_width,config.image_height,3), + pooling=None, + classes=1000, + classifier_activation=ACTIVATION, + ) + + + x = base_model.output + # x = GlobalAveragePooling2D()(x) + # add a fully-connected layer + x = tf.keras.layers.Dense(1000, activation='relu')(x) + # add logistic layer (complusory for predict classes) + predictions = tf.keras.layers.Dense(config.NUM_CLASSES, activation=ACTIVATION)(x) + + # this is the model we will train + model = tf.keras.models.Model(inputs=base_model.input, outputs=predictions) + + + return model diff --git a/models/testnet.py b/models/testnet.py new file mode 100755 index 0000000..75bf09e --- /dev/null +++ b/models/testnet.py @@ -0,0 +1,39 @@ +import tensorflow as tf +from config import NUM_CLASSES, image_width, image_height, channels + +def TestNet(): + if NUM_CLASSES == 2 : + activation = "sigmoid" + else: + activation = "softmax" + + model = tf.keras.Sequential([ + # layer 1 + tf.keras.layers.Conv2D(filters=4, + kernel_size=(3, 3), + strides=2, + padding="same", + activation=tf.keras.activations.relu, + input_shape=(image_height, image_width, channels)), + #tf.keras.layers.Dropout(rate=0.2), + tf.keras.layers.AveragePooling2D(pool_size=(2, 2), + strides=2, + padding="valid"), + tf.keras.layers.Conv2D(filters=4, + kernel_size=(3, 3), + strides=2, + padding="same", + activation=tf.keras.activations.relu), + tf.keras.layers.AveragePooling2D(pool_size=(2, 2), + strides=2, + padding="valid"), + tf.keras.layers.BatchNormalization(), + + # layer 2 + tf.keras.layers.Flatten(), + tf.keras.layers.Dense(units=NUM_CLASSES, + activation=activation), + # tf.keras.layers.Dropout(rate=0.5), + ]) + + return model \ No newline at end of file diff --git a/prepare_data.py b/prepare_data.py index 821080c..202f86d 100644 --- a/prepare_data.py +++ b/prepare_data.py @@ -4,7 +4,10 @@ def get_datasets(): # Preprocess the dataset train_datagen = tf.keras.preprocessing.image.ImageDataGenerator( - rescale=1.0 / 255.0 + zoom_range=0.5, + horizontal_flip=True, + vertical_flip=True, + # rescale=1.0 / 255.0 ) train_generator = train_datagen.flow_from_directory(config.train_dir, @@ -16,7 +19,10 @@ def get_datasets(): class_mode="categorical") valid_datagen = tf.keras.preprocessing.image.ImageDataGenerator( - rescale=1.0 /255.0 + zoom_range=0.5, + horizontal_flip=True, + vertical_flip=True, + # rescale=1.0 /255.0 ) valid_generator = valid_datagen.flow_from_directory(config.valid_dir, target_size=(config.image_height, config.image_width), @@ -27,7 +33,7 @@ def get_datasets(): class_mode="categorical" ) test_datagen = tf.keras.preprocessing.image.ImageDataGenerator( - rescale=1.0 /255.0 + # rescale=1.0 /255.0 ) test_generator = test_datagen.flow_from_directory(config.test_dir, target_size=(config.image_height, config.image_width), diff --git a/test_single_image.py b/test_single_image.py index 2004994..7ef51d5 100644 --- a/test_single_image.py +++ b/test_single_image.py @@ -1,9 +1,9 @@ import tensorflow as tf import config import numpy as np +import os - -def test_single_image(img_dir): +def test_single_image(img_dir, model): img_raw = tf.io.read_file(img_dir) img_tensor = tf.image.decode_jpeg(img_raw, channels=config.channels) img_tensor = tf.image.resize(img_tensor, [config.image_height, config.image_width]) @@ -11,16 +11,18 @@ def test_single_image(img_dir): img_numpy = img_tensor.numpy() img_numpy = (np.expand_dims(img_numpy, 0)) img_tensor = tf.convert_to_tensor(img_numpy, tf.float32) - # print(img_tensor.shape) - img = img_tensor / 255.0 - prob = model(img) + + # img_tensor = img_tensor / 255.0 # uncomment if model included rescale preprocessing layer + prob = model(tf.image.resize(img_tensor,[config.image_width,config.image_height])) # print(prob) + classification = np.argmax(prob) return classification if __name__ == '__main__': - model = tf.keras.models.load_model(config.model_dir) - classification = test_single_image(config.test_image_path) - print(classification) \ No newline at end of file + model = tf.keras.models.load_model(config.model_dir+config.model_save_name+".h5") + test_image = config.test_image_path + classification = test_single_image(test_image, model) + diff --git a/train.py b/train.py old mode 100644 new mode 100755 index 77d818d..8703284 --- a/train.py +++ b/train.py @@ -1,19 +1,39 @@ from __future__ import absolute_import, division, print_function import tensorflow as tf -from config import EPOCHS, BATCH_SIZE, model_dir +import pandas as pd +import config +from test_single_image import test_single_image +from evaluate import eval_model from prepare_data import get_datasets -from models.alexnet import AlexNet -from models.vgg16 import VGG16 -from models.vgg19 import VGG19 +from models.pretrained_models import pretrained_model + +# USAGE: python train.py +#((before run! please set config.py file and line 26 of train.py)) + +available_models=["Xception","TestNet", + "EfficientNetB0", "EfficientNetB1", "EfficientNetB2", + "EfficientNetB3", "EfficientNetB4", "EfficientNetB5", + "EfficientNetB6", "EfficientNetB7", + "EfficientNetV2B0", "EfficientNetV2B1", + "EfficientNetV2B2", "EfficientNetV2B3", + "EfficientNetV2S", "EfficientNetV2M", "EfficientNetV2L", + "VGG16","VGG19", + "DenseNet121", "DenseNet169", "DenseNet201", + "NASNetLarge","NASNetMobile", + "InceptionV3","InceptionResNetV2" + ] def get_model(): - # model = AlexNet() - model = VGG16() - # model = VGG19() + model = pretrained_model(model_name="TestNet", + load_weight=None) model.compile(loss=tf.keras.losses.categorical_crossentropy, - optimizer=tf.keras.optimizers.Adam(learning_rate=0.001), - metrics=['accuracy']) + optimizer=tf.keras.optimizers.Adam(learning_rate=0.005), + metrics=['accuracy', # add more metrics if you want + tf.keras.metrics.AUC(), + tf.keras.metrics.Precision(), + tf.keras.metrics.Recall(), + ]) return model if __name__ == '__main__': @@ -28,19 +48,45 @@ def get_model(): # Use command tensorboard --logdir "log" to start tensorboard tensorboard = tf.keras.callbacks.TensorBoard(log_dir='log') - callback_list = [tensorboard] + + model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint( + filepath=config.model_dir, + save_weights_only=True, + monitor='val_accuracy', + mode='max', + save_best_only=True + ) + + early_stop_callback = tf.keras.callbacks.EarlyStopping( + monitor="loss", + patience=20, + restore_best_weights=True + ) + callback_list = [tensorboard, model_checkpoint_callback, early_stop_callback] model = get_model() - model.summary() + #model.summary() # start training - model.fit_generator(train_generator, - epochs=EPOCHS, - steps_per_epoch=train_num // BATCH_SIZE, + history = model.fit(train_generator, + epochs=config.EPOCHS, + steps_per_epoch=train_num // config.BATCH_SIZE, validation_data=valid_generator, - validation_steps=valid_num // BATCH_SIZE, + validation_steps=valid_num // config.BATCH_SIZE, callbacks=callback_list) # save the whole model - model.save(model_dir) + model.save(config.model_dir+config.model_save_name+".h5") + + #write histry + hist_df = pd.DataFrame(history.history) + with open(config.model_dir+"train_history.csv", mode='w') as f: + hist_df.to_csv(f) + + # Evaluation + eval_model(model) + + # detect for samples + test_single_image(config.test_image_path, model) + print("end of training!!!") \ No newline at end of file