Update app.py

app.py

@@ -1,8 +1,11 @@
 import gradio as gr
 import torch
+import numpy as np
+from PIL import Image
+from roboflow import Roboflow
 from ultralytics import YOLO 
 
-file_path = 'best.pt'
+#file_path = 'best.pt'
 
 def load_model(file_path):
     # Load the Roboflow model
@@ -19,7 +22,8 @@ def load_model(file_path):
 
 def predict_fracture(image):
     # Preprocess the image for the Roboflow model
-    img_tensor = to_tensor(image).unsqueeze(0)  # Convert image to tensor and add batch dimension
+    img = Image.fromarray(image)
+    img_tensor = to_tensor(img).unsqueeze(0)  # Convert image to tensor and add batch dimension
 
     # Perform inference with the Roboflow model
     with torch.no_grad():
@@ -37,14 +41,19 @@ def predict_fracture(image):
             img_with_boxes.rectangle([xmin, ymin, xmax, ymax], outline=color, width=2)
             img_with_boxes.text((xmin, ymin), f"Fracture: {score:.2f}", font_size=12, color=color)
 
-    return img_with_boxes
-
+    return np.array(img_with_boxes)
 
+# Define the to_tensor function
+def to_tensor(image):
+    image = np.array(image) / 255.0
+    return torch.from_numpy(image.transpose((2, 0, 1))).float()
+    
 # Gradio Interface
 iface = gr.Interface(
     predict_fracture,
     inputs=gr.Image(),
     outputs=gr.Image(),
+    load_model=load_model,
     live=True,
     #capture_session=True,
     title="Bone Fracture Detection",