File size: 1,978 Bytes
07d1802 218faaa 07d1802 0e67741 07d1802 0664949 07d1802 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
import gradio as gr
import PIL.Image as Image
from ultralytics import YOLO
import spaces
import os
from huggingface_hub import hf_hub_download
# Helper function to download models from Hugging Face
def get_model_path(model_name):
model_cache_path = hf_hub_download(
repo_id="atalaydenknalbant/budgerigar_yolo_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) # Download model
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,
)
# Process results and convert to PIL Image
for r in results:
image_array = r.plot()
image = Image.fromarray(image_array[..., ::-1])
return image
# Define Gradio interface
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="Pretrained YOLO models for determining budgerigar gender based on cere color variations. 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() |