Update app.py

app.py

@@ -7,7 +7,7 @@ from sklearn.preprocessing import LabelEncoder
 from utils import extract_features
 
 def initialize_fallback_model():
-    """Creates and trains a simple fallback model"""
+    """Creates and trains a simple fallback model with image support"""
     print("Initializing fallback model...")
     
     # Simple training data
@@ -21,7 +21,12 @@ def initialize_fallback_model():
         '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(['StyleA', 'StyleB', 'StyleC'])  # Added mask_style
+        'mask_style': LabelEncoder().fit(['Glitter', 'Animal', 'Floral']),
+        'mask_images': {
+            0: 'masks/glitter.png',
+            1: 'masks/animal.png',
+            2: 'masks/floral.png'
+        }
     }
     
     return model, encoders
@@ -35,9 +40,9 @@ def safe_load_model():
         model = joblib.load('model/random_forest.pkl', mmap_mode='r')
         encoders = joblib.load('model/label_encoders.pkl', mmap_mode='r')
         
-        # Verify model is fitted
-        if not hasattr(model, 'classes_'):
-            raise ValueError("Model not properly trained")
+        # 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
@@ -47,7 +52,7 @@ def safe_load_model():
         return initialize_fallback_model()
 
 def recommend_mask(image):
-    """Process image and make prediction with error handling"""
+    """Process image and return both text and image recommendations"""
     try:
         # Extract features
         face_shape, skin_tone, face_size = extract_features(image)
@@ -59,33 +64,40 @@ def recommend_mask(image):
         
         # Predict
         prediction = model.predict([[face_encoded, skin_encoded, size_encoded]])[0]
-        return encoders["mask_style"].classes_[prediction]
-
-        # Get recommended mask image path
-        mask_image_path = encoders['mask_images'][prediction]
-        return (
-            encoders["mask_style"].classes_[prediction],  # Text
-            mask_image_path                               # Image
-        )
+        
+        # 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 "Error", "default_mask.png"  # Fallback
-
+        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(type="filepath"),
+    inputs=gr.Image(label="Upload Your Face", type="filepath"),
     outputs=[
         gr.Textbox(label="Recommended Style"),
-        gr.Image(label="Mask Preview")  # Add image output
+        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()