import gradio as gr
from ultralytics import YOLO
from PIL import Image
import io

# 加载 YOLOv8 模型
model = YOLO('best.pt')  # 替换为您的模型文件路径

def detect_objects(image):
    # 使用 YOLOv8 模型进行预测
    results = model(image)

    # 处理预测结果
    class_counts = {}
    output_text = ''
    annotated_image = Image.fromarray(results[0].plot())  # 将 NumPy 数组转换为 PIL 图像

    # 将 PIL 图像转换为 bytes 以供 Gradio 显示
    buffered = io.BytesIO()
    annotated_image.save(buffered, format="PNG")
    annotated_image_bytes = buffered.getvalue()

    # 将 bytes 转换为 PIL 图像对象
    annotated_image = Image.open(io.BytesIO(annotated_image_bytes))

    for result in results[0].boxes.data:
        class_id = int(result[5])
        class_name = model.names[class_id]
        confidence = result[4]
        x1, y1, x2, y2 = result[:4].cpu().numpy().astype(int)

        # 统计类别数量
        if class_name in class_counts:
            class_counts[class_name] += 1
        else:
            class_counts[class_name] = 1

        # 构建输出文本
        output_text += f'类别: {class_name}, 置信度: {confidence:.2f}, 坐标: ({x1}, {y1}), ({x2}, {y2})\n'

    # 构建每个类别数量的输出
    class_count_text = '\n'.join([f'{class_name}: {count}' for class_name, count in class_counts.items()])
    output_text += f'\n每个类别的数量:\n{class_count_text}'

    return annotated_image, output_text

# 创建 Gradio 界面
demo = gr.Interface(
    fn=detect_objects,
    inputs=gr.Image(type="pil"),
    outputs=[gr.Image(type="pil"), "text"],
    title="细胞检测",
    description="上传图片,进行物体检测并显示结果"
)

# 启动应用
demo.launch()