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, "No image provided"
    
    try:
        # 运行推理
        results = model(image)
        result = results[0]
        
        # 在图像上绘制检测框
        annotated_image = result.plot()
        annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
        
        # 获取检测结果统计
        num_detections = len(result.boxes)
        detection_info = f"检测到 {num_detections} 个目标"
        
        return annotated_image, detection_info
    except Exception as e:
        return None, f"Error: {str(e)}"

# 创建主题和样式
theme = gr.themes.Soft(
    primary_hue="indigo",
    secondary_hue="blue",
).set(
    body_background_fill="*neutral_50",
    block_background_fill="*neutral_100",
    block_label_background_fill="*primary_100",
    block_label_text_color="*primary_500",
    button_primary_background_fill="*primary_500",
    button_primary_background_fill_hover="*primary_600",
    button_primary_text_color="white",
    border_color_primary="*primary_300",
)

with gr.Blocks(theme=theme) as demo:
    gr.Markdown(
        """
        # 🐉 奶龙杀手 (NailongKiller)
        
        这是一个基于 YOLO 的奶龙检测系统。上传图片即可自动检测图中的奶龙。
        
        This is a YOLO-based Nailong detection system. Upload an image to detect Nailong automatically.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(
                label="输入图片 | Input Image",
                type="numpy",
                height=512,
                width=512,
            )
            
            with gr.Row():
                clear_btn = gr.Button("清除 | Clear", variant="secondary", size="lg")
                detect_btn = gr.Button("检测 | Detect", variant="primary", size="lg")
        
        with gr.Column(scale=1):
            output_image = gr.Image(
                label="检测结果 | Detection Result",
                height=512,
                width=512,
            )
            result_text = gr.Textbox(
                label="检测信息 | Detection Info",
                placeholder="等待检测...",
            )
    
    gr.Markdown(
        """
        ### 📝 使用说明 | Instructions
        
        1. 点击上传或拖拽图片到左侧输入区域
        2. 点击"检测"按钮开始识别
        3. 右侧将显示检测结果和统计信息
        
        ### ⚠️ 注意事项 | Notes
        
        - 支持常见图片格式 (jpg, png, etc.)
        - 建议上传清晰的图片以获得更好的检测效果
        - 图片会自动调整大小以优化性能
        
        ### 🔗 相关链接 | Links
        
        - [项目地址 | Project Repository](https://huggingface.co/spaces/Hakureirm/NailongKiller)
        - [YOLO Documentation](https://docs.ultralytics.com/)
        """
    )
    
    # 事件处理
    detect_btn.click(
        fn=detect_objects,
        inputs=input_image,
        outputs=[output_image, result_text],
    )
    
    clear_btn.click(
        lambda: (None, None, None),
        outputs=[input_image, output_image, result_text],
    )
    
    # 添加示例
    if os.path.exists("example1.jpg") and os.path.exists("example2.jpg"):
        gr.Examples(
            examples=["example1.jpg", "example2.jpg"],
            inputs=input_image,
            outputs=[output_image, result_text],
            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)