ultralytics 8.0.238__py3-none-any.whl → 8.0.239__py3-none-any.whl

ultralytics/data/converter.py

@@ -20,10 +20,98 @@ def coco91_to_coco80_class():
             corresponding 91-index class ID.
     """
     return [
-        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, None, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, None, 24, 25, None,
-        None, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, None, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50,
-        51, 52, 53, 54, 55, 56, 57, 58, 59, None, 60, None, None, 61, None, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72,
-        None, 73, 74, 75, 76, 77, 78, 79, None]
+        0,
+        1,
+        2,
+        3,
+        4,
+        5,
+        6,
+        7,
+        8,
+        9,
+        10,
+        None,
+        11,
+        12,
+        13,
+        14,
+        15,
+        16,
+        17,
+        18,
+        19,
+        20,
+        21,
+        22,
+        23,
+        None,
+        24,
+        25,
+        None,
+        None,
+        26,
+        27,
+        28,
+        29,
+        30,
+        31,
+        32,
+        33,
+        34,
+        35,
+        36,
+        37,
+        38,
+        39,
+        None,
+        40,
+        41,
+        42,
+        43,
+        44,
+        45,
+        46,
+        47,
+        48,
+        49,
+        50,
+        51,
+        52,
+        53,
+        54,
+        55,
+        56,
+        57,
+        58,
+        59,
+        None,
+        60,
+        None,
+        None,
+        61,
+        None,
+        62,
+        63,
+        64,
+        65,
+        66,
+        67,
+        68,
+        69,
+        70,
+        71,
+        72,
+        None,
+        73,
+        74,
+        75,
+        76,
+        77,
+        78,
+        79,
+        None,
+    ]
 
 
 def coco80_to_coco91_class():
@@ -42,16 +130,96 @@ def coco80_to_coco91_class():
         ```
     """
     return [
-        1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34,
-        35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63,
-        64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90]
-
-
-def convert_coco(labels_dir='../coco/annotations/',
-                 save_dir='coco_converted/',
-                 use_segments=False,
-                 use_keypoints=False,
-                 cls91to80=True):
+        1,
+        2,
+        3,
+        4,
+        5,
+        6,
+        7,
+        8,
+        9,
+        10,
+        11,
+        13,
+        14,
+        15,
+        16,
+        17,
+        18,
+        19,
+        20,
+        21,
+        22,
+        23,
+        24,
+        25,
+        27,
+        28,
+        31,
+        32,
+        33,
+        34,
+        35,
+        36,
+        37,
+        38,
+        39,
+        40,
+        41,
+        42,
+        43,
+        44,
+        46,
+        47,
+        48,
+        49,
+        50,
+        51,
+        52,
+        53,
+        54,
+        55,
+        56,
+        57,
+        58,
+        59,
+        60,
+        61,
+        62,
+        63,
+        64,
+        65,
+        67,
+        70,
+        72,
+        73,
+        74,
+        75,
+        76,
+        77,
+        78,
+        79,
+        80,
+        81,
+        82,
+        84,
+        85,
+        86,
+        87,
+        88,
+        89,
+        90,
+    ]
+
+
+def convert_coco(
+    labels_dir="../coco/annotations/",
+    save_dir="coco_converted/",
+    use_segments=False,
+    use_keypoints=False,
+    cls91to80=True,
+):
     """
     Converts COCO dataset annotations to a YOLO annotation format  suitable for training YOLO models.
 
@@ -75,76 +243,78 @@ def convert_coco(labels_dir='../coco/annotations/',
 
     # Create dataset directory
     save_dir = increment_path(save_dir)  # increment if save directory already exists
-    for p in save_dir / 'labels', save_dir / 'images':
+    for p in save_dir / "labels", save_dir / "images":
         p.mkdir(parents=True, exist_ok=True)  # make dir
 
     # Convert classes
     coco80 = coco91_to_coco80_class()
 
     # Import json
-    for json_file in sorted(Path(labels_dir).resolve().glob('*.json')):
-        fn = Path(save_dir) / 'labels' / json_file.stem.replace('instances_', '')  # folder name
+    for json_file in sorted(Path(labels_dir).resolve().glob("*.json")):
+        fn = Path(save_dir) / "labels" / json_file.stem.replace("instances_", "")  # folder name
         fn.mkdir(parents=True, exist_ok=True)
         with open(json_file) as f:
             data = json.load(f)
 
         # Create image dict
-        images = {f'{x["id"]:d}': x for x in data['images']}
+        images = {f'{x["id"]:d}': x for x in data["images"]}
         # Create image-annotations dict
         imgToAnns = defaultdict(list)
-        for ann in data['annotations']:
-            imgToAnns[ann['image_id']].append(ann)
+        for ann in data["annotations"]:
+            imgToAnns[ann["image_id"]].append(ann)
 
         # Write labels file
-        for img_id, anns in TQDM(imgToAnns.items(), desc=f'Annotations {json_file}'):
-            img = images[f'{img_id:d}']
-            h, w, f = img['height'], img['width'], img['file_name']
+        for img_id, anns in TQDM(imgToAnns.items(), desc=f"Annotations {json_file}"):
+            img = images[f"{img_id:d}"]
+            h, w, f = img["height"], img["width"], img["file_name"]
 
             bboxes = []
             segments = []
             keypoints = []
             for ann in anns:
-                if ann['iscrowd']:
+                if ann["iscrowd"]:
                     continue
                 # The COCO box format is [top left x, top left y, width, height]
-                box = np.array(ann['bbox'], dtype=np.float64)
+                box = np.array(ann["bbox"], dtype=np.float64)
                 box[:2] += box[2:] / 2  # xy top-left corner to center
                 box[[0, 2]] /= w  # normalize x
                 box[[1, 3]] /= h  # normalize y
                 if box[2] <= 0 or box[3] <= 0:  # if w <= 0 and h <= 0
                     continue
 
-                cls = coco80[ann['category_id'] - 1] if cls91to80 else ann['category_id'] - 1  # class
+                cls = coco80[ann["category_id"] - 1] if cls91to80 else ann["category_id"] - 1  # class
                 box = [cls] + box.tolist()
                 if box not in bboxes:
                     bboxes.append(box)
-                    if use_segments and ann.get('segmentation') is not None:
-                        if len(ann['segmentation']) == 0:
+                    if use_segments and ann.get("segmentation") is not None:
+                        if len(ann["segmentation"]) == 0:
                             segments.append([])
                             continue
-                        elif len(ann['segmentation']) > 1:
-                            s = merge_multi_segment(ann['segmentation'])
+                        elif len(ann["segmentation"]) > 1:
+                            s = merge_multi_segment(ann["segmentation"])
                             s = (np.concatenate(s, axis=0) / np.array([w, h])).reshape(-1).tolist()
                         else:
-                            s = [j for i in ann['segmentation'] for j in i]  # all segments concatenated
+                            s = [j for i in ann["segmentation"] for j in i]  # all segments concatenated
                             s = (np.array(s).reshape(-1, 2) / np.array([w, h])).reshape(-1).tolist()
                         s = [cls] + s
                         segments.append(s)
-                    if use_keypoints and ann.get('keypoints') is not None:
-                        keypoints.append(box + (np.array(ann['keypoints']).reshape(-1, 3) /
-                                                np.array([w, h, 1])).reshape(-1).tolist())
+                    if use_keypoints and ann.get("keypoints") is not None:
+                        keypoints.append(
+                            box + (np.array(ann["keypoints"]).reshape(-1, 3) / np.array([w, h, 1])).reshape(-1).tolist()
+                        )
 
             # Write
-            with open((fn / f).with_suffix('.txt'), 'a') as file:
+            with open((fn / f).with_suffix(".txt"), "a") as file:
                 for i in range(len(bboxes)):
                     if use_keypoints:
-                        line = *(keypoints[i]),  # cls, box, keypoints
+                        line = (*(keypoints[i]),)  # cls, box, keypoints
                     else:
-                        line = *(segments[i]
-                                 if use_segments and len(segments[i]) > 0 else bboxes[i]),  # cls, box or segments
-                    file.write(('%g ' * len(line)).rstrip() % line + '\n')
+                        line = (
+                            *(segments[i] if use_segments and len(segments[i]) > 0 else bboxes[i]),
+                        )  # cls, box or segments
+                    file.write(("%g " * len(line)).rstrip() % line + "\n")
 
-    LOGGER.info(f'COCO data converted successfully.\nResults saved to {save_dir.resolve()}')
+    LOGGER.info(f"COCO data converted successfully.\nResults saved to {save_dir.resolve()}")
 
 
 def convert_dota_to_yolo_obb(dota_root_path: str):
@@ -184,31 +354,32 @@ def convert_dota_to_yolo_obb(dota_root_path: str):
 
     # Class names to indices mapping
     class_mapping = {
-        'plane': 0,
-        'ship': 1,
-        'storage-tank': 2,
-        'baseball-diamond': 3,
-        'tennis-court': 4,
-        'basketball-court': 5,
-        'ground-track-field': 6,
-        'harbor': 7,
-        'bridge': 8,
-        'large-vehicle': 9,
-        'small-vehicle': 10,
-        'helicopter': 11,
-        'roundabout': 12,
-        'soccer-ball-field': 13,
-        'swimming-pool': 14,
-        'container-crane': 15,
-        'airport': 16,
-        'helipad': 17}
+        "plane": 0,
+        "ship": 1,
+        "storage-tank": 2,
+        "baseball-diamond": 3,
+        "tennis-court": 4,
+        "basketball-court": 5,
+        "ground-track-field": 6,
+        "harbor": 7,
+        "bridge": 8,
+        "large-vehicle": 9,
+        "small-vehicle": 10,
+        "helicopter": 11,
+        "roundabout": 12,
+        "soccer-ball-field": 13,
+        "swimming-pool": 14,
+        "container-crane": 15,
+        "airport": 16,
+        "helipad": 17,
+    }
 
     def convert_label(image_name, image_width, image_height, orig_label_dir, save_dir):
         """Converts a single image's DOTA annotation to YOLO OBB format and saves it to a specified directory."""
