Update app.py

app.py

@@ -6,62 +6,50 @@ from PIL import Image
 from roboflow import Roboflow
 from ultralytics import YOLO 
 
-def summary(model, input_size):
-    from torchsummary import summary
-    summary(model, input_size=input_size)
+# Initialize Roboflow
+rf = Roboflow(api_key="K1TXQnJq7EE7yoCf1g3C")
+project = rf.workspace().project("bone-fracture-detection-rkuqr")
+model = project.version("3").model
 
-def load_model(file_path):
-    # Load the Roboflow model
-    rf = Roboflow(api_key="K1TXQnJq7EE7yoCf1g3C")
-    project = rf.workspace("fyp-l87nq").project("bone-fracture-detection-rkuqr")
-    model = project.version(3).model
+def load_model():
+    # Load the model from Roboflow
+    yolov8_model = model.deploy()
 
-    # Load the model weights into a PyTorch model
-    pytorch_model = YOLO('data.yaml')
-    pytorch_model.load_state_dict(torch.load(file_path, map_location=torch.device('cpu')))
-    pytorch_model.eval()
+    return yolov8_model
 
-    return pytorch_model
+def predict_fracture(image_path):
+    # Load the model
+    yolov8_model = load_model()
 
-file_path = 'yolov8s.pt'
-model = load_model(file_path)
+    # Open the image
+    image = Image.open(image_path)
 
-def predict_fracture(image):
-    # Preprocess the image for the Roboflow model
-    img = Image.fromarray(image)
-    img_tensor = to_tensor(img).unsqueeze(0)  # Convert image to tensor and add batch dimension
+    # Perform inference
+    results = yolov8_model.predict(image_path, confidence=40, overlap=30)
 
-    # Perform inference with the Roboflow model
-    with torch.no_grad():
-        output = model(img_tensor)
-
-    # Postprocess the inference output
-    results = output[0]
+    # Display the results on the image
     img_with_boxes = image.copy()
-    for box in results:
-        label = int(box[5])
-        score = float(box[4])
-        if label == 0:  # Assuming 0 corresponds to the bone fracture class
-            color = "red" if score > 0.5 else "orange"  # Adjust the threshold as needed
-            xmin, ymin, xmax, ymax = box[:4].int().tolist()
+    for box in results["objects"]:
+        label = box["class"]
+        score = box["score"]
+        if label == "fracture":
+            color = "red"
+            xmin, ymin, xmax, ymax = box["bbox"]
             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 Image.fromarray((np.uint8(img_with_boxes)))
+    return 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(
-     fn=lambda *args, **kwargs: predict_fracture(args[0], load_model),
+    predict_fracture,
     inputs=gr.Image(),
     outputs=gr.Image(),
     live=True,
+    #capture_session=True,
     title="Bone Fracture Detection",
-    description="Upload an X-ray image to detect bone fractures using Roboflow's YOLOv8 model.",
+    description="Upload an X-ray image to detect bone fractures using YOLOv8.",
 )
 
 iface.launch()
+