CNN Tensorflow

The following code will be builg for using tensorflow keras to train a simple state-of-the-art CNN model in python.

Training

Import library

main.py


# Use tensorflow 2.13
import time
import tensorflow as tf
from tensorflow import keras
import matplotlib.pyplot as plt
import logging

Setup logging


t = time.time()
current_time = int(t * 1000)
model_name = str(current_time) + "_isStairModelMobieNet.h5"
log_name = str(current_time) + "_isStairModel.log"
 
logging.basicConfig(
    level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s', 
    datefmt='%Y-%m-%d %H:%M', 
    handlers=[logging.FileHandler(log_name, 'w', 'utf-8'), ]
)

Import dataset from folders

main.py


data_dir = "G:\\dataset\\someData"
 
batch_size = 128
imgSize = 224
split_rate = 0.25
dataSeed = 1238
 
train_ds = tf.keras.preprocessing.image_dataset_from_directory(
  data_dir, seed=dataSeed, batch_size=batch_size, validation_split=split_rate, subset="training",
  image_size=(imgSize, imgSize), color_mode='rgb'
)
 
valid_ds = tf.keras.preprocessing.image_dataset_from_directory(
  data_dir, seed=dataSeed, batch_size=batch_size, validation_split=split_rate, subset="validation",
  image_size=(imgSize, imgSize), color_mode='rgb'
)
 
classNum = len(train_ds.class_names)
classArr = train_ds.class_names
print(classNum)
print(train_ds.class_names)

Setup data augmentation

main.py


randomAugSeed = 123
 
data_augmentation = keras.Sequential([
  tf.keras.layers.RandomFlip( "horizontal", seed=randomAugSeed ),
  tf.keras.layers.RandomRotation(0.1, seed=randomAugSeed ),
  tf.keras.layers.RandomZoom(0.1, seed=randomAugSeed ),
  tf.keras.layers.RandomContrast(0.1, seed=randomAugSeed)
])
 
expendRound = 1
temp_ds = train_ds
for i in range(expendRound):
    train_ds = train_ds.concatenate(temp_ds)
 
train_ds = train_ds.map(lambda image,label:(data_augmentation(image),label))

Setup model

main.py


pretrained_model = tf.keras.applications.MobileNetV3Large(
    include_top=False,
    weights='imagenet',
    input_shape=(imgSize, imgSize, 3)
)
 
pretrained_model.trainable = True
 
model = tf.keras.Sequential([
    pretrained_model,
    tf.keras.layers.GlobalAveragePooling2D(),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(classNum, activation='softmax')
])

Setup callback and epochs, lr

main.py


epochsRound = 14
base_learning_rate = 0.0001
 
# save model callback
checkpoint_filepath = './tmp/checkpoint'
model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_filepath,
    save_weights_only=True,
    monitor='val_accuracy',
    mode='max',
    save_best_only=True
)
 
# set lr decay callback
def learnDecay(epoch, lr):
    if epoch < (epochsRound / 2):
        return lr
    else:
        return lr * tf.math.exp(-0.15)
 
learnDecayCallback = tf.keras.callbacks.LearningRateScheduler(learnDecay)
 
class LogCallback(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        # keys = list(logs.keys())
        # print("End epoch {} of training; got log keys: {}".format(epoch, keys))
        logging.info(f"End epoch {epoch} of training:")
        logging.info(f"Loss: { logs['loss'] } ")
        logging.info(f"accuracy: { logs['accuracy'] } ")
        logging.info(f"val_loss: { logs['val_loss'] } ")
        logging.info(f"val_accuracy: { logs['val_accuracy'] } ")
        logging.info(f"---------------------------------------")

Training and save model

main.py


model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=base_learning_rate),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),
    metrics=['accuracy']
)
 
history = model.fit(
    train_ds,
    epochs=epochsRound,
    validation_data=valid_ds, 
    callbacks=[ model_checkpoint_callback, learnDecayCallback, LogCallback() ]
)
 
model.load_weights(checkpoint_filepath)
test_loss, test_acc = model.evaluate(valid_ds, verbose=2)
print(test_acc)
 
model.save(model_name)

Plot result

main.py


acc = history.history['accuracy']
val_acc = history.history['val_accuracy']
loss = history.history['loss']
val_loss = history.history['val_loss']
 
epochs_range = range(epochsRound)
 
plt.figure(figsize=(8, 8))
plt.subplot(1, 2, 1)
plt.plot(epochs_range, acc, label='Training Accuracy')
plt.plot(epochs_range, val_acc, label='Validation Accuracy')
plt.legend(loc='lower right')
plt.title('Training and Validation Accuracy')
 
plt.subplot(1, 2, 2)
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()

Full code

main.py


# Use tensorflow 2.13
import time
import tensorflow as tf
from tensorflow import keras
import matplotlib.pyplot as plt
import logging
 
t = time.time()
current_time = int(t * 1000)
model_name = str(current_time) + "_isStairModelMobieNet.h5"
log_name = str(current_time) + "_isStairModel.log"
 
logging.basicConfig(
    level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s', 
    datefmt='%Y-%m-%d %H:%M', 
    handlers=[logging.FileHandler(log_name, 'w', 'utf-8'), ]
)
 
data_dir = "G:\\dataset\\someData"
 
batch_size = 128
imgSize = 224
split_rate = 0.25
dataSeed = 1238
 
