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

# 初始化 FastAPI 和模型
app = FastAPI()

# 检查 GPU 是否可用，并选择设备
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = YOLO('NailongKiller.yolo11n.pt').to(device)

@spaces.GPU
def predict(img):
    # 将 PIL 图像转换为 numpy 数组
    img_resized = np.array(Image.fromarray(img).resize((640, 640)))
    # 将 numpy 数组转换为 PyTorch 张量
    img_tensor = torch.tensor(img_resized, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(device)
    results = model.predict(img_tensor)
    return results[0].plot()

# Gradio 界面
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)