yolo数据集标记

2025-12-25 17:41:18 +08:00
parent 1a135cdda7
commit a0555ac6b5
22 changed files with 10194 additions and 1 deletions
@@ -0,0 +1,159 @@
+import logging
+import numpy as np
+import config
+
+
+def __convert_bbox_to_text(bbox, scale, x_max, y_max):
+    p0 = bbox[:2].astype(float)
+    p1 = p0 + bbox[2:].astype(float)
+    size = p1 - p0
+    center = p0 + (size / 2)
+    new_size = scale * size
+    p0 = center - new_size / 2
+    p1 = center + new_size / 2
+    scaled_bbox = np.array([p0, p1 - p0]).reshape(-1)
+    p0 = scaled_bbox[:2]
+    size = scaled_bbox[2:]
+    p1 = p0 + size
+    return "%d,%d,%d,%d" % (
+        int(max(p0[0], 0)),
+        int(max(p0[1], 0)),
+        int(min(p1[0], x_max)),
+        int(min(p1[1], y_max)))
+
+
+def __convert_bboxes_and_labels_to_text(bboxes, scale, max_x, max_y, labels):
+    assert (len(bboxes) == len(labels))
+    bboxes_text = ""
+    for i in range(len(bboxes)):
+        bbox = bboxes[i]
+        label = labels[i]
+        if bbox is None or label is None:
+            continue
+        bboxes_text += "%s,%s\n" % (__convert_bbox_to_text(bbox, scale, max_x, max_y), label)
+    return bboxes_text
+
+
+def __convert_rects_to_bboxes(rects):
+    bboxes = []
+    for rect in rects:
+        p0 = rect[:2]
+        p1 = rect[2:]
+        size = p1 - p0
+        bbox = np.array([p0, size]).reshape(-1)
+        bboxes.append(bbox)
+    return bboxes
+
+
+def validate_bboxes_text(s):
+    if s is None:
+        return ""
+    lines = s.split("\n")
+    for line in lines:
+        if len(line.strip()) > 0:
+            try:
+                parts = line.strip().split(',', 4)
+                if len(parts) != 5:
+                    raise ValueError("Line does not have enough parts for rect and label.")
+                rect_str = parts[:4]
+                np.array(rect_str, dtype=float).astype(int)
+            except Exception as e:
+                message = f"Error: Line '{line}' is not a valid bbox format. Details: {e}"
+                logging.critical(message)
+                raise ValueError(message)
+    return s
+
+def convert_text_to_rects_and_labels(bboxes_text):
+    rects = []
+    labels = []
+    object_ids = []
+    if not bboxes_text:
+        return rects, labels, object_ids
+
+    lines = [line for line in bboxes_text.split('\n') if line.strip()]
+    for line in lines:
+        try:
+            parts = line.strip().split(',', 5)
+            if len(parts) < 5:
+                logging.warning(f"Skipping malformed bbox line (not enough parts): '{line}'")
+                continue
+
+            rect_str = parts[:4]
+            label = parts[4]
+            object_id = parts[5] if len(parts) > 5 else None
+
+            coords = np.array(rect_str, dtype=float).astype(int)
+            x1, y1, x2, y2 = coords
+            rect = np.array([min(x1, x2), min(y1, y2), max(x1, x2), max(y1, y2)])
+
+            rects.append(rect)
+            labels.append(label)
+            object_ids.append(object_id)
+        except (ValueError, IndexError) as e:
+            logging.warning(f"Skipping malformed bbox line (parsing error): '{line}'. Error: {e}")
+            continue
+    return rects, labels, object_ids
+
+
+def count_boxes(bboxes_text):
+    if not bboxes_text:
+        return 0
+    return len([line for line in bboxes_text.split('\n') if line.strip()])
+
+
+def __convert_text_to_bboxes_and_labels(bboxes_text):
+    rects, labels, _ = convert_text_to_rects_and_labels(bboxes_text)
+    bboxes = __convert_rects_to_bboxes(rects)
+    return bboxes, labels
+
+
+def __scale_bboxes(bboxes, scale):
+    scaled_bboxes = []
+    for bbox in bboxes:
+        if bbox is None:
+            scaled_bboxes.append(None)
+        else:
+            p0 = bbox[:2].astype(float)
+            p1 = p0 + bbox[2:].astype(float)
+            size = p1 - p0
+            center = p0 + (size / 2)
+            new_size = scale * size
+            p0 = center - new_size / 2
+            p1 = center + new_size / 2
+            scaled_bboxes.append(np.array([p0, p1 - p0]).reshape(-1))
+    return scaled_bboxes
+
+
+def parse_bboxes_text(bboxes_text, scale=1):
+    bboxes_, labels = __convert_text_to_bboxes_and_labels(bboxes_text)
+    bboxes = __scale_bboxes(bboxes_, scale)
+    return bboxes, labels
+
+
+def extract_labels(bboxes_text):
+    _, labels, _ = convert_text_to_rects_and_labels(bboxes_text)
+    return labels
+
+
+def format_bboxes_text(bboxes, labels, scale, max_x, max_y):
+    return __convert_bboxes_and_labels_to_text(bboxes, 1 / scale, max_x, max_y, labels)
+
+
+def convert_to_yolo_format(bboxes_text, class_map, image_width, image_height):
+    rects, labels, _ = convert_text_to_rects_and_labels(bboxes_text)
+    yolo_lines = []
+    for i, rect in enumerate(rects):
+        label = labels[i]
+        if label not in class_map: continue
+        class_id = class_map[label]
+        x1, y1, x2, y2 = rect
+        box_width = float(x2 - x1)
+        box_height = float(y2 - y1)
+        x_center = float(x1) + (box_width / 2)
+        y_center = float(y1) + (box_height / 2)
+        x_center_norm = x_center / image_width
+        y_center_norm = y_center / image_height
+        width_norm = box_width / image_width
+        height_norm = box_height / image_height
+        yolo_lines.append(f"{class_id} {x_center_norm:.6f} {y_center_norm:.6f} {width_norm:.6f} {height_norm:.6f}")
+    return "\n".join(yolo_lines)