Update app.py

app.py

@@ -16,16 +16,7 @@ face_shape_descriptions = {
     "Square": "dengan rahang tegas dan dahi lebar."
 }
 
-# Frame recommendations
-glasses_recommendations = {
-    "Heart": "Frame Rimless",
-    "Oblong": "Frame Persegi Panjang",
-    "Oval": "Frame Bulat",
-    "Round": "Frame Kotak",
-    "Square": "Frame Oval"
-}
-
-# Frame images path (make sure these files are available!)
+# Frame images path
 glasses_images = {
     "Heart": "Glasses/RimlessFrame.jpg",
     "Oblong": "Glasses/RectangleFrame.jpg",
@@ -58,6 +49,22 @@ model.eval()
 # Initialize Mediapipe
 mp_face_detection = mp.solutions.face_detection.FaceDetection(model_selection=1, min_detection_confidence=0.5)
 
+# --- New: Decision tree function
+def recommend_glasses_tree(face_shape):
+    face_shape = face_shape.lower()
+    if face_shape == "square":
+        return ["Oval", "Round"]
+    elif face_shape == "round":
+        return ["Square", "Octagon", "Cat Eye"]
+    elif face_shape == "oval":
+        return ["Oval", "Pilot (Aviator)", "Cat Eye", "Round"]
+    elif face_shape == "heart":
+        return ["Pilot (Aviator)", "Cat Eye", "Round"]
+    elif face_shape == "oblong":
+        return ["Square", "Oval", "Pilot (Aviator)", "Cat Eye"]
+    else:
+        return []
+
 # Preprocess function
 def preprocess_image(image):
     img = np.array(image)
@@ -94,17 +101,24 @@ def predict(image):
             pred_label = id2label[pred_idx]
             pred_prob = probs[0][pred_idx].item() * 100
 
-        frame_recommendation = glasses_recommendations[pred_label]
+        # --- Use decision tree for recommendations
+        frame_recommendations = recommend_glasses_tree(pred_label)
+
         description = face_shape_descriptions[pred_label]
         frame_image_path = glasses_images.get(pred_label)
 
-        # Build the sentence
-        explanation = (f"Bentuk wajah kamu adalah {pred_label} ({pred_prob:.2f}%). "
-                       f"Kamu memiliki bentuk wajah {description} "
-                       f"Rekomendasi bentuk kacamata yang sesuai dengan wajah kamu adalah frame dengan bentuk {frame_recommendation}.")
+        # Build explanation text
+        if frame_recommendations:
+            recommended_frames = ', '.join(frame_recommendations)
+            explanation = (f"Bentuk wajah kamu adalah {pred_label} ({pred_prob:.2f}%). "
+                           f"Kamu memiliki bentuk wajah {description} "
+                           f"Rekomendasi bentuk kacamata yang sesuai dengan wajah kamu adalah: {recommended_frames}.")
+        else:
+            explanation = (f"Bentuk wajah kamu adalah {pred_label} ({pred_prob:.2f}%). "
+                           f"Tidak ada rekomendasi frame untuk bentuk wajah ini.")
 
         # Load frame image if available
-        if os.path.exists(frame_image_path):
+        if frame_image_path and os.path.exists(frame_image_path):
             frame_image = Image.open(frame_image_path)
         else:
             frame_image = None