train_ds = tf.keras.preprocessing.image_dataset_from_directory(
  data_dir, seed=dataSeed, batch_size=batch_size, validation_split=split_rate, subset="training",
  image_size=(imgSize, imgSize), color_mode='rgb'
)
 
valid_ds = tf.keras.preprocessing.image_dataset_from_directory(
  data_dir, seed=dataSeed, batch_size=batch_size, validation_split=split_rate, subset="validation",
  image_size=(imgSize, imgSize), color_mode='rgb'
)
 
classNum = len(train_ds.class_names)
classArr = train_ds.class_names
print(classNum)
print(train_ds.class_names)
 
randomAugSeed = 123
 
data_augmentation = keras.Sequential([
  tf.keras.layers.RandomFlip( "horizontal", seed=randomAugSeed ),
  tf.keras.layers.RandomRotation(0.1, seed=randomAugSeed ),
  tf.keras.layers.RandomZoom(0.1, seed=randomAugSeed ),
  tf.keras.layers.RandomContrast(0.1, seed=randomAugSeed)
])
 
expendRound = 1
temp_ds = train_ds
for i in range(expendRound):
    train_ds = train_ds.concatenate(temp_ds)
 
train_ds = train_ds.map(lambda image,label:(data_augmentation(image),label))
 
pretrained_model = tf.keras.applications.MobileNetV3Large(
    include_top=False,
    weights='imagenet',
    input_shape=(imgSize, imgSize, 3)
)
 
pretrained_model.trainable = True
 
model = tf.keras.Sequential([
    pretrained_model,
    tf.keras.layers.GlobalAveragePooling2D(),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(classNum, activation='softmax')
])
 
epochsRound = 14
base_learning_rate = 0.0001
 
# save model callback
checkpoint_filepath = './tmp/checkpoint'
model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_filepath,
    save_weights_only=True,
    monitor='val_accuracy',
    mode='max',
    save_best_only=True
)
 
# set lr decay callback
def learnDecay(epoch, lr):
    if epoch < (epochsRound / 2):
        return lr
    else:
        return lr * tf.math.exp(-0.15)
 
learnDecayCallback = tf.keras.callbacks.LearningRateScheduler(learnDecay)
 
class LogCallback(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        # keys = list(logs.keys())
        # print("End epoch {} of training; got log keys: {}".format(epoch, keys))
        logging.info(f"End epoch {epoch} of training:")
        logging.info(f"Loss: { logs['loss'] } ")
        logging.info(f"accuracy: { logs['accuracy'] } ")
        logging.info(f"val_loss: { logs['val_loss'] } ")
        logging.info(f"val_accuracy: { logs['val_accuracy'] } ")
        logging.info(f"---------------------------------------")
 
# model.load_weights(checkpoint_filepath)
 
model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=base_learning_rate),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),
    metrics=['accuracy']
)
 
history = model.fit(
    train_ds,
    epochs=epochsRound,
    validation_data=valid_ds, 
    callbacks=[ model_checkpoint_callback, learnDecayCallback, LogCallback() ]
)
 
model.load_weights(checkpoint_filepath)
test_loss, test_acc = model.evaluate(valid_ds, verbose=2)
print(test_acc)
 
model.save(model_name)
 
acc = history.history['accuracy']
val_acc = history.history['val_accuracy']
loss = history.history['loss']
val_loss = history.history['val_loss']
 
epochs_range = range(epochsRound)
 
plt.figure(figsize=(8, 8))
plt.subplot(1, 2, 1)
plt.plot(epochs_range, acc, label='Training Accuracy')
plt.plot(epochs_range, val_acc, label='Validation Accuracy')
plt.legend(loc='lower right')
plt.title('Training and Validation Accuracy')
 
plt.subplot(1, 2, 2)
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()

Predict


import tensorflow as tf
import shutil
from pathlib import Path
 
import numpy as np
import os
 
def predictDataset(
    original_data_dir: str,
    model_path: str,
    output_data_dir: str,
    labels: list[str],
) -> None:
    """ Generate predict result with regarding model """
 
    batch_size = 64
    imgSize = 240
 
    model = tf.keras.models.load_model(model_path)
 
    first_stuff = Path(os.path.join(original_data_dir, os.listdir(original_data_dir)[0]))
    labels_mode = "inferred" if first_stuff.is_dir() else None
 
    valid_ds = tf.keras.preprocessing.image_dataset_from_directory(
        original_data_dir, labels=labels_mode, shuffle=False,
        image_size=(imgSize, imgSize), batch_size=batch_size, color_mode='rgb'
    )
 
    file_paths: list[str] = valid_ds.file_paths
 
    result = model.predict(
        valid_ds, batch_size=batch_size,
    )
 
    # [file_path, file predicted result]
    labels_list = [ [file_paths[ind], labels[np.argsort(v)[::-1][0]]] for ind, v in enumerate(result) ]
 
if __name__ == "__main__":
    
    model_dir = "E:\\dir\\pythonCode\\isStairModelv5.h5"
    original_data_dir = "G:\\dir\\original_data_dir"
    labels = ['isStair', 'notStair']
    
    output_data_dir = "G:\\dir\\output_data_dir"
 
    predictDataset(
        original_data_dir,
        model_dir,
        output_data_dir,
        labels
    )