app.py

import gradio as gr
from ultralytics import YOLO
import torch
import cv2
import numpy as np
from fastapi import FastAPI, File, UploadFile
from PIL import Image
import io

# 初始化 FastAPI
app = FastAPI()

# 加载模型
model = YOLO("NailongKiller.yolo11n.pt")

def detect_objects(image):
    # 运行推理
    results = model(image)
    
    # 获取第一个结果
    result = results[0]
    
    # 在图像上绘制检测框
    annotated_image = result.plot()
    
    # 转换为RGB格式
    annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
    
    return annotated_image

# API端点
@app.post("/detect/")
async def detect_api(file: UploadFile = File(...)):
    # 读取上传的图片
    contents = await file.read()
    image = Image.open(io.BytesIO(contents))
    
    # 转换为numpy数组
    image_np = np.array(image)
    
    # 运行推理
    results = model(image_np)
    result = results[0]
    
    # 获取检测结果
    detections = []
    for box in result.boxes:
        detection = {
            "bbox": box.xyxy[0].tolist(),  # 边界框坐标
            "confidence": float(box.conf[0]),  # 置信度
            "class": int(box.cls[0])  # 类别
        }
        detections.append(detection)
    
    return {"detections": detections}

# 创建Gradio界面
demo = gr.Interface(
    fn=detect_objects,
    inputs=gr.Image(type="numpy"),
    outputs=gr.Image(),
    title="奶龙杀手 (Nailong Killer)",
    description="上传图片来检测奶龙 (Upload an image to detect Nailong)",
    examples=["example1.jpg", "example2.jpg"]
)

# 将Gradio接口挂载到FastAPI
app = gr.mount_gradio_app(app, demo, path="/")

# 启动应用
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)