Spaces:
Sleeping
Sleeping
File size: 1,570 Bytes
11c00e3 4c65693 0475832 11c00e3 4c65693 11c00e3 4c65693 11c00e3 4c65693 11c00e3 4c65693 11c00e3 4c65693 0475832 11c00e3 0475832 11c00e3 0475832 11c00e3 0475832 11c00e3 0475832 11c00e3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | from PIL import Image, ImageDraw, ImageFont
import gradio as gr
from ultralytics import YOLO
# Class mapping and colors
class_names = {0: "weed", 1: "wheat"}
class_colors = {
"weed": "red",
"wheat": "green"
}
# Load ONNX model
model = YOLO("best.onnx")
def detect(image):
original_image = image.copy()
results = model.predict(image, imgsz=640)[0]
preds = results.boxes.data
draw = ImageDraw.Draw(original_image)
try:
font = ImageFont.truetype("arial.ttf", 16)
except:
font = ImageFont.load_default()
output = []
for box in preds.cpu():
x1, y1, x2, y2, conf, cls = box.tolist()
if conf >= 0.5:
class_id = int(cls)
class_name = class_names.get(class_id, "unknown")
color = class_colors.get(class_name, "blue")
label = f"{class_name} ({conf:.2f})"
draw.rectangle([x1, y1, x2, y2], outline=color, width=2)
draw.text((x1, y1 - 10), label, fill=color, font=font)
output.append({
"x1": int(x1), "y1": int(y1),
"x2": int(x2), "y2": int(y2),
"confidence": round(conf, 2),
"class": class_name
})
return output # Return only JSON for Android API usage
# Launch Gradio as API
gr.Interface(
fn=detect,
inputs=gr.Image(type="pil"),
outputs=gr.JSON(label="Detection Results"),
title="YOLOv11M Weed Detection API",
description="Send an image to detect weed (🔴) and wheat (🟢).",
allow_flagging="never"
).launch()
|