File size: 1,370 Bytes
beb303e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 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!") |