train_model.py

# -*- coding: utf-8 -*-
"""classification_model.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1IFyPpQyAcLU_DDdue91KwMMQqFqixrUd
"""

import os
import math
import random
import shutil
import matplotlib.pyplot as plt
import tensorflow as tf
import tensorflow_hub as hub
from google.colab import drive
from tensorflow import keras
from tensorflow.keras import backend as K
from tensorflow.keras.backend import clear_session
from tensorflow.keras.callbacks import ModelCheckpoint, LearningRateScheduler, EarlyStopping, ReduceLROnPlateau
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.callbacks import Callback
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
from tensorflow.keras.models import Sequential
from tensorflow.keras.optimizers import Adam

# Constants
SOURCE_DATA_DIR = '/content/drive/Shareddrives/ML Dataset/'
ROOT_DIR = '/content/dataset'
TRAIN_DIR = os.path.join(ROOT_DIR, 'train')
VALID_DIR = os.path.join(ROOT_DIR, 'validation')
EVAL_DIR = os.path.join(ROOT_DIR, 'evaluation')
SPLIT_SIZE = 0.75
BATCH_SIZE = 32
TARGET_SIZE = (450, 450)
EPOCHS = 200
STEPS = 10
SAVED_MODEL_PATH = '/content/drive/Shareddrives/ML Dataset/transfer_learning/models/mobilenet_v3/1/'

class MyCallback(Callback):
    def __init__(self, thereshold):
        super().__init__()
        self.thereshold = thereshold

    def on_epoch_end(self, epoch, logs={}):
        if logs.get('accuracy') > self.thereshold:
            print("\nReached target metric. Stopping training!")
            self.model.stop_training = True

class PrintLrCallback(Callback):
  def on_epoch_end(self, epoch, logs=None):
    lr = float(K.get_value(self.model.optimizer.lr))
    print(f"Epoch {epoch+1}: Learning rate = {lr}")

def step_decay(epoch):
	initial_lr = 0.01
	drop = 0.5
	epochs_drop = 1
	lr = initial_lr * math.pow(drop, math.floor((1+epoch)/epochs_drop))
	return lr

# Function to create directories
def create_dirs(path, dir_names):
    for dir_name in dir_names:
        os.makedirs(os.path.join(path, dir_name), exist_ok=True)

# Function to split data
def split_data(source_dir, train_dir, valid_dir, split_size):
    files = os.listdir(source_dir)
    random.shuffle(files)
    split_point = int(len(files) * split_size)
    for i, file in enumerate(files):
        if os.path.getsize(os.path.join(source_dir, file)) > 0:
            if i < split_point:
                shutil.copy(os.path.join(source_dir, file), os.path.join(train_dir, file))
            else:
                shutil.copy(os.path.join(source_dir, file), os.path.join(valid_dir, file))

# Function to create data generators
def create_generators(train_dir, valid_dir, target_size, batch_size):
    train_datagen = ImageDataGenerator(rescale=1.0/255., shear_range=0.1, rotation_range=5, horizontal_flip=True)
    valid_datagen = ImageDataGenerator(rescale=1.0/255.)
    train_generator = train_datagen.flow_from_directory(directory=train_dir, batch_size=batch_size, class_mode='categorical', target_size=target_size)
    valid_generator = valid_datagen.flow_from_directory(directory=valid_dir, batch_size=batch_size, class_mode='categorical', target_size=target_size)
    return train_generator, valid_generator

# Function to create model
def create_model(input_shape):
    model = Sequential([
        Conv2D(16, (3, 3), activation='relu', input_shape=input_shape),
        MaxPooling2D(2, 2),
        Conv2D(32, (3, 3), activation='relu'),
        MaxPooling2D(2, 2),
        Conv2D(64, (3, 3), activation='relu'),
        MaxPooling2D(2, 2),
        Conv2D(128, (3, 3), activation='relu'),
        MaxPooling2D(2, 2),
        Flatten(),
        Dense(256, activation='relu'),
        Dropout(0.2),
        Dense(3, activation='softmax')
    ])
    model.compile(loss='categorical_crossentropy', optimizer=Adam(), metrics=['accuracy'])
    return model

# Function to create transfer learning model
def create_transfer_model(input_shape, SAVED_MODEL_PATH):
    model = Sequential([
        hub.KerasLayer(SAVED_MODEL_PATH, input_shape=input_shape, trainable=False),
        Dense(512, activation='relu'),
        Dropout(0.2),
        Dense(3, activation='softmax')
    ])
    model.compile(loss='categorical_crossentropy', optimizer=Adam(), metrics=['accuracy'])
    return model

# Function to plot metrics
def plot_metrics(history):
    metrics = ['accuracy', 'loss']
    for metric in metrics:
        plt.plot(history.history[metric], 'r', label='Training ' + metric)
        plt.plot(history.history['val_' + metric], 'b', label='Validation ' + metric)
        plt.title('Training and validation ' + metric)
        plt.legend()
        plt.show()

# Function to evaluate model
def evaluate_model(model, eval_dir, datasets, target_size, batch_size):
    eval_datagen = ImageDataGenerator(rescale=1.0/255.)
    for dataset in datasets:
        shutil.rmtree(eval_dir, ignore_errors=True)
        shutil.copytree(VALID_DIR, eval_dir)
        for directory in os.listdir(eval_dir):
            if directory != dataset:
                for file in os.listdir(os.path.join(eval_dir, directory))[1:]:
                    os.remove(os.path.join(eval_dir, directory, file))
        eval_generator = eval_datagen.flow_from_directory(directory=eval_dir, batch_size=batch_size, class_mode='categorical', target_size=target_size)
        print(f"Loss and Accuracy of {dataset}:")
        print(model.evaluate(eval_generator))
        print()

