app.py

import gradio as gr
from ultralytics import YOLO
from fastapi import FastAPI
from PIL import Image
import torch
import spaces
import numpy as np

app = FastAPI()
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = YOLO('NailongKiller.yolo11n.pt').to(device)

@spaces.GPU
def predict(img):
    # 优化图像预处理
    img_resized = np.array(Image.fromarray(img).resize((640, 640)))
    # 规范化像素值到 0-1 范围
    img_tensor = torch.tensor(img_resized, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).div(255.0).to(device)
    
    # 设置模型预测参数以加快后处理速度
    results = model.predict(
        img_tensor,
        conf=0.50,  # 提高置信度阈值
        iou=0.45,   # 调整 IOU 阈值
        max_det=100 # 限制最大检测数量
    )
    return results[0].plot()

demo = gr.Interface(
    fn=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(server_name="0.0.0.0", server_port=7860)