Update app.py

app.py

@@ -26,6 +26,15 @@ glasses_images = {
     "Pilot (Aviator)": "glasses/aviator.jpg"
 }
 
+# Ensure the 'glasses' directory exists and contains the images
+if not os.path.exists("glasses"):
+    os.makedirs("glasses")
+    # Create dummy image files if they don't exist
+    for _, path in glasses_images.items():
+        if not os.path.exists(path):
+            dummy_image = Image.new('RGB', (200, 100), color='gray')
+            dummy_image.save(path)
+
 id2label = {0: 'Heart', 1: 'Oblong', 2: 'Oval', 3: 'Round', 4: 'Square'}
 label2id = {v: k for k, v in id2label.items()}
 
@@ -42,10 +51,14 @@ model = SwinForImageClassification.from_pretrained(
 )
 
 # Load your trained weights
-state_dict = torch.load('LR-0001-adamW-32-64swin.pth', map_location=device)
-model.load_state_dict(state_dict, strict=False)
-model.to(device)
-model.eval()
+# Ensure 'LR-0001-adamW-32-64swin.pth' is in the same directory or provide the correct path
+if os.path.exists('LR-0001-adamW-32-64swin.pth'):
+    state_dict = torch.load('LR-0001-adamW-32-64swin.pth', map_location=device)
+    model.load_state_dict(state_dict, strict=False)
+    model.to(device)
+    model.eval()
+else:
+    print("Warning: Trained weights file 'LR-0001-adamW-32-64swin.pth' not found. Using pre-trained weights only.")
 
 # Initialize Mediapipe
 mp_face_detection = mp.solutions.face_detection.FaceDetection(model_selection=1, min_detection_confidence=0.5)
@@ -105,8 +118,18 @@ def predict(image):
         # --- 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)
+        description = face_shape_descriptions.get(pred_label, "tidak dikenali")
+        frame_images = []
+
+        # Load images for all recommended frames
+        for frame in frame_recommendations:
+            frame_image_path = glasses_images.get(frame)
+            if frame_image_path and os.path.exists(frame_image_path):
+                try:
+                    frame_image = Image.open(frame_image_path)
+                    frame_images.append(frame_image)
+                except Exception as e:
+                    print(f"Error loading image for {frame}: {e}")
 
         # Build explanation text
         if frame_recommendations:
@@ -118,29 +141,32 @@ def predict(image):
             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 frame_image_path and os.path.exists(frame_image_path):
-            frame_image = Image.open(frame_image_path)
-        else:
-            frame_image = None
-
-        return pred_label, explanation, frame_image
+        return pred_label, explanation, frame_images
 
+    except ValueError as ve:
+        return "Error", str(ve), []
     except Exception as e:
-        return "Error", str(e), None
+        return "Error", f"Terjadi kesalahan: {str(e)}", []
 
 # Gradio Interface
-iface = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs=[
-        gr.Textbox(label="Bentuk Wajah Terdeteksi"),
-        gr.Textbox(label="Rekomendasi dan Penjelasan"),
-        gr.Image(label="Gambar Frame Rekomendasi")
-    ],
-    title="Rekomendasi Kacamata Berdasarkan Bentuk Wajah",
-    description="Upload foto wajahmu untuk mendapatkan rekomendasi bentuk kacamata yang sesuai!"
-)
+with gr.Blocks(theme=gr.themes.Soft()) as iface:
+    gr.Markdown("# Program Rekomendasi Kacamata Berdasarkan Bentuk Wajah")
+    gr.Markdown("Upload foto wajahmu untuk mendapatkan rekomendasi bentuk kacamata yang sesuai.")
+
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(type="pil")
+            with gr.Row():
+                upload_button = gr.UploadButton("Unggah Gambar", file_types=["image"])
+                clear_button = gr.Button("Ganti")
+        with gr.Column():
+            detected_shape = gr.Textbox(label="Bentuk Wajah Terdeteksi")
+            explanation_output = gr.Textbox(label="Penjelasan")
+            recommendation_gallery = gr.Gallery(label="Rekomendasi Kacamata", columns=3)
+
+    upload_button.upload(predict, inputs=upload_button, outputs=[detected_shape, explanation_output, recommendation_gallery])
+    clear_button.click(lambda: (None, "", []), inputs=None, outputs=[image_input, detected_shape, recommendation_gallery])
+    image_input.upload(predict, inputs=image_input, outputs=[detected_shape, explanation_output, recommendation_gallery])
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()