Update app.py

app.py

@@ -1,41 +1,35 @@
 import torch
+import ultralytics
 from ultralytics import YOLO
 from ultralytics.nn.tasks import DetectionModel
 from ultralytics.nn.modules.conv import Conv
 import torch.nn as nn
+import cv2
 import gradio as gr
 
-# ✅ Whitelist all necessary classes for PyTorch 2.6 safe loading
+# ---- FIX for PyTorch 2.6+ ----
 torch.serialization.add_safe_globals([DetectionModel, nn.Sequential, Conv])
 
-# Load YOLOv8 model
-model = YOLO("best.pt")  # Change to your model file path
+# ---- Load trained YOLO model ----
+model = YOLO("best.pt")  # Ensure your model file is in the same folder
 
-# Detection function
-def detect_fracture(image):
-    results = model.predict(source=image, save=False)
-    annotated_img = results[0].plot()
+# ---- Prediction function ----
+def predict(image):
+    # Run inference
+    results = model.predict(source=image, conf=0.25)
+    # Draw boxes on the image
+    result_image = results[0].plot()
+    # Convert BGR → RGB for Gradio
+    return cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB)
 
-    detections = []
-    for box in results[0].boxes:
-        cls_id = int(box.cls[0])
-        label = model.names[cls_id]
-        conf = float(box.conf[0])
-        detections.append(f"{label}: {conf:.2f}")
-
-    return annotated_img, "\n".join(detections) if detections else "No fracture detected"
-
-# Gradio app
-demo = gr.Interface(
-    fn=detect_fracture,
-    inputs=gr.Image(type="numpy", label="Upload Bone X-ray"),
-    outputs=[
-        gr.Image(type="numpy", label="Detection Result"),
-        gr.Textbox(label="Detected Fractures & Confidence Scores")
-    ],
+# ---- Gradio Interface ----
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="filepath", label="Upload Bone X-ray"),
+    outputs=gr.Image(type="numpy", label="Detection Result"),
     title="Human Bone Fracture Detection",
-    description="Upload an X-ray image to detect types of human bone fractures using a YOLOv8 model."
+    description="Upload an X-ray image to detect types of human bone fractures using YOLOv8."
 )
 
 if __name__ == "__main__":
-    demo.launch(debug=True)
+    iface.launch()