RICHERGIRL commited on
Commit
ff5af75
·
verified ·
1 Parent(s): 642e4b0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +50 -17
app.py CHANGED
@@ -2,32 +2,65 @@ import gradio as gr
2
  import cv2
3
  import numpy as np
4
  import joblib
5
- from utils import extract_features # Your feature extraction logic
 
6
 
7
- # Load model and encoders
8
- model = joblib.load("model/random_forest.pkl")
9
- label_encoders = joblib.load("model/label_encoders.pkl")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
10
 
11
  def recommend_mask(image):
12
- # Extract face shape, skin tone, face size from image
13
- face_shape, skin_tone, face_size = extract_features(image)
14
-
15
- # Label encode features
16
- face_encoded = label_encoders["face_shape"].transform([face_shape])[0]
17
- skin_encoded = label_encoders["skin_tone"].transform([skin_tone])[0]
18
- size_encoded = label_encoders["face_size"].transform([face_size])[0]
19
-
20
- # Predict mask style
21
- prediction = model.predict([[face_encoded, skin_encoded, size_encoded]])[0]
22
- return prediction
 
 
 
 
 
 
23
 
24
- # Gradio Interface
 
 
 
25
  demo = gr.Interface(
26
  fn=recommend_mask,
27
  inputs=gr.Image(label="Upload Your Face", type="filepath"),
28
  outputs=gr.Textbox(label="Recommended Mask Style"),
29
  title="🎭 AI Party Mask Recommender",
30
  description="Upload a photo to get a personalized mask recommendation!",
 
31
  )
32
 
33
- demo.launch()
 
 
2
  import cv2
3
  import numpy as np
4
  import joblib
5
+ import os
6
+ from utils import extract_features # Make sure this exists in utils.py
7
 
8
+ def safe_load_model():
9
+ """Safely loads model files with error handling"""
10
+ try:
11
+ # Verify files exist
12
+ if not all(os.path.exists(f'model/{f}') for f in ['random_forest.pkl', 'label_encoders.pkl']):
13
+ raise FileNotFoundError("Model files missing")
14
+
15
+ # Load with mmap_mode for Hugging Face
16
+ model = joblib.load('model/random_forest.pkl', mmap_mode='r')
17
+ encoders = joblib.load('model/label_encoders.pkl', mmap_mode='r')
18
+ print("Model loaded successfully!")
19
+ return model, encoders
20
+
21
+ except Exception as e:
22
+ print(f"Model loading failed: {str(e)}")
23
+ # Fallback simple model
24
+ from sklearn.ensemble import RandomForestClassifier
25
+ from sklearn.preprocessing import LabelEncoder
26
+ print("Using fallback model")
27
+ return RandomForestClassifier(n_estimators=10), {
28
+ 'face_shape': LabelEncoder().fit(['Oval', 'Round', 'Square']),
29
+ 'skin_tone': LabelEncoder().fit(['Fair', 'Medium', 'Dark']),
30
+ 'face_size': LabelEncoder().fit(['Small', 'Medium', 'Large'])
31
+ }
32
 
33
  def recommend_mask(image):
34
+ """Process image and make prediction"""
35
+ try:
36
+ # Extract features
37
+ face_shape, skin_tone, face_size = extract_features(image)
38
+
39
+ # Encode features
40
+ face_encoded = encoders["face_shape"].transform([face_shape])[0]
41
+ skin_encoded = encoders["skin_tone"].transform([skin_tone])[0]
42
+ size_encoded = encoders["face_size"].transform([face_size])[0]
43
+
44
+ # Predict
45
+ prediction = model.predict([[face_encoded, skin_encoded, size_encoded]])[0]
46
+ return encoders["mask_style"].classes_[prediction]
47
+
48
+ except Exception as e:
49
+ print(f"Prediction error: {str(e)}")
50
+ return f"Error: Could not process image - {str(e)}"
51
 
52
+ # Initialize model and encoders
53
+ model, encoders = safe_load_model()
54
+
55
+ # Create Gradio interface
56
  demo = gr.Interface(
57
  fn=recommend_mask,
58
  inputs=gr.Image(label="Upload Your Face", type="filepath"),
59
  outputs=gr.Textbox(label="Recommended Mask Style"),
60
  title="🎭 AI Party Mask Recommender",
61
  description="Upload a photo to get a personalized mask recommendation!",
62
+ examples=[["example_face.jpg"]] if os.path.exists("example_face.jpg") else None
63
  )
64
 
65
+ if __name__ == "__main__":
66
+ demo.launch(server_name="0.0.0.0", server_port=7860)