import gradio as gr
import torch
import yaml
import numpy as np
from PIL import Image
from roboflow import Roboflow
from ultralytics import YOLO 

# Initialize Roboflow
rf = Roboflow(api_key="K1TXQnJq7EE7yoCf1g3C")
project = rf.workspace().project("bone-fracture-detection-rkuqr")
model = project.version("3").model

def load_model():
    # Load the model from Roboflow
    yolov8_model = model.deploy()

    return yolov8_model

def predict_fracture(image_path):
    # Load the model
    yolov8_model = load_model()

    # Open the image
    image = Image.open(image_path)

    # Perform inference
    results = yolov8_model.predict(image_path, confidence=40, overlap=30)

    # Display the results on the image
    img_with_boxes = image.copy()
    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 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.",
)

iface.launch()