YOLO数据处理：通过iou合并yolo标签

import os

from tqdm import tqdm


def calculate_iou(box1, box2):
    # box = [center_x, center_y, width, height]
    # 将 YOLO 格式转换为 xmin, ymin, xmax, ymax
    box1_x1 = box1[0] - box1[2] / 2
    box1_y1 = box1[1] - box1[3] / 2
    box1_x2 = box1[0] + box1[2] / 2
    box1_y2 = box1[1] + box1[3] / 2

    box2_x1 = box2[0] - box2[2] / 2
    box2_y1 = box2[1] - box2[3] / 2
    box2_x2 = box2[0] + box2[2] / 2
    box2_y2 = box2[1] + box2[3] / 2

    # 计算交集
    inter_x1 = max(box1_x1, box2_x1)
    inter_y1 = max(box1_y1, box2_y1)
    inter_x2 = min(box1_x2, box2_x2)
    inter_y2 = min(box1_y2, box2_y2)

    inter_area = max(0, inter_x2 - inter_x1) * max(0, inter_y2 - inter_y1)

    # 计算并集
    box1_area = (box1_x2 - box1_x1) * (box1_y2 - box1_y1)
    box2_area = (box2_x2 - box2_x1) * (box2_y2 - box2_y1)

    union_area = box1_area + box2_area - inter_area

    # 返回IOU值
    return inter_area / union_area if union_area > 0 else 0


def merge_two_files(labels1, labels2, iou_threshold=0.5):
    merged = []
    used_labels2 = set()  # 记录已使用的 labels2 索引

    for box1 in labels1:
        merged_box = box1
        for j, box2 in enumerate(labels2):
            if j in used_labels2:
                continue

            # 只合并同类的框
            if box1[0] == box2[0]:
                iou = calculate_iou(box1[1:], box2[1:])
                if iou >= iou_threshold:
                    # 合并框，通过平均它们的位置和大小
                    merged_box = [
                        box1[0],  # 保持 class_id 不变
                        (box1[1] + box2[1]) / 2,
                        (box1[2] + box2[2]) / 2,
                        (box1[3] + box2[3]) / 2,
                        (box1[4] + box2[4]) / 2
                    ]
                    used_labels2.add(j)  # 标记 box2 已使用
                    break

        merged.append(merged_box)

    # 添加 labels2 中未使用的框
    for i, box2 in enumerate(labels2):
        if i not in used_labels2:
            merged.append(box2)

    return merged


def process_two_files(file1_path, file2_path, output_path, iou_threshold=0.5):
    # 读取两个标签文件，如果存在
    labels1 = []
    labels2 = []

    if os.path.exists(file1_path):
        with open(file1_path, 'r') as f1:
            labels1 = [list(map(float, line.split())) for line in f1]

    if os.path.exists(file2_path):
        with open(file2_path, 'r') as f2:
            labels2 = [list(map(float, line.split())) for line in f2]

    # 如果存在标签，则进行合并
    if labels1 or labels2:  # 确保有标签可合并
        merged_labels = merge_two_files(labels1, labels2, iou_threshold)

        # 将合并后的标签保存到输出文件
        with open(output_path, 'w') as f_out:
            for label in merged_labels:
                f_out.write(' '.join(map(str, label)) + '\n')
    else:
        print(f"{file1_path} 和 {file2_path} 中没有找到标签。")


def process_folder(input_folder1, input_folder2, output_folder, iou_threshold=0.5):
    # 如果输出文件夹不存在，创建该文件夹
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # 遍历第一个文件夹中的所有 .txt 文件
    for filename in tqdm(os.listdir(input_folder1)):
        if filename.endswith('.txt'):
            file1_path = os.path.join(input_folder1, filename)
            file2_path = os.path.join(input_folder2, filename)  # 根据名称匹配文件
            output_path = os.path.join(output_folder, filename)

            # 如果第二个文件夹中没有对应文件，则直接输出第一个文件夹中的内容
            if not os.path.exists(file2_path):
                with open(file1_path, 'r') as f_in, open(output_path, 'w') as f_out:
                    f_out.write(f_in.read())
            else:
                # 处理并合并文件
                process_two_files(file1_path, file2_path, output_path, iou_threshold)


if __name__ == '__main__':
    input_folder1 = r"D:\BaiduNetdiskDownload\txt"
    input_folder2 = r"D:\BaiduNetdiskDownload\predict_labels"
    output_folder = r"out"

    # 执行批量处理
    process_folder(input_folder1, input_folder2, output_folder, iou_threshold=0.7)