test_final.py

import os
import pandas as pd
import numpy as np
import tensorflow as tf
import cv2
from sklearn.model_selection import train_test_split
from tensorflow.keras.preprocessing.image import ImageDataGenerator

def main():
    model = tf.keras.models.load_model("./model_mobilenet_25epochs.h5")
    
    test_folder = './test_mtcnn/'

    predictions = []
    files = os.listdir(test_folder)
    num_files = len(files)
    class_to_name = list(get_class_to_name().keys())
    
    preprocessed_images = []
    for filename in range(4977):
        base_filename = test_folder + str(filename) + ".jpg"
        #if os.path.exists(base_filename):
        img = cv2.imread(test_folder + str(filename) + ".jpg")
        img = cv2.resize(img, (224, 224))
        img = img.astype('float32') / 255.0
        preprocessed_images.append(img)

    print("Predicting....")
    predictions = model.predict(np.array(preprocessed_images))
    df_pred = []
    curr_idx = 0
    print("Sorting files...")
    for filename in range(4977):
        print(f"Current file... {filename}.jpg")
        base_filename = test_folder + str(filename) + ".jpg"
        # if os.path.exists(base_filename):
        prediction = predictions[curr_idx]
        predicted_class = np.argmax(prediction)
        df_pred.append((filename, class_to_name[predicted_class]))
        curr_idx += 1

    df = pd.DataFrame(df_pred, columns=['Id', 'Category'])

    # Save the DataFrame to a CSV file
    df.to_csv('predictions_epoch25.csv', index=False)

def get_class_to_name():
    with open("./new_train_mtcnn.csv", "r") as csv_file:
        df = pd.read_csv(csv_file, delimiter=',')
        
   # Split data into training and validation sets
    train_df, val_df = train_test_split(df, test_size=0.2, random_state=42)


    # Create training ImageDataGenerator, allows code to go through smaller chunks of training data 
    # at a time instead of loading all data at once
    train_datagen = ImageDataGenerator(
        rescale=1./255,
        rotation_range=20,
        width_shift_range=0.2,
        height_shift_range=0.2,
        shear_range=0.2,
        zoom_range=0.2,
        horizontal_flip=True,
        fill_mode='nearest'
    )

    # Load and augment training data
    train_generator = train_datagen.flow_from_dataframe(
        dataframe=train_df,
        directory='./train_mtcnn/',
        x_col='File Name',
        y_col='Category',
        target_size=(224, 224),
        batch_size=32,
        class_mode='categorical',
        shuffle=True,
        seed=42
    )

    class_to_name = train_generator.class_indices
    return class_to_name

main()