Spaces:
Runtime error
Runtime error
import gradio as gr | |
import cv2 | |
import os | |
from ultralyticsplus import YOLO, render_result | |
model_path= 'best.pt' | |
def preds_image(image, conf_thres, iou_thres): | |
model = YOLO(model_path) | |
result = model.predict(image, | |
conf= conf_thres, | |
iou= iou_thres) | |
box = result[0].boxes | |
print("Object type: ", box.cls) | |
print("Coordinates: ", box.xyxy) | |
print("Probability: ", box.conf) | |
render = render_result(model=model, image=image, result=result[0]) | |
return render | |
inputs_image = [ | |
gr.Image(label="Input Image"), | |
gr.Slider(minimum=0.0, maximum=1.0, value=0.40, step=0.05, label="Confidence threshold"), | |
gr.Slider(minimum=0.0, maximum=1.0, value=0.45, step=0.05, label="IOU threshold") | |
] | |
outputs_image = [ | |
gr.Image(label="Output Image"), | |
] | |
interface_image = gr.Interface( | |
fn=preds_image, | |
inputs=inputs_image, | |
outputs=outputs_image, | |
title="Chili Leaf Disease Detector (Image)" | |
) | |
def preds_video(video): | |
video_path = video | |
video_path_out = '{}_out.mp4'.format(video_path) | |
cap = cv2.VideoCapture(video_path) | |
ret, frame = cap.read() | |
H, W, _ = frame.shape | |
out = cv2.VideoWriter(video_path_out, cv2.VideoWriter_fourcc(*'MP4V'), int(cap.get(cv2.CAP_PROP_FPS)), (W, H)) | |
model= YOLO(model_path) | |
threshold = 0.4 | |
while ret: | |
results = model(frame)[0] | |
for result in results.boxes.data.tolist(): | |
x1, y1, x2, y2, score, class_id = result | |
if score > threshold: | |
cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 0, 255), 4) | |
#object details | |
org = (int(x1), int(y1 - 10)) | |
font = cv2.FONT_HERSHEY_SIMPLEX | |
fontscale = 0.5 | |
color = (0, 0, 0) | |
cv2.putText(frame, results.names[int(class_id)].lower(), org, | |
font, fontscale, color, 1, cv2.LINE_AA) | |
out.write(frame) | |
ret, frame = cap.read() | |
cap.release() | |
out.release() | |
cv2.destroyAllWindows() | |
return video_path_out | |
inputs_video = gr.Video(label= 'Original chili leaf video') | |
outputs_video = gr.Video(label= 'Predicted leaf') | |
interface_video = gr.Interface( | |
fn=preds_video, | |
inputs=inputs_video, | |
outputs=outputs_video, | |
title="Chili Leaf Disease Detector (Video)" | |
) | |
gr.TabbedInterface( | |
[interface_image, interface_video], | |
tab_names=['Image inference', 'Video inference'] | |
).queue().launch() |