Update app.py

app.py

@@ -3,20 +3,38 @@ import gradio as gr
 from ultralytics import YOLO
 from PIL import Image
 
-# Load YOLOv8 model 
-model = YOLO('best.pt')
+# Load YOLOv8 model
+model = YOLO('/content/drive/MyDrive/yolov8_models/best.pt')
 
-# Folder with example test images 
-test_images_folder = 'test_images'
+# 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))
 
-# Predict using model
+# 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']
+
 def predict_image(image_path):
     results = model(image_path)
     img_array = results[0].plot(conf=False, labels=True, boxes=True)
-    return Image.fromarray(img_array)
+    
+    # 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
+    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
 
-# Run prediction logic
+# 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)
@@ -24,24 +42,25 @@ def run_prediction(uploaded_image, selected_image):
         image_path = os.path.join(test_images_folder, selected_image)
         return predict_image(image_path)
     else:
-        return None
+        return None, ""
 
-# Gradio Interface
+# Gradio interface
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("## 🦷 Dental Segmentation with YOLOv8")
-    gr.Markdown("Upload your own image or select a test image to run detection.")
+    gr.Markdown("Upload your own image or choose a test image from the list below.")
 
     with gr.Column():
         uploaded_image = gr.Image(label="Upload your image (optional)", type="filepath")
-        selected_image = gr.Dropdown(choices=test_images, label="Or choose a test image")
+        selected_image = gr.Dropdown(choices=test_images, label="...or select a test image")
 
-    gr.Markdown("###  Prediction Output")
-    output_image = gr.Image(label="Predicted Output")
+    gr.Markdown("### Prediction Result")
+    output_image = gr.Image(label="Predicted Image")
+    output_text = gr.Textbox(label="Detected & Missing Tooth Classes", lines=5)
 
     gr.Button("Run prediction").click(
         fn=run_prediction,
         inputs=[uploaded_image, selected_image],
-        outputs=output_image
+        outputs=[output_image, output_text]
     )
 
 demo.launch()