|
|
@@ -0,0 +1,37 @@
|
|
|
+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
|
