Create train_model.py

train_model.py

@@ -0,0 +1,44 @@
+import pandas as pd
+from sklearn.ensemble import RandomForestClassifier
+from sklearn.preprocessing import LabelEncoder
+import joblib
+import os
+
+# 1. Create dataset
+data = {
+    'face_shape': ['Oval', 'Round', 'Square'] * 100,
+    'skin_tone': ['Fair', 'Medium', 'Dark'] * 100,
+    'face_size': ['Small', 'Medium', 'Large'] * 100,
+    'mask_style': ['Glitter', 'Animal', 'Floral'] * 100
+}
+df = pd.DataFrame(data)
+
+# 2. Initialize encoders
+encoders = {
+    'face_shape': LabelEncoder().fit(df['face_shape']),
+    'skin_tone': LabelEncoder().fit(df['skin_tone']),
+    'face_size': LabelEncoder().fit(df['face_size']),
+    'mask_style': LabelEncoder().fit(df['mask_style']),
+    'mask_images': {
+        0: 'masks/glitter.png',
+        1: 'masks/animal.png',
+        2: 'masks/floral.png'
+    }
+}
+
+# 3. Train model
+model = RandomForestClassifier(n_estimators=150, random_state=42)
+model.fit(
+    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'])
+    }),
+    encoders['mask_style'].transform(df['mask_style'])
+)
+
+# 4. Save models
+os.makedirs('model', exist_ok=True)
+joblib.dump(model, 'model/random_forest.pkl')
+joblib.dump(encoders, 'model/label_encoders.pkl')
+print("Model trained and saved!")