Update Training/Code/train.py

Training/Code/train.py

@@ -1,11 +1,11 @@
 import os
 import numpy as np
-from keras.models import Model
-from keras.layers import Dense, Dropout, GlobalAveragePooling2D, Input
-from keras.optimizers import Adam
-from keras_preprocessing.image import ImageDataGenerator
-from keras.applications import MobileNetV2
-from keras.callbacks import EarlyStopping, ModelCheckpoint
+from tensorflow.keras.models import Model
+from tensorflow.keras.layers import Dense, Dropout, GlobalAveragePooling2D, Input
+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
 
 # Define paths
 base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '../'))
@@ -24,17 +24,20 @@ train_datagen = ImageDataGenerator(
 )
 val_datagen = ImageDataGenerator(rescale=1./255)
 
+# Use a larger image size for better accuracy
+img_size = 128
+
 train_generator = train_datagen.flow_from_directory(
-    train_dir, target_size=(96, 96), batch_size=32, color_mode='rgb', class_mode='categorical')
+    train_dir, target_size=(img_size, img_size), batch_size=32, color_mode='rgb', class_mode='categorical')
 
 validation_generator = val_datagen.flow_from_directory(
-    val_dir, target_size=(96, 96), batch_size=32, color_mode='rgb', class_mode='categorical')
+    val_dir, target_size=(img_size, img_size), batch_size=32, color_mode='rgb', class_mode='categorical')
 
 # Load base model
-base_model = MobileNetV2(include_top=False, input_shape=(96, 96, 3), weights='imagenet')
+base_model = MobileNetV2(include_top=False, input_shape=(img_size, img_size, 3), weights='imagenet')
 base_model.trainable = False  # Freeze base layers
 
-# Add custom layers
+# Add custom classification head
 x = base_model.output
 x = GlobalAveragePooling2D()(x)
 x = Dense(256, activation='relu')(x)
@@ -47,11 +50,11 @@ model.compile(optimizer=Adam(learning_rate=0.0001), loss='categorical_crossentro
 # Callbacks
 callbacks = [
     EarlyStopping(monitor='val_loss', patience=5, restore_best_weights=True),
-    ModelCheckpoint('best_model.h5', monitor='val_loss', save_best_only=True)
+    ModelCheckpoint('best_model.keras', monitor='val_loss', save_best_only=True)
 ]
 
 # Train the model
 model.fit(train_generator, validation_data=validation_generator, epochs=30, callbacks=callbacks)
 
-# Save model
+# Save the final model
 model.save("emotion_model.keras")