import gradio as gr
from ultralytics import YOLO

# 加载模型
model = YOLO('nailong_yolo11.pt')

def predict(img):
    # 运行推理
    results = model.predict(img)
    # 返回第一个结果的绘制图像
    return results[0].plot()

# 创建 Gradio 演示
demo = gr.Interface(
    predict,
    inputs=[
        gr.Image(label="输入图片")
    ],
    outputs=[
        gr.Image(label="检测结果", type="numpy")
    ],
    title="🐉 奶龙杀手 (NailongKiller)",
    description="上传图片来检测奶龙 | Upload an image to detect Nailong",
    # examples=[
    #     ["example1.jpg"]
    # ],
    # cache_examples=True
)

# 启动应用
if __name__ == "__main__":
    demo.launch()