Update app.py

app.py

@@ -3,48 +3,48 @@ import gradio as gr
 from ultralytics import YOLO
 from PIL import Image
 
+# Adjust these paths based on where you uploaded the files in your HF repo
+MODEL_PATH = 'best.pt'  # Put your best.pt at repo root or adjust if in a folder
+TEST_IMAGES_FOLDER = 'test_images'  # Folder inside your repo with test images
+
 # Load YOLOv8 model
-model = YOLO('/content/drive/MyDrive/yolov8_models/best.pt')
+model = YOLO(MODEL_PATH)
 
-# Folder with test images
-test_images_folder = '/content/Instance_seg_teeth/Dataset/yolo_test_dataset/test/images'
-test_images = sorted(os.listdir(test_images_folder))
+# List all test images available (make sure folder exists and has images)
+test_images = sorted(os.listdir(TEST_IMAGES_FOLDER)) if os.path.exists(TEST_IMAGES_FOLDER) else []
 
-# List of all tooth classes (from your dataset)
-tooth_classes = ['11', '12', '13', '14', '15', '16', '17', '18', '21', '22', '23', '24', '25', '26', '27', '28', 
-                 '31', '32', '33', '34', '35', '36', '37', '38', '41', '42', '43', '44', '45', '46', '47', '48']
+# List of all tooth classes (your dataset classes)
+tooth_classes = ['11', '12', '13', '14', '15', '16', '17', '18', 
+                 '21', '22', '23', '24', '25', '26', '27', '28', 
+                 '31', '32', '33', '34', '35', '36', '37', '38', 
+                 '41', '42', '43', '44', '45', '46', '47', '48']
 
 def predict_image(image_path):
     results = model(image_path)
     img_array = results[0].plot(conf=False, labels=True, boxes=True)
-    
-    # Extract predicted classes (as numbers)
-    pred_classes = results[0].boxes.cls.cpu().numpy().astype(int)  # class indices
-    
-    # Map indices to class names
-    detected_classes = sorted(set([tooth_classes[i] for i in pred_classes]))
-    
-    # Classes not detected
+
+    # Extract predicted classes (indices)
+    pred_classes = results[0].boxes.cls.cpu().numpy().astype(int) if len(results[0].boxes) > 0 else []
+
+    detected_classes = sorted(set([tooth_classes[i] for i in pred_classes])) if len(pred_classes) > 0 else []
     missing_classes = sorted(set(tooth_classes) - set(detected_classes))
-    
+
     detected_str = ", ".join(detected_classes) if detected_classes else "None"
     missing_str = ", ".join(missing_classes) if missing_classes else "None"
-    
+
     info_text = f"Detected tooth classes:\n{detected_str}\n\nMissing tooth classes:\n{missing_str}"
-    
+
     return Image.fromarray(img_array), info_text
 
-# Logic: use uploaded image if available, otherwise selected image
 def run_prediction(uploaded_image, selected_image):
     if uploaded_image is not None:
         return predict_image(uploaded_image)
     elif selected_image is not None:
-        image_path = os.path.join(test_images_folder, selected_image)
+        image_path = os.path.join(TEST_IMAGES_FOLDER, selected_image)
         return predict_image(image_path)
     else:
         return None, ""
 
-# Gradio interface
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("## 🦷 Dental Segmentation with YOLOv8")
     gr.Markdown("Upload your own image or choose a test image from the list below.")