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, image_path):
    """
    从标签文件中加载指定图片二维码嵌入坐标及所属类别
    :param watermark_txt: 标签文件
    :param image_path: 图片路径
    :return: [x1, y1, x2, y2, cls]
    """
    with open(watermark_txt, 'r') as f:
        for line in f.readlines():
            parts = line.strip().split()
            filename = parts[0]
            filename = os.path.basename(filename)
            if filename == os.path.basename(image_path):
                x_center, y_center, w, h = map(float, parts[1:5])
                cls = int(float(parts[5]))  # 转换类别为整数
                return [x_center, y_center, w, h, cls]
    return []