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