Update Training/Code/train.py

Training/Code/train.py

@@ -1,72 +1,57 @@
-# Importing Libraries
-import os
-import numpy as np
-import cv2
-from keras.models import Sequential
-from keras.layers import Dense, Dropout, Flatten, Conv2D, MaxPooling2D, Input
-from keras.optimizers import Adam
-from keras_preprocessing.image import ImageDataGenerator
-from keras.callbacks import EarlyStopping, ModelCheckpoint
-
-# Define paths using os.path for portability
-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')
-
-# Data augmentation and rescaling
-train_datagen = ImageDataGenerator(
-    rescale=1./255,
-    rotation_range=30,
-    zoom_range=0.2,
-    horizontal_flip=True,
-    shear_range=0.2,
-    width_shift_range=0.2,
-    height_shift_range=0.2
-)
-val_datagen = ImageDataGenerator(rescale=1./255)
-
-train_generator = train_datagen.flow_from_directory(
-    train_dir, target_size=(48, 48), batch_size=64, color_mode='grayscale', class_mode='categorical')
-
-validation_generator = val_datagen.flow_from_directory(
-    val_dir, target_size=(48, 48), batch_size=64, color_mode='grayscale', class_mode='categorical')
-
-# Building the Convolutional Network Architecture
-emotion_model = Sequential()
-emotion_model.add(Input(shape=(48, 48, 1)))
-emotion_model.add(Conv2D(32, kernel_size=(3, 3), activation='relu'))
-emotion_model.add(Conv2D(64, kernel_size=(3, 3), activation='relu'))
-emotion_model.add(MaxPooling2D(pool_size=(2, 2)))
-emotion_model.add(Dropout(0.25))
-
-emotion_model.add(Conv2D(128, kernel_size=(3, 3), activation='relu'))
-emotion_model.add(MaxPooling2D(pool_size=(2, 2)))
-emotion_model.add(Conv2D(128, kernel_size=(3, 3), activation='relu'))
-emotion_model.add(MaxPooling2D(pool_size=(2, 2)))
-emotion_model.add(Dropout(0.25))
-
-emotion_model.add(Flatten())
-emotion_model.add(Dense(1024, activation='relu'))
-emotion_model.add(Dropout(0.5))
-emotion_model.add(Dense(7, activation='softmax'))
-
-# Compile the model
-emotion_model.compile(loss='categorical_crossentropy', optimizer=Adam(learning_rate=0.0001), metrics=['accuracy'])
-
-# Define callbacks
-callbacks = [
-    EarlyStopping(monitor='val_loss', patience=5, restore_best_weights=True),
-    ModelCheckpoint('best_model.h5', monitor='val_loss', save_best_only=True)
-]
-
-# Train the model
-emotion_model_info = emotion_model.fit(
-    train_generator,
-    epochs=50,
-    validation_data=validation_generator,
-    callbacks=callbacks
-)
-
-# Save the full model
-emotion_model.save("emotion_model.keras")
-
+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
+
+# 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')
+
+# Image generators with augmentation
+train_datagen = ImageDataGenerator(
+    rescale=1./255,
+    rotation_range=30,
+    zoom_range=0.2,
+    horizontal_flip=True,
+    shear_range=0.2,
+    width_shift_range=0.2,
+    height_shift_range=0.2
+)
+val_datagen = ImageDataGenerator(rescale=1./255)
+
+train_generator = train_datagen.flow_from_directory(
+    train_dir, target_size=(96, 96), 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')
+
+# Load base model
+base_model = MobileNetV2(include_top=False, input_shape=(96, 96, 3), weights='imagenet')
+base_model.trainable = False  # Freeze base layers
+
+# Add custom layers
+x = base_model.output
+x = GlobalAveragePooling2D()(x)
+x = Dense(256, activation='relu')(x)
+x = Dropout(0.5)(x)
+predictions = Dense(7, activation='softmax')(x)
+
+model = Model(inputs=base_model.input, outputs=predictions)
+model.compile(optimizer=Adam(learning_rate=0.0001), loss='categorical_crossentropy', metrics=['accuracy'])
+
+# Callbacks
+callbacks = [
+    EarlyStopping(monitor='val_loss', patience=5, restore_best_weights=True),
+    ModelCheckpoint('best_model.h5', monitor='val_loss', save_best_only=True)
+]
+
+# Train the model
+model.fit(train_generator, validation_data=validation_generator, epochs=30, callbacks=callbacks)
+
+# Save model
+model.save("emotion_model.keras")