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()