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

# ---- FIX for PyTorch 2.6+ ----
torch.serialization.add_safe_globals([DetectionModel, nn.Sequential, Conv])

# ---- Load trained YOLO model ----
model = YOLO("res.pt")  # Ensure your model file is in the same folder

# ---- 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)

# ---- 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 human bone fractures using YOLOv8.<br><br>"
        "📸 **Dataset Source:** "
        "<a href='https://www.kaggle.com/datasets/jockeroika/human-bone-fractures-image-dataset' target='_blank'>"
        "Human Bone Fractures Image Dataset</a>"
    )
)

if __name__ == '__main__':
    iface.launch()