Spaces:
Sleeping
Sleeping
import gradio as gr | |
from ultralytics import YOLO | |
from PIL import Image | |
import numpy as np | |
import cv2 | |
# Load the YOLOv8 model (ensure this path is correct) | |
model = YOLO('yolov8n.pt') | |
def identify_disease(image): | |
# Convert the image to RGB if it's not | |
if image.mode != 'RGB': | |
image = image.convert('RGB') | |
# Perform inference | |
results = model(image) | |
predictions = results[0] | |
# Check if there are any detections | |
if len(predictions.boxes) == 0: | |
# No detections, return the image with a message | |
annotated_image = np.array(image) | |
cv2.putText(annotated_image, "No disease detected", (10, 30), | |
cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2) | |
annotated_image = Image.fromarray(annotated_image) | |
return annotated_image, [{"Disease": "None", "Confidence": "N/A"}] | |
# Extract predictions | |
boxes = predictions.boxes | |
labels = boxes.cls.cpu().numpy() | |
scores = boxes.conf.cpu().numpy() | |
class_names = model.names | |
# Annotate image with bounding boxes and labels | |
annotated_image = np.array(image) | |
for box, label, score in zip(boxes.xyxy.cpu().numpy(), labels, scores): | |
x1, y1, x2, y2 = map(int, box) | |
class_name = class_names[int(label)] | |
confidence = f"{score * 100:.2f}%" | |
annotated_image = cv2.rectangle(annotated_image, (x1, y1), (x2, y2), (0, 255, 0), 2) | |
annotated_image = cv2.putText(annotated_image, f"{class_name} {confidence}", (x1, y1 - 10), | |
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) | |
# Convert annotated image back to PIL format | |
annotated_image = Image.fromarray(annotated_image) | |
# Prepare results for display | |
results_list = [{"Disease": class_names[int(label)], "Confidence": f"{score * 100:.2f}%"} for label, score in zip(labels, scores)] | |
return annotated_image, results_list | |
# Define Gradio interface with updated syntax | |
interface = gr.Interface( | |
fn=identify_disease, | |
inputs=gr.Image(type="pil"), | |
outputs=[ | |
gr.Image(type="pil", label="Annotated Image"), | |
gr.Dataframe(headers=["Disease", "Confidence"], label="Predictions") | |
], | |
title="Leaf Disease Identifier with YOLOv8", | |
description="Upload an image of a leaf, and this tool will identify the disease with confidence scores." | |
) | |
# Launch the app | |
interface.launch() | |