Update Training/Code/train.py

Training/Code/train.py

@@ -1,60 +1,99 @@
 import os
 import numpy as np
+import tensorflow as tf
+from tensorflow.keras.preprocessing.image import ImageDataGenerator
+from tensorflow.keras.applications import EfficientNetV2B1
 from tensorflow.keras.models import Model
-from tensorflow.keras.layers import Dense, Dropout, GlobalAveragePooling2D, Input
+from tensorflow.keras.layers import Dense, Dropout, GlobalAveragePooling2D
+from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, ReduceLROnPlateau
 from tensorflow.keras.optimizers import Adam
-from tensorflow.keras.preprocessing.image import ImageDataGenerator
-from tensorflow.keras.applications import MobileNetV2
-from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
+from sklearn.utils.class_weight import compute_class_weight
+
+# ==================== Paths ====================
+train_dir = "/content/combine_dataset/train"
+val_dir = "/content/combine_dataset/test"
 
-# Define paths
-base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '../'))
-train_dir = os.path.join(base_dir, 'Data/train')
-val_dir = os.path.join(base_dir, 'Data/test')
+# ==================== Parameters ====================
+img_size = (192, 192)  # Recommended for EfficientNetV2B1
+batch_size = 32
+epochs = 30
+num_classes = 7
 
-# Image generators with augmentation
+# ==================== Data Augmentation ====================
 train_datagen = ImageDataGenerator(
     rescale=1./255,
-    rotation_range=30,
-    zoom_range=0.2,
+    rotation_range=10,
+    zoom_range=0.1,
+    width_shift_range=0.05,
+    height_shift_range=0.05,
+    brightness_range=[0.9, 1.1],
     horizontal_flip=True,
-    shear_range=0.2,
-    width_shift_range=0.2,
-    height_shift_range=0.2
+    fill_mode='nearest'
 )
-val_datagen = ImageDataGenerator(rescale=1./255)
 
-# Use a larger image size for better accuracy
-img_size = 128
+val_datagen = ImageDataGenerator(rescale=1./255)
 
 train_generator = train_datagen.flow_from_directory(
-    train_dir, target_size=(img_size, img_size), batch_size=32, color_mode='rgb', class_mode='categorical')
+    train_dir,
+    target_size=img_size,
+    batch_size=batch_size,
+    class_mode='categorical',
+    shuffle=True
+)
 
-validation_generator = val_datagen.flow_from_directory(
-    val_dir, target_size=(img_size, img_size), batch_size=32, color_mode='rgb', class_mode='categorical')
+val_generator = val_datagen.flow_from_directory(
+    val_dir,
+    target_size=img_size,
+    batch_size=batch_size,
+    class_mode='categorical',
+    shuffle=False
+)
 
-# Load base model
-base_model = MobileNetV2(include_top=False, input_shape=(img_size, img_size, 3), weights='imagenet')
-base_model.trainable = False  # Freeze base layers
+# ==================== Compute Class Weights ====================
+labels = train_generator.classes
+class_weights = compute_class_weight(class_weight='balanced', classes=np.unique(labels), y=labels)
+class_weights = dict(enumerate(class_weights))
 
-# Add custom classification head
+# ==================== Build Model ====================
+base_model = EfficientNetV2B1(include_top=False, input_shape=(192, 192, 3), weights='imagenet')
 x = base_model.output
 x = GlobalAveragePooling2D()(x)
-x = Dense(256, activation='relu')(x)
-x = Dropout(0.5)(x)
-predictions = Dense(7, activation='softmax')(x)
+x = Dropout(0.4)(x)
+output = Dense(num_classes, activation='softmax')(x)
+model = Model(inputs=base_model.input, outputs=output)
+
+# ==================== Compile Model ====================
+optimizer = Adam(learning_rate=1e-5)
+model.compile(optimizer=optimizer, loss='categorical_crossentropy', metrics=['accuracy'])
 
-model = Model(inputs=base_model.input, outputs=predictions)
-model.compile(optimizer=Adam(learning_rate=0.0001), loss='categorical_crossentropy', metrics=['accuracy'])
+# ==================== Callbacks ====================
+checkpoint = ModelCheckpoint(
+    "/content/emotion_model.keras",
+    monitor='val_accuracy',
+    save_best_only=True,
+    verbose=1
+)
 
-# Callbacks
-callbacks = [
-    EarlyStopping(monitor='val_loss', patience=5, restore_best_weights=True),
-    ModelCheckpoint('best_model.keras', monitor='val_loss', save_best_only=True)
-]
+early_stop = EarlyStopping(
+    monitor='val_loss',
+    patience=7,
+    restore_best_weights=True,
+    verbose=1
+)
 
-# Train the model
-model.fit(train_generator, validation_data=validation_generator, epochs=30, callbacks=callbacks)
+lr_schedule = ReduceLROnPlateau(
+    monitor='val_loss',
+    factor=0.5,
+    patience=3,
+    verbose=1,
+    min_lr=1e-6
+)
 
-# Save the final model
-model.save("emotion_model.keras")
+# ==================== Train Model ====================
+model.fit(
+    train_generator,
+    validation_data=val_generator,
+    epochs=epochs,
+    callbacks=[checkpoint, early_stop, lr_schedule],
+    class_weight=class_weights
+)