| import gradio as gr |
| import joblib |
| import os |
| from utils import extract_features |
|
|
| |
| if not os.path.exists('model/random_forest.pkl'): |
| print("Training model...") |
| os.system('python train_model.py') |
|
|
| |
| model = joblib.load('model/random_forest.pkl') |
| encoders = joblib.load('model/label_encoders.pkl') |
|
|
| def recommend_mask(image): |
| try: |
| face_shape, skin_tone, face_size = extract_features(image) |
| prediction = model.predict([[ |
| encoders['face_shape'].transform([face_shape])[0], |
| encoders['skin_tone'].transform([skin_tone])[0], |
| encoders['face_size'].transform([face_size])[0] |
| ]])[0] |
| return ( |
| encoders['mask_style'].classes_[prediction], |
| encoders['mask_images'][prediction] |
| ) |
| except Exception as e: |
| return "Error", "masks/default.png" |
|
|
| demo = gr.Interface( |
| fn=recommend_mask, |
| inputs=gr.Image(type="filepath"), |
| outputs=[ |
| gr.Textbox(label="Recommended Style"), |
| gr.Image(label="Mask Preview") |
| ], |
| title="🎭 Mask Recommender" |
| ) |
|
|
| if __name__ == "__main__": |
| demo.launch() |