sionova's picture
Add project files
c1b46d3
import gradio as gr
from ultralytics import YOLO
from PIL import Image
CONF_THRESH=0.332
IOU=0.4
model = YOLO("best.pt")
def predict(image):
results = model(image, verbose=False, conf=CONF_THRESH, iou=IOU, agnostic_nms=True, rect=False)
results[0].names = {0: 'bee', 1: 'stylopidae', 2: 'meloidae larvae'}
bgr_result_img = results[0].plot()
rgb_result_img = Image.fromarray(bgr_result_img[..., ::-1])
return rgb_result_img
def main():
with gr.Blocks() as demo:
gr.HTML("<h1>Deep Learning Based Detection of Wild Bee Parasites under Natural Conditions</h1>")
with gr.Row():
image_input = gr.Image(type="filepath", label="Upload Image")
image_output = gr.Image(label="Prediction")
gr.Examples(
examples=["stylopidae.jpg", "larvae.jpg"],
inputs=image_input
)
submit_btn = gr.Button("Detect")
submit_btn.click(fn=predict, inputs=image_input, outputs=[image_output])
gr.HTML("""
<div>
Sources:
<ul>
<li><a href="https://observation.org/photos/111615265/">Stylopidae</a></li>
<li><a href="https://observation.org/photos/67773412/">Meloidae larvae</a></li>
</ul>
</div>
<img src="https://upload.wikimedia.org/wikipedia/commons/9/98/TU_Ilmenau_Logo_black_green.svg" alt="TU Ilmenau" style="position: absolute; bottom: 0; right: 0; width: 20vw; max-width: 250px; height: auto;">
""")
demo.launch()
if __name__ == '__main__':
main()