|
|
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) |
|
|
|
|
|
|
|
|
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() |
