Update app.py

app.py

@@ -1,74 +1,33 @@
-import pandas as pd
-from sklearn.preprocessing import LabelEncoder
-from sklearn.ensemble import RandomForestClassifier
-from sklearn.model_selection import train_test_split
 import gradio as gr
-
-# Function to load/generate data
-def load_data():
-    data = {
-        'face_shape': ['round', 'oval', 'square', 'heart', 'oval', 'round', 'square', 'heart', 'oval', 'square'],
-        'skin_tone': ['fair', 'medium', 'deep', 'cool', 'warm', 'fair', 'medium', 'deep', 'cool', 'warm'],
-        'face_size': ['small', 'medium', 'large', 'medium', 'small', 'large', 'small', 'large', 'medium', 'small'],
-        'mask_style': ['glitter_cat', 'gold_venetian', 'black_minimal', 'floral_masquerade', 'gold_venetian',
-                       'glitter_cat', 'black_minimal', 'floral_masquerade', 'gold_venetian', 'black_minimal']
-    }
-    return pd.DataFrame(data)
-
-# Load data
-df = load_data()
-
-# Label encode categorical features
-label_encoders = {}
-for column in ['face_shape', 'skin_tone', 'face_size', 'mask_style']:
-    le = LabelEncoder()
-    df[column] = le.fit_transform(df[column])
-    label_encoders[column] = le
-
-# Split features and target
-X = df[['face_shape', 'skin_tone', 'face_size']]
-y = df['mask_style']
-
-# Split into train and test sets
-X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.5, random_state=42, stratify=y)
-
-# Train Random Forest
-rf_model = RandomForestClassifier(n_estimators=100, random_state=42)
-rf_model.fit(X_train, y_train)
-
-# Define the recommendation function
-def recommend_mask(face_shape, skin_tone, face_size):
-    input_data = pd.DataFrame({
-        'face_shape': [face_shape],
-        'skin_tone': [skin_tone],
-        'face_size': [face_size]
-    })
-    for col in input_data.columns:
-        input_data[col] = label_encoders[col].transform(input_data[col])
-    prediction_encoded = rf_model.predict(input_data)
-    predicted_label = label_encoders['mask_style'].inverse_transform(prediction_encoded)
-    return predicted_label[0]
-
-# Get unique values for dropdown choices
-face_shapes_labels = df['face_shape'].unique().tolist()
-face_shapes_labels = label_encoders['face_shape'].inverse_transform(face_shapes_labels).tolist()
-skin_tones_labels = df['skin_tone'].unique().tolist()
-skin_tones_labels = label_encoders['skin_tone'].inverse_transform(skin_tones_labels).tolist()
-face_sizes_labels = df['face_size'].unique().tolist()
-face_sizes_labels = label_encoders['face_size'].inverse_transform(face_sizes_labels).tolist()
-
-
-# Define Gradio interface
-iface = gr.Interface(
+import cv2
+import numpy as np
+import joblib
+from utils import extract_features  # Your feature extraction logic
+
+# Load model and encoders
+model = joblib.load("model/random_forest.pkl")
+label_encoders = joblib.load("model/label_encoders.pkl")
+
+def recommend_mask(image):
+    # Extract face shape, skin tone, face size from image
+    face_shape, skin_tone, face_size = extract_features(image)
+    
+    # Label encode features
+    face_encoded = label_encoders["face_shape"].transform([face_shape])[0]
+    skin_encoded = label_encoders["skin_tone"].transform([skin_tone])[0]
+    size_encoded = label_encoders["face_size"].transform([face_size])[0]
+    
+    # Predict mask style
+    prediction = model.predict([[face_encoded, skin_encoded, size_encoded]])[0]
+    return prediction
+
+# Gradio Interface
+demo = gr.Interface(
     fn=recommend_mask,
-    inputs=[
-        gr.Dropdown(choices=face_shapes_labels, label="Face Shape"),
-        gr.Dropdown(choices=skin_tones_labels, label="Skin Tone"),
-        gr.Dropdown(choices=face_sizes_labels, label="Face Size"),
-    ],
-    outputs="text",
-    title="🎭 Party Face Mask Recommender",
-    description="Get personalized party face mask recommendations based on your facial features."
+    inputs=gr.Image(label="Upload Your Face", type="filepath"),
+    outputs=gr.Textbox(label="Recommended Mask Style"),
+    title="🎭 AI Party Mask Recommender",
+    description="Upload a photo to get a personalized mask recommendation!",
 )
 
-iface.launch(share=True, server_name="0.0.0.0", server_port=7860)
+demo.launch()