import os


def expand_bbox(label_path, k=1.1):
    """
    扩大YOLO标签的边界框
    :param label_path: 标签文件路径（.txt）
    :param k: 扩大系数（k>1）
    """
    with open(label_path, 'r') as f:
        lines = f.readlines()

    new_lines = []
    for line in lines:
        line = line.strip()
        if not line:
            continue
        # 解析标签
        class_id, xc, yc, w, h = line.split()
        xc = float(xc)
        yc = float(yc)
        w = float(w)
        h = float(h)

        # 计算新宽高
        new_w = w * k
        new_h = h * k

        # 计算边界
        x1 = xc - new_w / 2
        y1 = yc - new_h / 2
        x2 = xc + new_w / 2
        y2 = yc + new_h / 2

        # 截断超出图像的部分（0~1范围）
        x1 = max(0.0, x1)
        y1 = max(0.0, y1)
        x2 = min(1.0, x2)
        y2 = min(1.0, y2)

        # 重新计算中心和宽高
        new_xc = (x1 + x2) / 2
        new_yc = (y1 + y2) / 2
        new_w = x2 - x1
        new_h = y2 - y1

        # 保留6位小数，拼接新标签
        new_line = f"{class_id} {new_xc:.6f} {new_yc:.6f} {new_w:.6f} {new_h:.6f}\n"
        new_lines.append(new_line)

    # 写入新标签（覆盖原文件，或改为新路径）
    with open(label_path, 'w') as f:
        f.writelines(new_lines)


# 批量处理文件夹中的所有标签
label_dir = r"D:\DataPreHandler\yuanshi_data\images\val\labels"  # 替换为你的标签文件夹路径
k = 1.25  # 扩大系数，根据需求调整

for filename in os.listdir(label_dir):
    if filename.endswith('.txt'):
        label_path = os.path.join(label_dir, filename)
        expand_bbox(label_path, k)

print("处理完成！")