Update app.py

app.py

@@ -3,20 +3,20 @@ import gradio as gr
 from ultralytics import YOLO
 from PIL import Image
 
-# Load YOLOv8 model
-model = YOLO('/content/drive/MyDrive/yolov8_models/best.pt')
+# Load YOLOv8 model (make sure best.pt is in the same folder as this script)
+model = YOLO('best.pt')
 
-# Folder with test images
+# Folder with example test images (must be inside your repo for HF Spaces)
 test_images_folder = 'test_images'
 test_images = sorted(os.listdir(test_images_folder))
 
-# Prediction function
+# Predict using model
 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)
 
-# Logic: use uploaded image if available, otherwise selected image
+# Run prediction logic
 def run_prediction(uploaded_image, selected_image):
     if uploaded_image is not None:
         return predict_image(uploaded_image)
@@ -26,17 +26,17 @@ def run_prediction(uploaded_image, selected_image):
     else:
         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 choose a test image from the list below.")
+    gr.Markdown("Upload your own image or select a test image to run detection.")
 
     with gr.Column():
         uploaded_image = gr.Image(label="Upload your image (optional)", type="filepath")
-        selected_image = gr.Dropdown(choices=test_images, label="...or select a test image")
+        selected_image = gr.Dropdown(choices=test_images, label="Or choose a test image")
 
-    gr.Markdown("### Prediction Result")
-    output_image = gr.Image(label="Predicted Image")
+    gr.Markdown("### 📸 Prediction Output")
+    output_image = gr.Image(label="Predicted Output")
 
     gr.Button("Run prediction").click(
         fn=run_prediction,