| import pandas as pd |
| from sklearn.ensemble import RandomForestClassifier |
| from sklearn.preprocessing import LabelEncoder |
| import joblib |
| import os |
|
|
| |
| 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) |
|
|
| |
| 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' |
| } |
| } |
|
|
| |
| 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']) |
| ) |
|
|
| |
| 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!") |
