Spaces:
Configuration error
Configuration error
| import gradio as gr | |
| import joblib | |
| import os | |
| import numpy as np | |
| from sklearn.ensemble import RandomForestClassifier | |
| from sklearn.preprocessing import LabelEncoder | |
| from utils import extract_features | |
| def initialize_fallback_model(): | |
| """Creates and trains a simple fallback model with image support""" | |
| print("Initializing fallback model...") | |
| # Simple training data | |
| X = np.array([[0,0,0], [1,1,1], [2,2,2]]) # Dummy encoded features | |
| y = np.array([0, 1, 0]) # Dummy target | |
| model = RandomForestClassifier(n_estimators=10) | |
| model.fit(X, y) | |
| encoders = { | |
| 'face_shape': LabelEncoder().fit(['Oval', 'Round', 'Square']), | |
| 'skin_tone': LabelEncoder().fit(['Fair', 'Medium', 'Dark']), | |
| 'face_size': LabelEncoder().fit(['Small', 'Medium', 'Large']), | |
| 'mask_style': LabelEncoder().fit(['Glitter', 'Animal', 'Floral']), | |
| 'mask_images': { | |
| 0: 'masks/glitter.png', | |
| 1: 'masks/animal.png', | |
| 2: 'masks/floral.png' | |
| } | |
| } | |
| return model, encoders | |
| def safe_load_model(): | |
| """Safely loads model files with comprehensive fallback""" | |
| try: | |
| if not all(os.path.exists(f'model/{f}') for f in ['random_forest.pkl', 'label_encoders.pkl']): | |
| raise FileNotFoundError("Model files missing") | |
| model = joblib.load('model/random_forest.pkl', mmap_mode='r') | |
| encoders = joblib.load('model/label_encoders.pkl', mmap_mode='r') | |
| # Verify model is fitted and has image mappings | |
| if not hasattr(model, 'classes_') or 'mask_images' not in encoders: | |
| raise ValueError("Model or encoders incomplete") | |
| print("Main model loaded successfully!") | |
| return model, encoders | |
| except Exception as e: | |
| print(f"Loading failed: {str(e)}") | |
| return initialize_fallback_model() | |
| def recommend_mask(image): | |
| """Process image and return both text and image recommendations""" | |
| try: | |
| # Extract features | |
| face_shape, skin_tone, face_size = extract_features(image) | |
| # Encode features | |
| face_encoded = encoders["face_shape"].transform([face_shape])[0] | |
| skin_encoded = encoders["skin_tone"].transform([skin_tone])[0] | |
| size_encoded = encoders["face_size"].transform([face_size])[0] | |
| # Predict | |
| prediction = model.predict([[face_encoded, skin_encoded, size_encoded]])[0] | |
| # Get both text and image recommendations | |
| mask_style = encoders["mask_style"].classes_[prediction] | |
| mask_image = encoders['mask_images'].get(prediction, 'masks/default.png') | |
| # Verify image exists | |
| if not os.path.exists(mask_image): | |
| print(f"Warning: Mask image not found at {mask_image}") | |
| mask_image = 'masks/default.png' | |
| return mask_style, mask_image | |
| except Exception as e: | |
| print(f"Prediction error: {str(e)}") | |
| return "Basic Mask (Fallback)", "masks/default.png" | |
| # Initialize model and encoders | |
| model, encoders = safe_load_model() | |
| # Verify masks directory exists | |
| os.makedirs('masks', exist_ok=True) | |
| # Create Gradio interface | |
| demo = gr.Interface( | |
| fn=recommend_mask, | |
| inputs=gr.Image(label="Upload Your Face", type="filepath"), | |
| outputs=[ | |
| gr.Textbox(label="Recommended Style"), | |
| gr.Image(label="Mask Preview", type="filepath") | |
| ], | |
| title="🎭 AI Party Mask Recommender", | |
| description="Upload a photo to get a personalized mask recommendation!", | |
| examples=[["example_face.jpg"]] if os.path.exists("example_face.jpg") else None | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |