Importing required libraries

import os import random import numpy as np import pandas as pd import tensorflow as tf import matplotlib.pyplot as plt from tensorflow.keras.models import Sequential, Model from tensorflow.keras.layers import Dense, Conv2D, Flatten, Dropout, MaxPooling2D, BatchNormalization, GlobalAveragePooling2D from tensorflow.keras.preprocessing.image import ImageDataGenerator from tensorflow.keras.preprocessing import image from tensorflow.keras.applications.vgg16 import VGG16

Setting up directories and parameters

train_dir = 'C:/Users/chaofan/Desktop/class/class/archive/seg_train/seg_train' # Training directory val_dir = 'C:/Users/chaofan/Desktop/class/class/archive/seg_test/seg_test' # Validation directory classes = os.listdir(train_dir) # List of classes batch_size = 64 # Batch size IMG_HEIGHT = 150 # Image height IMG_WIDTH = 150 # Image width epochs = 10 # Number of training epochs

Data generators for image preprocessing

train_image_generator = ImageDataGenerator( rescale=1./255, # Normalize pixel values horizontal_flip=True # Apply horizontal flip augmentation )

val_image_generator = ImageDataGenerator( rescale=1./255 # Normalize pixel values )

Generate training data

train_data_gen = train_image_generator.flow_from_directory( batch_size=batch_size, # Batch size directory=train_dir, # Training directory shuffle=True, # Shuffle data target_size=(IMG_HEIGHT, IMG_WIDTH), # Image size class_mode='categorical' # Classification mode )

Generate validation data

val_data_gen = val_image_generator.flow_from_directory( batch_size=batch_size, # Batch size directory=val_dir, # Validation directory shuffle=True, # Shuffle data target_size=(IMG_HEIGHT, IMG_WIDTH), # Image size class_mode='categorical' # Classification mode )

total_train = len(train_data_gen) # Total number of training samples total_val = len(val_data_gen) # Total number of validation samples

print('Total training batches: ', total_train) print('Total validation batches: ', total_val)

Load pre-trained VGG16 model

base_model = VGG16(weights='imagenet', include_top=False, input_shape=(IMG_HEIGHT, IMG_WIDTH, 3))

Freeze all layers in the base model

for layer in base_model.layers: layer.trainable = False

Build the classifier model on top of the base model

model = Sequential() model.add(base_model) model.add(Flatten()) model.add(Dense(1024, activation='relu')) model.add(Dense(512, activation='relu')) model.add(Dense(64, activation='relu')) model.add(Dense(len(classes), activation='softmax'))

Compile the model

model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

Print model summary

model.summary()

Train the model

history = model.fit( train_data_gen, # Training data generator steps_per_epoch=total_train, # Steps per epoch epochs=epochs, # Number of epochs validation_data=val_data_gen, # Validation data generator validation_steps=total_val # Validation steps per epoch )

Plot training and validation accuracy

acc = history.history['accuracy'] val_acc = history.history['val_accuracy'] loss = history.history['loss'] val_loss = history.history['val_loss']

plt.figure(figsize=(8, 8)) plt.subplot(2, 1, 1) plt.plot(acc, label='Training Accuracy') plt.plot(val_acc, label='Validation Accuracy') plt.legend(loc='lower right') plt.ylabel('Accuracy') plt.ylim([min(plt.ylim()), 1]) plt.title('Training and Validation Accuracy')

plt.subplot(2, 1, 2) plt.plot(loss, label='Training Loss') plt.plot(val_loss, label='Validation Loss') plt.legend(loc='upper right') plt.ylabel('Cross Entropy') plt.ylim([0, 1.0]) plt.title('Training and Validation Loss') plt.xlabel('Epoch')

plt.show()