train_CNN_1.py

"""
CNN Training Code
Version: 1
Dataset: all
accuracy: 0.2318 - loss: 2.8027 - val_accuracy: 0.2111 - val_loss: 2.9446
Epoch: 61/100
Batch-siez: 64
"""

import os
import pandas as pd
import cv2
import numpy as np
from sklearn.preprocessing import LabelEncoder
import tensorflow as tf
from tensorflow.keras import layers, models
from tensorflow.keras.utils import Sequence
from tensorflow.keras.callbacks import EarlyStopping
import pickle

# ======================= CONFIG =======================
TARGET_SIZE = (128, 128)  # Smaller input size for faster training
BATCH_SIZE = 64           # Larger batch size to leverage GPU parallelism

# ======================= GENERATOR =======================
class CardDatasetGenerator(tf.keras.utils.Sequence):
    def __init__(self, dataframe, label_encoder, batch_size=64, image_size=(128, 128), shuffle=True):
        self.df = dataframe.copy()
        self.label_encoder = label_encoder
        self.batch_size = batch_size
        self.image_size = image_size
        self.shuffle = shuffle
        self.indices = np.arange(len(self.df))
        self.on_epoch_end()

    def __len__(self):
        return int(np.ceil(len(self.df) / self.batch_size))

    def on_epoch_end(self):
        if self.shuffle:
            np.random.shuffle(self.indices)

    def __getitem__(self, index):
        batch_indices = self.indices[index * self.batch_size:(index + 1) * self.batch_size]
        batch_df = self.df.iloc[batch_indices]

        images, labels = [], []

        for _, row in batch_df.iterrows():
            filename, label, root = row.get("filename"), row.get("class"), row.get("root")
            if pd.isna(filename) or pd.isna(label) or pd.isna(root):
                continue

            img_path = os.path.join(root, filename)
            if not os.path.exists(img_path):
                continue

            img = cv2.imread(img_path)
            if img is None:
                continue

            if len(img.shape) == 2 or img.shape[2] == 1:
                img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
            else:
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

            img = cv2.resize(img, self.image_size)
            img = img.astype(np.float32) / 255.0

            images.append(img)
            labels.append(label)

        if not images:
            raise ValueError(f"No valid images found in batch {index}.")

        X = np.array(images)
        y = self.label_encoder.transform(labels)
        return X, y

# ======================= MAIN =======================
if __name__ == "__main__":
    df1 = pd.read_csv(r"C:\Users\zizhu\.ssh\ps_Project\.kaggle\archive\train_labels.csv")
    df2_train = pd.read_csv(r"C:\Users\zizhu\.ssh\ps_Project\stacked_cards_db\train_cards_label_translated.csv")
    df2_test = pd.read_csv(r"C:\Users\zizhu\.ssh\ps_Project\stacked_cards_db\test_cards_label_translated.csv")
    df3 = pd.read_csv(r"C:\Users\zizhu\.ssh\ps_Project\tree_cards_database\cards.csv")

    df3 = df3.rename(columns={"filename": "filename", "class": "class"})
    df3 = df3[["filename", "class"]]

    df1["root"] = df1["filename"].apply(
        lambda x: r"C:\Users\zizhu\.ssh\ps_Project\.kaggle\archive\train\train" if os.path.exists(
            os.path.join(r"C:\Users\zizhu\.ssh\ps_Project\.kaggle\archive\train\train", x)) else
        r"C:\Users\zizhu\.ssh\ps_Project\.kaggle\archive\test\test")

    df2_train["root"] = r"C:\Users\zizhu\.ssh\ps_Project\stacked_cards_db\train_zipped"
    df2_test["root"] = r"C:\Users\zizhu\.ssh\ps_Project\stacked_cards_db\test_zipped"
    df3["root"] = r"C:\Users\zizhu\.ssh\ps_Project\tree_cards_database"

    df = pd.concat([df1, df2_train, df2_test, df3], ignore_index=True)

    label_encoder = LabelEncoder()
    label_encoder.fit(df["class"])

    with open("label_encoder.pkl", "wb") as f:
        pickle.dump(label_encoder, f)

    df = df.sample(frac=1, random_state=42).reset_index(drop=True)
    val_split = int(0.8 * len(df))
    train_df, val_df = df.iloc[:val_split], df.iloc[val_split:]

    train_gen = CardDatasetGenerator(train_df, label_encoder=label_encoder, batch_size=BATCH_SIZE, image_size=TARGET_SIZE)
    val_gen = CardDatasetGenerator(val_df, label_encoder=label_encoder, batch_size=BATCH_SIZE, image_size=TARGET_SIZE, shuffle=False)

    model = models.Sequential([
        layers.Input(shape=(TARGET_SIZE[0], TARGET_SIZE[1], 3)),
        layers.Conv2D(32, (3, 3), activation='relu'),
        layers.MaxPooling2D((2, 2)),
        layers.Conv2D(64, (3, 3), activation='relu'),
        layers.MaxPooling2D((2, 2)),
        layers.Conv2D(128, (3, 3), activation='relu'),
        layers.GlobalAveragePooling2D(),
        layers.Dense(256, activation='relu'),
        layers.Dropout(0.4),
        layers.Dense(len(label_encoder.classes_), activation='softmax')
    ])

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

    early_stop = EarlyStopping(monitor='val_accuracy', patience=5, restore_best_weights=True)
    model.fit(
    train_gen,
    validation_data=val_gen,
    epochs=100,
    callbacks=[early_stop]
    )   

    model.save("model_CNN_1.keras")