import os


def parse_labels(file_path):
    categories = {}

    with open(file_path, 'r') as file:
        for line in file:
            parts = line.strip().split()
            file_path = parts[0]
            category = int(float(parts[-1]))

            if category not in categories:
                categories[category] = []
            categories[category].append(file_path)

    return categories


def load_watermark_info(watermark_txt, img_width, img_height):
    watermark_boxes = {}
    with open(watermark_txt, 'r') as f:
        for line in f.readlines():
            parts = line.strip().split()
            filename = parts[0]
            filename = os.path.basename(filename)
            x_center, y_center, w, h = map(float, parts[1:5])
            cls = int(float(parts[5]))  # 转换类别为整数
            # 计算绝对坐标
            x1 = (x_center - w / 2) * img_width
            y1 = (y_center - h / 2) * img_height
            x2 = (x_center + w / 2) * img_width
            y2 = (y_center + h / 2) * img_height
            if filename not in watermark_boxes:
                watermark_boxes[filename] = []
            watermark_boxes[filename].append([x1, y1, x2, y2, cls])
    return watermark_boxes