import os
from ultralytics import YOLO

def auto_label(image_dir, output_dir, model_path='model/best.pt', conf=0.5, save_segmentation=True):
    """
    使用 YOLO 模型自动标注图片并生成 TXT 标签文件。
    """
    # 检查模型文件是否存在
    if not os.path.exists(model_path):
        print(f"错误: 模型文件不存在 -> {model_path}")
        # 尝试使用根目录下的 yolov8n-seg.pt 作为备选
        fallback_model = 'yolov8n-seg.pt'
        if os.path.exists(fallback_model):
            print(f"尝试使用备选模型 -> {fallback_model}")
            model_path = fallback_model
        else:
            return

    # 加载模型
    print(f"正在加载模型: {model_path} ...")
    model = YOLO(model_path)

    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"创建输出目录: {output_dir}")

    # 支持的图片格式
    img_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.webp')
    
    # 获取图片列表
    image_files = [f for f in os.listdir(image_dir) if f.lower().endswith(img_extensions)]
    total_files = len(image_files)
    
    print(f"找到 {total_files} 张图片，开始自动标注...")

    count = 0
    for idx, img_file in enumerate(image_files):
        img_path = os.path.join(image_dir, img_file)
        txt_filename = os.path.splitext(img_file)[0] + ".txt"
        txt_path = os.path.join(output_dir, txt_filename)

        # 推理
        # stream=True 可以节省内存，但对于单张处理差异不大
        results = model(img_path, conf=conf, verbose=False)

        for result in results:
            with open(txt_path, 'w') as f:
                # 优先尝试保存分割掩码
                if save_segmentation and result.masks is not None:
                    # result.masks.xyn 获取归一化的多边形坐标片段
                    if hasattr(result.masks, 'xyn'):
                        for i, seg in enumerate(result.masks.xyn):
                            cls = int(result.boxes.cls[i]) # 对应的类别
                            # 格式: class x1 y1 x2 y2 ...
                            # flatten() 将数组展平，tolist() 转为列表
                            coords = " ".join([f"{p[0]:.6f} {p[1]:.6f}" for p in seg])
                            f.write(f"{cls} {coords}\n")
                    else:
                        pass # 兼容性处理
                
                # 如果没有分割结果或未启用分割，保存检测框
                elif result.boxes is not None:
                    for box in result.boxes:
                        cls = int(box.cls)
                        # xywhn: x_center, y_center, width, height (normalized)
                        x, y, w, h = box.xywhn[0].tolist()
                        f.write(f"{cls} {x:.6f} {y:.6f} {w:.6f} {h:.6f}\n")
        
        count += 1
        if count % 10 == 0:
            print(f"进度: {count}/{total_files}")

    print(f"完成！已生成 {count} 个标签文件。保存位置: {output_dir}")

if __name__ == "__main__":
    # --- 配置区域 (请在此修改路径) ---
    
    # 待标注图片文件夹 (例如: d:\data\20251204)
    IMAGE_DIR = r'd:\data\20251204' 
    
    # 标签保存文件夹 (例如: d:\data\20251204\labels_auto)
    OUTPUT_DIR = r'd:\data\20251204\labels_auto' 
    
    # 模型路径 (例如: d:\data\model\best.pt)
    MODEL_PATH = r'd:\data\best.pt'
    
    # 置信度阈值
    CONF_THRESHOLD = 0.5
    
    # 是否保存分割点 (True=保存多边形, False=只保存矩形框)
    SAVE_SEGMENTATION = True 
    # --------------------------------
    
    # 检查输入目录是否存在，防止报错
    if not os.path.exists(IMAGE_DIR):
        print(f"错误: 图片目录不存在 -> {IMAGE_DIR}")
        print("请打开脚本修改 IMAGE_DIR 为正确的路径。")
    else:
        auto_label(IMAGE_DIR, OUTPUT_DIR, MODEL_PATH, CONF_THRESHOLD, SAVE_SEGMENTATION)