from ultralytics import YOLO
import cv2
import os

def detect_and_save(image_path, output_dir="result_yolov8", save_txt=False):
    """
    使用 YOLOv8 检测图片中的目标并保存结果
    
    参数:
        image_path: 输入图片路径
        output_dir: 结果保存目录
    """
    # 加载预训练模型 (自动下载 yolov8n.pt 若不存在)
    model = YOLO("./yolov8n.pt")
    
    # 进行目标检测
    results = model(image_path)
    
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 处理每个检测结果 (单张图片返回一个结果列表)
    for i, r in enumerate(results):
        # 1. 保存带标注的图片
        annotated_img = r.plot()  # 生成带标注框的BGR图像
        img_name = os.path.basename(image_path)
        img_output_path = os.path.join(output_dir, f"annotated_{img_name}.jpg")
        print(f"保存带标注的图片：{img_output_path}")
        cv2.imwrite(img_output_path, annotated_img)
        print(f"✅ 标注图片已保存至: {img_output_path}")
        
        if save_txt:
            # 2. 保存检测结果(txt文件，YOLO格式)
            txt_path = os.path.join(output_dir, f"results_{os.path.splitext(img_name)[0]}.txt")
            r.save_txt(txt_path)
            print(f"📝 检测结果已保存至: {txt_path}")
        
        # 3. 控制台打印检测信息
        boxes = r.boxes
        print("\n检测到的目标:")
        print(f"{'类别':<10}{'置信度':<15}{'坐标(xywh)':<30}")
        print("-" * 55)
        for box in boxes:
            cls_id = int(box.cls)
            class_name = model.names[cls_id]

            # 统计类别信息
            class_set.add(class_name)

            conf = float(box.conf)
            bbox = [int(coord) for coord in box.xywh[0]]  # 获取xywh格式坐标
            print(f"{class_name:<10}{conf:.4f}{'':<5}{bbox}")
        print("-" * 55)


def detect_file(path_file, output_dir):
    t = os.listdir(path_file)
    for i in t:
        detect_and_save(os.path.join(path_file, i), os.path.join(output_dir, path_file.split("/")[-1]))


def detect_file_list(path_file, output_dir):
    t = os.listdir(path_file)
    for i in t:
        detect_file(os.path.join(path_file, i), output_dir)
    

class_set = set()

if __name__ == "__main__":
    import os
    
    path_file = "data_image"
    output_dir = "result_yolov8"

    detect_file_list(path_file, output_dir)

    print(class_set)

    # # 输入图片路径 (可替换为你的图片路径)
    # input_image = "data_image/1e4c75b76e531606e2adc491a8f09ae8/frame_000150.jpg"  # 确保图片存在
    # output_dir = "result_yolov8"

    
    # # 执行检测并保存结果
    # detect_and_save(input_image, output_dir)