Eco-Try-webapp / train_model.py
EcoTry's picture
Upload train_model.py
77ee48b verified
Raw
History Blame Contribute Delete
1.21 kB
# pages/train_model.py
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.neural_network import MLPClassifier
from sklearn.pipeline import Pipeline
import joblib
# Load dataset
data = pd.read_csv("size_dataset.csv")
# Features
X = data[[
"gender",
"shoulder",
"chest",
"waist",
"hip",
"chest_depth",
"hip_depth",
"height",
"weight"
]]
# Target
y = data["size"]
# Encode labels
le = LabelEncoder()
y_encoded = le.fit_transform(y)
# Train-test split
X_train, X_test, y_train, y_test = train_test_split(
X, y_encoded, test_size=0.2, random_state=42
)
# Scaling + MLP in Pipeline
model = Pipeline([
("scaler", StandardScaler()),
("mlp", MLPClassifier(
hidden_layer_sizes=(128, 64, 32),
max_iter=2000,
random_state=42
))
])
# Train
model.fit(X_train, y_train)
# Accuracy
accuracy = model.score(X_test, y_test)
print("Model Accuracy:", accuracy)
# Save model + encoder
joblib.dump(model, "size_model.pkl")
joblib.dump(le, "label_encoder.pkl")
print("Model saved successfully!")