drive.mount('/content/drive')

# Create directories
if os.path.exists(ROOT_DIR):
  shutil.rmtree(ROOT_DIR)

dir_names = ['without_both', 'either_banner_or_logo', 'both_banner_and_logo']
create_dirs(TRAIN_DIR, dir_names)
create_dirs(VALID_DIR, dir_names)

# Split data
for dir_name in dir_names:
    split_data(os.path.join(SOURCE_DATA_DIR, dir_name), os.path.join(TRAIN_DIR, dir_name), os.path.join(VALID_DIR, dir_name), SPLIT_SIZE)

"""# Create and compile conv model"""

clear_session()

TARGET_SIZE = (450, 450)

train_generator, valid_generator = create_generators(TRAIN_DIR, VALID_DIR, TARGET_SIZE, BATCH_SIZE)

model = create_model(TARGET_SIZE + (3,))

model.summary()

# Train model
history = model.fit(train_generator,
                    epochs=EPOCHS,
                    validation_data=valid_generator,
                    steps_per_epoch=STEPS,
                    callbacks=[ModelCheckpoint(r'/content/drive/Shareddrives/ML Dataset/export/models/conv/2/', monitor='val_accuracy', save_best_only=True, verbose=1),
                               ReduceLROnPlateau(monitor='accuracy', factor=0.8, verbose=1, patience=5, min_lr=0.00005)])

# Plot metrics
plot_metrics(history)

# Evaluate model
evaluate_model(model, EVAL_DIR, dir_names, TARGET_SIZE, BATCH_SIZE)

"""# Create and train transfer learning model
## MobileNet

"""

TARGET_SIZE = (224, 224)
train_generator, valid_generator = create_generators(TRAIN_DIR, VALID_DIR, TARGET_SIZE, BATCH_SIZE)

SAVED_MODEL_PATH = '/content/drive/Shareddrives/ML Dataset/transfer_learning/models/mobilenet_v3/1/'

transfer_mobilenet = create_transfer_model(TARGET_SIZE + (3,), SAVED_MODEL_PATH)

transfer_mobilenet.summary()

# Train transfer learning model
transfer_mobilenet_history = transfer_mobilenet.fit(train_generator,
                                                    epochs=EPOCHS,
                                                    validation_data=valid_generator,
                                                    steps_per_epoch=STEPS,
                                                    callbacks=[ModelCheckpoint(r'/content/drive/Shareddrives/ML Dataset/export/models/mobilenet_tl/2/', monitor='val_accuracy', save_best_only=True, verbose=1),
                                                               ReduceLROnPlateau(monitor='accuracy', factor=0.8, verbose=1, patience=5, min_lr=0.00005)])


# Plot metrics
plot_metrics(transfer_mobilenet_history)

# Evaluate transfer learning model
evaluate_model(transfer_mobilenet, EVAL_DIR, dir_names, TARGET_SIZE, BATCH_SIZE)

"""# Inception"""

TARGET_SIZE = (299, 299)
train_generator, valid_generator = create_generators(TRAIN_DIR, VALID_DIR, TARGET_SIZE, BATCH_SIZE)

SAVED_MODEL_PATH = '/content/drive/Shareddrives/ML Dataset/transfer_learning/models/inception_v3/1/'

transfer_inception = create_transfer_model(TARGET_SIZE + (3,), SAVED_MODEL_PATH)

transfer_inception.summary()

# Train transfer learning model
transfer_inception_history = transfer_inception.fit(train_generator,
                                                    epochs=EPOCHS,
                                                    validation_data=valid_generator,
                                                    steps_per_epoch=STEPS,
                                                    callbacks=[ModelCheckpoint(r'/content/drive/Shareddrives/ML Dataset/export/models/inception_tl/2/', monitor='val_accuracy', save_best_only=True, verbose=1),
                                                    ReduceLROnPlateau(monitor='accuracy', factor=0.8, verbose=1, patience=5, min_lr=0.00005)])



# Plot metrics
plot_metrics(transfer_inception_history)

# Evaluate transfer learning model
evaluate_model(transfer_inception, EVAL_DIR, dir_names, TARGET_SIZE, BATCH_SIZE)

"""## ResNet

"""

TARGET_SIZE = (224, 224)
train_generator, valid_generator = create_generators(TRAIN_DIR, VALID_DIR, TARGET_SIZE, BATCH_SIZE)

SAVED_MODEL_PATH = '/content/drive/Shareddrives/ML Dataset/transfer_learning/models/resnet_v2/1/'

transfer_resnet = create_transfer_model(TARGET_SIZE + (3,), SAVED_MODEL_PATH)

transfer_resnet.summary()

# Train transfer learning model
transfer_resnet_history = transfer_resnet.fit(train_generator,
                                                    epochs=EPOCHS,
                                                    validation_data=valid_generator,
                                                    steps_per_epoch=STEPS,
                                                    callbacks=[ModelCheckpoint(r'/content/drive/Shareddrives/ML Dataset/export/models/resnet_tl/2/', monitor='val_accuracy', save_best_only=True, verbose=1),
                                                    ReduceLROnPlateau(monitor='accuracy', factor=0.8, verbose=1, patience=5, min_lr=0.000005)])



# Plot metrics
plot_metrics(transfer_resnet_history)

# Evaluate transfer learning model
evaluate_model(transfer_resnet, EVAL_DIR, dir_names, TARGET_SIZE, BATCH_SIZE)