sad-v2 / model.py
cheenchan's picture
--model
f977328 verified
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler
import tensorflow as tf
# Load the dataset
data = pd.read_csv('emo-final.csv')
# Separate features (X) and target labels (y)
X = data[['spO2', 'heart-rate', 'body-temperature']]
y = data[['anger', 'fear', 'sadness', 'disgust', 'surprise', 'anticipation', 'trust', 'joy']]
# Split the dataset into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# Normalize features using Min-Max scaling
scaler = MinMaxScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)
model = tf.keras.Sequential([
tf.keras.layers.Dense(64, activation='relu', input_shape=(3,)),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dense(8, activation='softmax') # Output layer with 8 units for 8 emotions
])
# Compile the model
model.compile(optimizer='adam', loss='mse')
# Train the model
model.fit(X_train_scaled, y_train, epochs=50, batch_size=32, validation_data=(X_test_scaled, y_test))
# Save the trained model
model.save('models/emotion_model.h5')