from ultralytics import YOLO
import cv2
import numpy as np

def detect_and_visualize(image_path, model_path):
    # 加载YOLOv8模型
    model = YOLO(model_path)  # 例如 'yolov8n.pt', 'yolov8s.pt' 等
    
    # 读取图片
    image = cv2.imread(image_path)
    
    # 运行检测
    results = model(image)
    
    # 获取第一帧的结果
    result = results[0]
    
    # 在原图上绘制检测结果
    for box in result.boxes:
        # 获取边界框坐标
        x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
        
        # 获取置信度
        conf = float(box.conf[0])
        
        # 获取类别ID和名称
        cls_id = int(box.cls[0])
        cls_name = result.names[cls_id]
        
        # 为每个类别生成不同的颜色
        color = tuple(np.random.randint(0, 255, 3).tolist())
        
        # 绘制边界框
        cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
        
        # 准备标签文本
        label = f'{cls_name} {conf:.2f}'
        
        # 计算标签大小
        (label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        
        # 绘制标签背景
        cv2.rectangle(image, (x1, y1-label_height-5), (x1+label_width, y1), color, -1)
        
        # 绘制标签文本
        cv2.putText(image, label, (x1, y1-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
    
    # 保存结果图片
    output_path = 'output_detected.jpg'
    cv2.imwrite(output_path, image)
    print(f"检测结果已保存至: {output_path}")

# 使用示例
if __name__ == "__main__":
    image_path = "./test_math.png"  # 替换为你的图片路径
    model_path = "docgenome_object_detection_yolov8.pt"  # 替换为你的模型权重路径
    detect_and_visualize(image_path, model_path)