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):
    # 将输入图像转换为PIL Image对象
    input_image = Image.fromarray(img)
    original_size = input_image.size
    
    # 计算填充尺寸
    max_size = max(original_size)
    pad_w = max_size - original_size[0]
    pad_h = max_size - original_size[1]
    
    # 创建方形画布并保持宽高比
    padded_img = Image.new('RGB', (max_size, max_size), (114, 114, 114))  # 使用灰色填充
    padded_img.paste(input_image, (pad_w//2, pad_h//2))
    
    # 转换为numpy数组并进行预测
    img_array = np.array(padded_img)
    
    # 进行预测
    results = model.predict(
        img_array,
        conf=0.30,  # 置信度阈值
        iou=0.45,   # IOU阈值
        max_det=20 # 最大检测数量
    )
    
    # 获取预测结果
    result_img = results[0].plot()
    
    # 裁剪回原始尺寸
    if pad_w > 0 or pad_h > 0:
        result_img = result_img[pad_h//2:pad_h//2 + original_size[1], 
                              pad_w//2:pad_w//2 + original_size[0]]
    
    return result_img

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)