RICHERGIRL commited on
Commit
3ee9576
·
verified ·
1 Parent(s): 7fed2a1

Create model/random_forest.pkl

Browse files
Files changed (1) hide show
  1. model/random_forest.pkl +49 -0
model/random_forest.pkl ADDED
@@ -0,0 +1,49 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import pandas as pd
2
+ import numpy as np
3
+ from sklearn.ensemble import RandomForestClassifier
4
+ from sklearn.preprocessing import LabelEncoder
5
+ from sklearn.model_selection import train_test_split
6
+ import joblib
7
+ import os
8
+
9
+ # Sample dataset (replace with your actual data)
10
+ data = {
11
+ 'face_shape': ['Oval', 'Round', 'Square', 'Oval', 'Round', 'Square'] * 10,
12
+ 'skin_tone': ['Fair', 'Medium', 'Dark', 'Medium', 'Dark', 'Fair'] * 10,
13
+ 'face_size': ['Small', 'Medium', 'Large', 'Medium', 'Large', 'Small'] * 10,
14
+ 'mask_style': ['Glitter', 'Animal', 'Floral', 'Animal', 'Glitter', 'Floral'] * 10
15
+ }
16
+ df = pd.DataFrame(data)
17
+
18
+ # Initialize LabelEncoders
19
+ encoders = {
20
+ 'face_shape': LabelEncoder().fit(df['face_shape']),
21
+ 'skin_tone': LabelEncoder().fit(df['skin_tone']),
22
+ 'face_size': LabelEncoder().fit(df['face_size']),
23
+ 'mask_style': LabelEncoder().fit(df['mask_style'])
24
+ }
25
+
26
+ # Encode features
27
+ X = pd.DataFrame({
28
+ 'face_shape': encoders['face_shape'].transform(df['face_shape']),
29
+ 'skin_tone': encoders['skin_tone'].transform(df['skin_tone']),
30
+ 'face_size': encoders['face_size'].transform(df['face_size'])
31
+ })
32
+ y = encoders['mask_style'].transform(df['mask_style'])
33
+
34
+ # Train-test split
35
+ X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
36
+
37
+ # Train Random Forest
38
+ model = RandomForestClassifier(n_estimators=100, random_state=42)
39
+ model.fit(X_train, y_train)
40
+
41
+ # Create model directory
42
+ os.makedirs('model', exist_ok=True)
43
+
44
+ # Save model and encoders
45
+ joblib.dump(model, 'model/random_forest.pkl')
46
+ joblib.dump(encoders, 'model/label_encoders.pkl')
47
+
48
+ print("Model trained and saved successfully!")
49
+ print(f"Test Accuracy: {model.score(X_test, y_test):.2f}")