File size: 1,163 Bytes
691ea33 | 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 | import gradio as gr
import joblib
import os
from utils import extract_features
# Auto-train if model doesn't exist
if not os.path.exists('model/random_forest.pkl'):
print("Training model...")
os.system('python train_model.py')
# Load model
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() |