|
import gradio as gr |
|
import PIL.Image as Image |
|
from ultralytics import YOLO |
|
import spaces |
|
import os |
|
from huggingface_hub import hf_hub_download |
|
|
|
|
|
def get_model_path(model_name): |
|
model_cache_path = hf_hub_download( |
|
repo_id="atalaydenknalbant/budgerigar_gender_models", |
|
filename=model_name |
|
) |
|
return model_cache_path |
|
|
|
@spaces.GPU |
|
def yolo_inference(images, model_id, conf_threshold, iou_threshold, max_detection): |
|
model_path = get_model_path(model_id) |
|
model = YOLO(model_path) |
|
results = model.predict( |
|
source=images, |
|
conf=conf_threshold, |
|
iou=iou_threshold, |
|
imgsz=640, |
|
max_det=max_detection, |
|
show_labels=True, |
|
show_conf=True, |
|
) |
|
|
|
|
|
for r in results: |
|
image_array = r.plot() |
|
image = Image.fromarray(image_array[..., ::-1]) |
|
return image |
|
|
|
|
|
interface = gr.Interface( |
|
fn=yolo_inference, |
|
inputs=[ |
|
gr.Image(type="pil", label="Upload Image"), |
|
gr.Dropdown( |
|
choices=['budgerigar_yolo11x.pt', 'budgerigar_yolov9e.pt'], |
|
label="Model Name", |
|
value="budgerigar_yolo11x.pt", |
|
), |
|
gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence Threshold"), |
|
gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU Threshold"), |
|
gr.Slider(minimum=1, maximum=300, step=1, value=300, label="Max Detection"), |
|
], |
|
outputs=gr.Image(type="pil", label="Annotated Image"), |
|
cache_examples=True, |
|
title="Budgerigar Gender Determination", |
|
description="Upload image(s) for inference", |
|
examples=[ |
|
["Male.png", "budgerigar_yolov9e.pt", 0.25, 0.45, 300], |
|
["Female.png", "budgerigar_yolo11x.pt", 0.25, 0.45, 300], |
|
], |
|
) |
|
interface.launch() |