Spaces:
Runtime error
Runtime error
import gradio as gr | |
import cv2 | |
import requests | |
import os | |
from ultralytics import YOLO | |
# Define the directory containing the YOLO model files (*.pt) | |
model_weights_dir = 'best_weights' | |
# List all "*.pt" files in the specified directory | |
model_paths = [os.path.join(model_weights_dir, filename) for filename in os.listdir(model_weights_dir) if filename.endswith('.pt')] | |
# Initialize YOLO models based on the discovered model paths | |
models = [YOLO(model_path) for model_path in model_paths] | |
# Extract model names from paths (remove directory and ".pt" extension) | |
model_names = [os.path.splitext(os.path.basename(model_path))[0] for model_path in model_paths] | |
path = [['plot.JPG']] | |
def show_preds_image(image_path, selection): | |
image = cv2.imread(image_path) | |
outputs = models[selection].predict(source=image_path) | |
results = outputs[0].cpu().numpy() | |
for i, det in enumerate(results.boxes.xyxy): | |
cv2.rectangle( | |
image, | |
(int(det[0]), int(det[1])), | |
(int(det[2]), int(det[3])), | |
color=(0, 0, 255), | |
thickness=2, | |
lineType=cv2.LINE_AA | |
) | |
box_count = len(results) #counting the number of boxes detected | |
return cv2.cvtColor(image, cv2.COLOR_BGR2RGB), box_count | |
inputs = [ | |
gr.components.Image(type="filepath", label="Input Image"), | |
gr.components.Dropdown(choices=list(zip(model_names, range(len(models)))), label="Select Model"), | |
] | |
outputs_image = [ | |
gr.components.Image(type="numpy", label="Output Image"), | |
gr.components.Textbox(label="Box Count", type="auto"), | |
] | |
# Create the Gradio interface | |
interface_image = gr.Interface( | |
fn=show_preds_image, | |
inputs=inputs, | |
outputs=outputs_image, | |
title="Paddy Growth Stage Recognition", | |
description="Select an image and a YOLO model to detect the growth stage", | |
examples=path, | |
cache_examples=False, | |
) | |
gr.TabbedInterface( | |
[interface_image], | |
tab_names=['Image inference'] | |
).queue().launch() |