Upload train_model.py

train_model.py

@@ -0,0 +1,61 @@
+# -*- coding: utf-8 -*-
+"""train_model.py
+
+Automatically generated by Colab.
+
+Original file is located at
+    https://colab.research.google.com/drive/1Oq8iUsTw8pCPeB1qLiaYLCI6QRVjO05g
+"""
+
+pip install joblib
+
+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
+
+# Sample dataset (replace with your actual data)
+data = {
+    'face_shape': ['Oval', 'Round', 'Square', 'Oval', 'Round', 'Square'] * 10,
+    'skin_tone': ['Fair', 'Medium', 'Dark', 'Medium', 'Dark', 'Fair'] * 10,
+    'face_size': ['Small', 'Medium', 'Large', 'Medium', 'Large', 'Small'] * 10,
+    'mask_style': ['Glitter', 'Animal', 'Floral', 'Animal', 'Glitter', 'Floral'] * 10
+}
+df = pd.DataFrame(data)
+
+# Initialize LabelEncoders
+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'])
+}
+
+# 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'])
+
+# Train-test split
+X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
+
+# Train Random Forest
+model = RandomForestClassifier(n_estimators=100, random_state=42)
+model.fit(X_train, y_train)
+
+# Create model directory
+os.makedirs('model', exist_ok=True)
+
+# Save model and encoders
+joblib.dump(model, 'model/random_forest.pkl')
+joblib.dump(encoders, 'model/label_encoders.pkl')
+
+print("Model trained and saved successfully!")
+print(f"Test Accuracy: {model.score(X_test, y_test):.2f}")
+