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('nailong_yolo11.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)
    result = results[0]
    
    # 获取预测结果
    result_img = result.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]]
    
    # 处理检测信息
    info = {
        "detected": len(result.boxes) > 0,
        "count": len(result.boxes),
        "confidences": []
    }
    
    if info["detected"]:
        # 获取每个检测框的置信度
        for box in result.boxes:
            conf = float(box.conf[0])
            info["confidences"].append(f"{conf:.2%}")
    
    # 生成输出文本
    output_text = f"""检测结果:
- 是否检测到目标: {'是' if info['detected'] else '否'}
- 检测到的目标数量: {info['count']}"""
    
    if info["confidences"]:
        output_text += f"\n- 置信度: {', '.join(info['confidences'])}"
    
    return result_img, output_text

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(label="输入图片"),
    outputs=[
        gr.Image(label="检测结果", type="numpy"),
        gr.Textbox(label="检测信息")
    ],
    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)