Spaces:
Runtime error
Runtime error
File size: 1,697 Bytes
cf1ca8a f6f8722 cf1ca8a |
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 |
import torch
import gradio as gr
from PIL import Image, ImageDraw
import numpy as np
import torchvision.transforms as T
# Load your model
model = torch.hub.load('ultralytics/yolov8', 'custom', path='best.pt') # Adjust path if necessary
model.eval()
# Define your classes
classes = [
"Apple", "Banana", "Beetroot", "Bitter_Gourd", "Bottle_Gourd", "Cabbage",
"Capsicum", "Carrot", "Cauliflower", "Cherry", "Chilli", "Coconut",
"Cucumber", "EggPlant", "Ginger", "Grape", "Green_Orange", "Kiwi",
"Maize", "Mango", "Melon", "Okra", "Onion", "Orange", "Peach", "Pear",
"Peas", "Pineapple", "Pomegranate", "Potato", "Radish", "Strawberry",
"Tomato", "Turnip", "Watermelon", "walnut", "almond"
]
# Define the inference function
def detect(image):
# Transform the image to tensor
transform = T.Compose([T.ToTensor()])
input_tensor = transform(image).unsqueeze(0)
# Perform inference
with torch.no_grad():
detections = model(input_tensor)[0]
# Draw bounding boxes and labels on the image
draw = ImageDraw.Draw(image)
for detection in detections:
# Each detection includes [x1, y1, x2, y2, confidence, class]
x1, y1, x2, y2, conf, cls = detection
if conf >= 0.5: # Consider detections with confidence >= 0.5
label = classes[int(cls)]
draw.rectangle(((x1, y1), (x2, y2)), outline="red", width=2)
draw.text((x1, y1), f"{label} ({conf:.2f})", fill="red")
return np.array(image)
# Create a Gradio interface
iface = gr.Interface(
fn=detect,
inputs=gr.inputs.Image(source="webcam", tool="editor"),
outputs="image",
live=True,
)
# Launch the app
iface.launch()
|