12345678910111213141516171819202122232425262728293031323334353637 |
- 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
|