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):
    if image is None:
        return None
    
    # 运行推理
    results = model(image)
    result = results[0]
    annotated_image = result.plot()
    annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
    
    return annotated_image

# 创建自定义CSS
custom_css = """
#component-0 {
    max-width: 900px;
    margin: auto;
    padding: 20px;
    border-radius: 15px;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
}
.gradio-container {
    font-family: 'Source Sans Pro', sans-serif;
}
.gr-button {
    background: linear-gradient(45deg, #FF6B6B, #FF8E53) !important;
    border: none !important;
}
.gr-button:hover {
    transform: translateY(-2px);
    box-shadow: 0 5px 15px rgba(255, 107, 107, 0.4);
}
"""

# 创建Gradio界面
with gr.Blocks(css=custom_css) as demo:
    gr.Markdown(
        """
        # 🐉 奶龙杀手 (Nailong Killer)
        ## 一个基于YOLO的奶龙检测系统
        """
    )
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                label="上传图片 (Upload Image)", 
                type="numpy",
                tool="select"
            )
            with gr.Row():
                clear_btn = gr.Button("清除 (Clear)", variant="secondary")
                submit_btn = gr.Button("检测 (Detect)", variant="primary")
                
        with gr.Column():
            output_image = gr.Image(label="检测结果 (Detection Result)")
            
    gr.Markdown(
        """
        ### 使用说明 (Instructions):
        1. 点击上传或拖拽图片到左侧区域
        2. 点击"检测"按钮开始识别
        3. 右侧将显示检测结果
        
        ### 注意事项 (Notes):
        - 支持常见图片格式 (jpg, png, etc.)
        - 建议上传清晰的图片以获得更好的检测效果
        """
    )
    
    # 事件处理
    submit_btn.click(
        fn=detect_objects,
        inputs=input_image,
        outputs=output_image
    )
    
    clear_btn.click(
        lambda: (None, None),
        outputs=[input_image, output_image]
    )
    
    # 添加示例
    gr.Examples(
        examples=["example1.jpg", "example2.jpg"],
        inputs=input_image,
        outputs=output_image,
        fn=detect_objects,
        cache_examples=True
    )

# API端点
@app.post("/detect/")
async def detect_api(file: UploadFile = File(...)):
    contents = await file.read()
    image = Image.open(io.BytesIO(contents))
    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接口挂载到FastAPI
app = gr.mount_gradio_app(app, demo, path="/")

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