Update train_model.py

train_model.py

@@ -1,33 +1,59 @@
 # train_model.py
 import pandas as pd
+import numpy as np
 from sklearn.ensemble import RandomForestClassifier
 from sklearn.preprocessing import LabelEncoder
+from sklearn.model_selection import train_test_split
 import joblib
 import os
 
-# 1. Load/Sample Data
+# 1. Create Sample Dataset (REPLACE WITH YOUR ACTUAL DATA)
 data = {
-    'face_shape': ['Oval', 'Round', 'Square']*10,
-    'skin_tone': ['Fair', 'Medium', 'Dark']*10,
-    'face_size': ['Small', 'Medium', 'Large']*10,
-    'mask_style': ['StyleA', 'StyleB', 'StyleC']*10
+    'face_shape': ['Oval', 'Round', 'Square'] * 50,
+    'skin_tone': ['Fair', 'Medium', 'Dark'] * 50,
+    'face_size': ['Small', 'Medium', 'Large'] * 50,
+    'mask_style': ['Glitter', 'Animal', 'Floral'] * 50
 }
 df = pd.DataFrame(data)
 
-# 2. Create and Save Encoders
+# 2. Initialize Label Encoders
 encoders = {
-    col: LabelEncoder().fit(df[col].unique()) 
-    for col in ['face_shape', 'skin_tone', 'face_size', 'mask_style']
+    'face_shape': LabelEncoder().fit(df['face_shape'].unique()),
+    'skin_tone': LabelEncoder().fit(df['skin_tone'].unique()),
+    'face_size': LabelEncoder().fit(df['face_size'].unique()),
+    'mask_style': LabelEncoder().fit(df['mask_style'].unique())
 }
 
-# 3. Train Model
-X = pd.DataFrame({col: encoders[col].transform(df[col]) for col in encoders if col != 'mask_style'})
+# 3. Encode Features
+X = pd.DataFrame({
+    'face_shape': encoders['face_shape'].transform(df['face_shape']),
+    'skin_tone': encoders['skin_tone'].transform(df['skin_tone']),
+    'face_size': encoders['face_size'].transform(df['face_size'])
+})
 y = encoders['mask_style'].transform(df['mask_style'])
 
-model = RandomForestClassifier()
-model.fit(X, y)
+# 4. Train/Test Split
+X_train, X_test, y_train, y_test = train_test_split(
+    X, y, test_size=0.2, random_state=42
+)
 
-# 4. Save to model/ directory
+# 5. Train Model
+model = RandomForestClassifier(
+    n_estimators=100,
+    max_depth=5,
+    random_state=42
+)
+model.fit(X_train, y_train)
+
+# 6. Evaluate
+print(f"Training Accuracy: {model.score(X_train, y_train):.2f}")
+print(f"Test Accuracy: {model.score(X_test, y_test):.2f}")
+
+# 7. Save to model/ Directory
 os.makedirs('model', exist_ok=True)
 joblib.dump(model, 'model/random_forest.pkl')
-joblib.dump(encoders, 'model/label_encoders.pkl')
+joblib.dump(encoders, 'model/label_encoders.pkl')
+
+print("\nModel and encoders saved to model/ directory!")
+print("Face Shape Classes:", encoders['face_shape'].classes_)
+print("Mask Style Classes:", encoders['mask_style'].classes_)