kylecsnow's picture
Update app.py
432d40e verified
from ultralytics import YOLO
import gradio as gr
model = YOLO("./runs/detect/train18/weights/best.pt")
def yolo_predict(image):
"""Run YOLOv8 inference and return annotated image with results"""
results = model(image)
print(results)
annotated_image = results[0].plot()
# Get prediction details
boxes = results[0].boxes
prediction_details = []
for box in boxes:
class_id = int(box.cls[0].item())
class_name = model.names[class_id]
confidence = round(box.conf[0].item(), 2)
coords = box.xyxy[0].tolist() # [x1, y1, x2, y2]
prediction_details.append({
"class": class_name,
"confidence": confidence,
"bbox": coords
})
return annotated_image
# return annotated_image, prediction_details
with gr.Blocks() as demo:
gr.Markdown("# YOLOv8 Object Detection")
gr.Markdown(
"""
This application uses a YOLOv8m model fine-tuned specifically to detect red blood
cells, white blood cells, and platelets in images of blood cells. This version
was trained using the `keremberke/blood-cell-object-detection` dataset on huggingface.com.
"""
)
gr.Interface(
fn=yolo_predict,
inputs=gr.Image(label="Input Image",type="pil"),
outputs=[
gr.Image(label="Detected Objects"),
# gr.JSON(label="Detection Details")
],
# title="YOLOv8 Object Detection",
# # description="Upload an image to detect objects using YOLOv8",
description='Select an example image below (none of which were included in model training or validation), or upload your own image. Then, click "Submit" to see the model in action.',
examples=[
"./bloodcell-examples/image_0.jpg",
"./bloodcell-examples/image_1.jpg",
"./bloodcell-examples/image_2.jpg",
"./bloodcell-examples/image_3.jpg",
"./bloodcell-examples/image_4.jpg",
],
)
demo.launch(
show_error=True,
height=900,
width="80%",
# width="100%",
# share=True,
)