Update app.py

app.py

@@ -2,121 +2,44 @@ import gradio as gr
 import torch
 from ultralytics import YOLO 
 
-
-import io
-import requests
-from google.colab import auth
-from googleapiclient.http import MediaIoBaseDownload
-from googleapiclient.discovery import build
-
-def download_from_google_drive(id, mime_type):
-    auth.authenticate_user()
-    drive_service = build('drive', 'v3')
-    request = drive_service.files().export_media(fileId=id, mimeType=mime_type)
-    file = io.BytesIO()
-    downloader = MediaIoBaseDownload(file, request)
-    done = False
-    while done is False:
-        status, done = downloader.next_chunk()
-        print("Download %d%%." % int(status.progress() * 100))
-    file.seek(0)
-    return file
-
-def download_image_from_google_drive(shareable_link, image_name):
-    file_id = shareable_link.split("/")[5]
-    request = requests.get(shareable_link)
-    request.raise_for_status()
-    file_metadata = request.json()
-    mime_type = file_metadata["mimeType"]
-    image_file = download_from_google_drive(file_id, mime_type)
-    with open(image_name, "wb") as f:
-        f.write(image_file.read())
-
-download_image_from_google_drive('https://drive.google.com/file/d/1D0YH45dh5vubZsg54nAMBRw8RRKxlheC/view?usp=drive_link', 'one.jpg')
-download_image_from_google_drive('https://drive.google.com/file/d/18xJsjWJqSqA2ca6e2LO570rMoQeJKjGo/view?usp=drive_link', 'two.jpg')
-download_image_from_google_drive('https://drive.google.com/file/d/1PPQC0qkmoziNNEtKv7uPv9wh72hxNKMo/view?usp=drive_link', 'three.jpg')
-
-
-def yoloV8_func(image: gr.inputs.Image = None,
-                image_size: gr.inputs.Slider = 640,
-                conf_threshold: gr.inputs.Slider = 0.4,
-                iou_threshold: gr.inputs.Slider = 0.50):
-    """This function performs YOLOv8 object detection on the given image.
-
-    Args:
-        image (gr.inputs.Image, optional): Input image to detect objects on. Defaults to None.
-        image_size (gr.inputs.Slider, optional): Desired image size for the model. Defaults to 640.
-        conf_threshold (gr.inputs.Slider, optional): Confidence threshold for object detection. Defaults to 0.4.
-        iou_threshold (gr.inputs.Slider, optional): Intersection over Union threshold for object detection. Defaults to 0.50.
-    """
-    # Load the YOLOv8 model from the 'best.pt' checkpoint
-    model_path = "best.pt"
-    model = YOLO(model_path)
-
-    # Perform object detection on the input image using the YOLOv8 model
-    results = model.predict(image,
-                            conf=conf_threshold,
-                            iou=iou_threshold,
-                            imgsz=image_size)
-
-    # Print the detected objects' information (class, coordinates, and probability)
-    box = results[0].boxes
-    print("Object type:", box.cls)
-    print("Coordinates:", box.xyxy)
-    print("Probability:", box.conf)
-
-    # Custom rendering function using OpenCV
-    def render_custom(model, image, result):
-        class_ids = result.boxes.cls.cpu().numpy()
-        confidences = result.boxes.conf.cpu().numpy()
-        x1y1x2y2 = result.boxes.xyxy.cpu().numpy()
-
-        for i in range(len(class_ids)):
-            class_id = class_ids[i]
-            confidence = confidences[i]
-            x1, y1, x2, y2 = x1y1x2y2[i]
-
-            label = f"{model.names[class_id]} {confidence:.2f}"
-            color = (0, 255, 0)  # Green color for bounding boxes
-
-            cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
-            cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
-
-        return image
-
-    # Render the output image with bounding boxes around detected objects
-    custom_render = render_custom(model, image, results[0])
-    return custom_render
-
-
-inputs = [
-    gr.inputs.Image(type="filepath", label="Input Image"),
-    gr.inputs.Slider(minimum=320, maximum=1280, default=640,
-                     step=32, label="Image Size"),
-    gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.25,
-                     step=0.05, label="Confidence Threshold"),
-    gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.45,
-                     step=0.05, label="IOU Threshold"),
-]
-
-
-outputs = gr.outputs.Image(type="filepath", label="Output Image")
-
-title = "Bone Fracture Detection"
-
-
-examples = [['one.jpg', 640, 0.5, 0.7],
-            ['two.jpg', 800, 0.5, 0.6],
-            ['three.jpg', 900, 0.5, 0.8]]
-
-yolo_app = gr.Interface(
-    fn=yoloV8_func,
-    inputs=inputs,
-    outputs=outputs,
-    title=title,
-    examples=examples,
-    cache_examples=True,
+file_path = 'best.pt'
+
+def load_model(file_path):
+    # load the model weights from the file
+    yolov8_model = torch.load(file_path)
+
+    yolov8_model.eval()
+
+    return yolov8_model
+
+def predict_fracture(image):
+    # Preprocess the image for YOLOv8
+    img_tensor = to_tensor(image).unsqueeze(0)  # Convert image to tensor and add batch dimension
+    results = yolov8_model(img_tensor)  # Perform inference
+
+    # Display the results on the image
+    img_with_boxes = image.copy()
+    for box in results.xyxy[0]:
+        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()
+            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
+
+
+# Gradio Interface
+iface = gr.Interface(
+    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 YOLOv8.",
 )
 
-# Launch the Gradio interface in debug mode with queue enabled
-yolo_app.launch(debug=True, enable_queue=True)
+iface.launch()