You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
303 lines
12 KiB
303 lines
12 KiB
# Ultralytics YOLO 🚀, AGPL-3.0 license
|
|
|
|
import json
|
|
from collections import defaultdict
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from tqdm import tqdm
|
|
|
|
from ultralytics.utils.checks import check_requirements
|
|
from ultralytics.utils.files import make_dirs
|
|
|
|
|
|
def coco91_to_coco80_class():
|
|
"""Converts 91-index COCO class IDs to 80-index COCO class IDs.
|
|
|
|
Returns:
|
|
(list): A list of 91 class IDs where the index represents the 80-index class ID and the value is the
|
|
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]
|
|
|
|
|
|
def convert_coco(labels_dir='../coco/annotations/', use_segments=False, use_keypoints=False, cls91to80=True):
|
|
"""Converts COCO dataset annotations to a format suitable for training YOLOv5 models.
|
|
|
|
Args:
|
|
labels_dir (str, optional): Path to directory containing COCO dataset annotation files.
|
|
use_segments (bool, optional): Whether to include segmentation masks in the output.
|
|
use_keypoints (bool, optional): Whether to include keypoint annotations in the output.
|
|
cls91to80 (bool, optional): Whether to map 91 COCO class IDs to the corresponding 80 COCO class IDs.
|
|
|
|
Raises:
|
|
FileNotFoundError: If the labels_dir path does not exist.
|
|
|
|
Example Usage:
|
|
convert_coco(labels_dir='../coco/annotations/', use_segments=True, use_keypoints=True, cls91to80=True)
|
|
|
|
Output:
|
|
Generates output files in the specified output directory.
|
|
"""
|
|
|
|
save_dir = make_dirs('yolo_labels') # output directory
|
|
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
|
|
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']}
|
|
# Create image-annotations dict
|
|
imgToAnns = defaultdict(list)
|
|
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']
|
|
|
|
bboxes = []
|
|
segments = []
|
|
keypoints = []
|
|
for ann in anns:
|
|
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[: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
|
|
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:
|
|
segments.append([])
|
|
continue
|
|
if isinstance(ann['segmentation'], dict):
|
|
ann['segmentation'] = rle2polygon(ann['segmentation'])
|
|
if 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 = (np.array(s).reshape(-1, 2) / np.array([w, h])).reshape(-1).tolist()
|
|
s = [cls] + s
|
|
if s not in segments:
|
|
segments.append(s)
|
|
if use_keypoints and ann.get('keypoints') is not None:
|
|
k = (np.array(ann['keypoints']).reshape(-1, 3) / np.array([w, h, 1])).reshape(-1).tolist()
|
|
k = box + k
|
|
keypoints.append(k)
|
|
|
|
# Write
|
|
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
|
|
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')
|
|
|
|
|
|
def convert_dota_to_yolo_obb(dota_root_path: str):
|
|
"""
|
|
Converts DOTA dataset annotations to YOLO OBB (Oriented Bounding Box) format.
|
|
|
|
The function processes images in the 'train' and 'val' folders of the DOTA dataset. For each image, it reads the
|
|
associated label from the original labels directory and writes new labels in YOLO OBB format to a new directory.
|
|
|
|
Args:
|
|
dota_root_path (str): The root directory path of the DOTA dataset.
|
|
|
|
Example:
|
|
```python
|
|
from ultralytics.data.converter import convert_dota_to_yolo_obb
|
|
|
|
convert_dota_to_yolo_obb('path/to/DOTA')
|
|
```
|
|
|
|
Notes:
|
|
The directory structure assumed for the DOTA dataset:
|
|
- DOTA
|
|
- images
|
|
- train
|
|
- val
|
|
- labels
|
|
- train_original
|
|
- val_original
|
|
|
|
After the function execution, the new labels will be saved in:
|
|
- DOTA
|
|
- labels
|
|
- train
|
|
- val
|
|
"""
|
|
dota_root_path = Path(dota_root_path)
|
|
|
|
# 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}
|
|
|
|
def convert_label(image_name, image_width, image_height, orig_label_dir, save_dir):
|
|
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:
|
|
lines = f.readlines()
|
|
for line in lines:
|
|
parts = line.strip().split()
|
|
if len(parts) < 9:
|
|
continue
|
|
class_name = parts[8]
|
|
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]
|
|
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
|
|
|
|
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':
|
|
continue
|
|
image_name_without_ext = image_path.stem
|
|
img = cv2.imread(str(image_path))
|
|
h, w = img.shape[:2]
|
|
convert_label(image_name_without_ext, w, h, orig_label_dir, save_dir)
|
|
|
|
|
|
def rle2polygon(segmentation):
|
|
"""
|
|
Convert Run-Length Encoding (RLE) mask to polygon coordinates.
|
|
|
|
Args:
|
|
segmentation (dict, list): RLE mask representation of the object segmentation.
|
|
|
|
Returns:
|
|
(list): A list of lists representing the polygon coordinates for each contour.
|
|
|
|
Note:
|
|
Requires the 'pycocotools' package to be installed.
|
|
"""
|
|
check_requirements('pycocotools')
|
|
from pycocotools import mask
|
|
|
|
m = mask.decode(segmentation)
|
|
m[m > 0] = 255
|
|
contours, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_KCOS)
|
|
polygons = []
|
|
for contour in contours:
|
|
epsilon = 0.001 * cv2.arcLength(contour, True)
|
|
contour_approx = cv2.approxPolyDP(contour, epsilon, True)
|
|
polygon = contour_approx.flatten().tolist()
|
|
polygons.append(polygon)
|
|
return polygons
|
|
|
|
|
|
def min_index(arr1, arr2):
|
|
"""
|
|
Find a pair of indexes with the shortest distance between two arrays of 2D points.
|
|
|
|
Args:
|
|
arr1 (np.array): A NumPy array of shape (N, 2) representing N 2D points.
|
|
arr2 (np.array): A NumPy array of shape (M, 2) representing M 2D points.
|
|
|
|
Returns:
|
|
(tuple): A tuple containing the indexes of the points with the shortest distance in arr1 and arr2 respectively.
|
|
"""
|
|
dis = ((arr1[:, None, :] - arr2[None, :, :]) ** 2).sum(-1)
|
|
return np.unravel_index(np.argmin(dis, axis=None), dis.shape)
|
|
|
|
|
|
def merge_multi_segment(segments):
|
|
"""
|
|
Merge multiple segments into one list by connecting the coordinates with the minimum distance between each segment.
|
|
This function connects these coordinates with a thin line to merge all segments into one.
|
|
|
|
Args:
|
|
segments (List[List]): Original segmentations in COCO's JSON file.
|
|
Each element is a list of coordinates, like [segmentation1, segmentation2,...].
|
|
|
|
Returns:
|
|
s (List[np.ndarray]): A list of connected segments represented as NumPy arrays.
|
|
"""
|
|
s = []
|
|
segments = [np.array(i).reshape(-1, 2) for i in segments]
|
|
idx_list = [[] for _ in range(len(segments))]
|
|
|
|
# record the indexes with min distance between each segment
|
|
for i in range(1, len(segments)):
|
|
idx1, idx2 = min_index(segments[i - 1], segments[i])
|
|
idx_list[i - 1].append(idx1)
|
|
idx_list[i].append(idx2)
|
|
|
|
# use two round to connect all the segments
|
|
for k in range(2):
|
|
# forward connection
|
|
if k == 0:
|
|
for i, idx in enumerate(idx_list):
|
|
# middle segments have two indexes
|
|
# reverse the index of middle segments
|
|
if len(idx) == 2 and idx[0] > idx[1]:
|
|
idx = idx[::-1]
|
|
segments[i] = segments[i][::-1, :]
|
|
|
|
segments[i] = np.roll(segments[i], -idx[0], axis=0)
|
|
segments[i] = np.concatenate([segments[i], segments[i][:1]])
|
|
# deal with the first segment and the last one
|
|
if i in [0, len(idx_list) - 1]:
|
|
s.append(segments[i])
|
|
else:
|
|
idx = [0, idx[1] - idx[0]]
|
|
s.append(segments[i][idx[0]:idx[1] + 1])
|
|
|
|
else:
|
|
for i in range(len(idx_list) - 1, -1, -1):
|
|
if i not in [0, len(idx_list) - 1]:
|
|
idx = idx_list[i]
|
|
nidx = abs(idx[1] - idx[0])
|
|
s.append(segments[i][nidx:])
|
|
return s
|