Spaces:
Sleeping
Sleeping
File size: 1,487 Bytes
e55378f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
import os
import numpy as np
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
from tensorflow.keras.optimizers import Adam
from utils.preprocess import load_images_from_folder # adjust import if needed
# Parameters
DATASET_PATH = './data/train'
IMG_SIZE = (48, 48) # Example size (width, height)
NUM_CLASSES = 7 # Number of emotion classes, adjust accordingly
EPOCHS = 30
BATCH_SIZE = 64
# Load data
x_data, y_data, emotions = load_images_from_folder(DATASET_PATH, IMG_SIZE)
# Preprocess labels
y_data_cat = to_categorical(y_data, num_classes=NUM_CLASSES)
# Normalize images
x_data = x_data.astype('float32') / 255.0
# Build a simple CNN model
model = Sequential([
Conv2D(32, (3,3), activation='relu', input_shape=(IMG_SIZE[0], IMG_SIZE[1], 1)),
MaxPooling2D(2,2),
Conv2D(64, (3,3), activation='relu'),
MaxPooling2D(2,2),
Conv2D(128, (3,3), activation='relu'),
MaxPooling2D(2,2),
Flatten(),
Dense(128, activation='relu'),
Dropout(0.5),
Dense(NUM_CLASSES, activation='softmax')
])
model.compile(optimizer=Adam(), loss='categorical_crossentropy', metrics=['accuracy'])
# Train
model.fit(x_data, y_data_cat, epochs=EPOCHS, batch_size=BATCH_SIZE, validation_split=0.2)
# Save model
os.makedirs('./models', exist_ok=True)
model.save('./models/emotion_model.h5')
print("Training complete and model saved.")
|