-        orig_label_path = orig_label_dir / f'{image_name}.txt'
-        save_path = save_dir / f'{image_name}.txt'
+        orig_label_path = orig_label_dir / f"{image_name}.txt"
+        save_path = save_dir / f"{image_name}.txt"
 
-        with orig_label_path.open('r') as f, save_path.open('w') as g:
+        with orig_label_path.open("r") as f, save_path.open("w") as g:
             lines = f.readlines()
             for line in lines:
                 parts = line.strip().split()
@@ -218,20 +389,21 @@ def convert_dota_to_yolo_obb(dota_root_path: str):
                 class_idx = class_mapping[class_name]
                 coords = [float(p) for p in parts[:8]]
                 normalized_coords = [
-                    coords[i] / image_width if i % 2 == 0 else coords[i] / image_height for i in range(8)]
-                formatted_coords = ['{:.6g}'.format(coord) for coord in normalized_coords]
+                    coords[i] / image_width if i % 2 == 0 else coords[i] / image_height for i in range(8)
+                ]
+                formatted_coords = ["{:.6g}".format(coord) for coord in normalized_coords]
                 g.write(f"{class_idx} {' '.join(formatted_coords)}\n")
 
-    for phase in ['train', 'val']:
-        image_dir = dota_root_path / 'images' / phase
-        orig_label_dir = dota_root_path / 'labels' / f'{phase}_original'
-        save_dir = dota_root_path / 'labels' / phase
+    for phase in ["train", "val"]:
+        image_dir = dota_root_path / "images" / phase
+        orig_label_dir = dota_root_path / "labels" / f"{phase}_original"
+        save_dir = dota_root_path / "labels" / phase
 
         save_dir.mkdir(parents=True, exist_ok=True)
 
         image_paths = list(image_dir.iterdir())
-        for image_path in TQDM(image_paths, desc=f'Processing {phase} images'):
-            if image_path.suffix != '.png':
+        for image_path in TQDM(image_paths, desc=f"Processing {phase} images"):
+            if image_path.suffix != ".png":
                 continue
             image_name_without_ext = image_path.stem
             img = cv2.imread(str(image_path))
@@ -293,7 +465,7 @@ def merge_multi_segment(segments):
                     s.append(segments[i])
                 else:
                     idx = [0, idx[1] - idx[0]]
-                    s.append(segments[i][idx[0]:idx[1] + 1])
+                    s.append(segments[i][idx[0] : idx[1] + 1])
 
         else:
             for i in range(len(idx_list) - 1, -1, -1):