Danung's picture
Add new file
06f4efc
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()