ma4389's picture
Update app.py
8082eee verified
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()