"""
coco_to_yolo.py
Converts a COCO-formatted dataset to a YOLO-formatted dataset, flattening
the dataset (to a single folder) in the process.
If the input and output folders are the same, writes .txt files to the input folder,
and neither moves nor modifies images.
Currently ignores segmentation masks, and errors if an annotation has a
segmentation polygon but no bbox.
Has only been tested on a handful of COCO Camera Traps data sets; if you
use it for more general COCO conversion, YMMV.
"""
#%% Imports and constants
import json
import os
import shutil
import sys
import argparse
from collections import defaultdict
from tqdm import tqdm
from megadetector.utils.path_utils import safe_create_link,find_images
#%% Support functions
[docs]
def write_yolo_dataset_file(yolo_dataset_file,
dataset_base_dir,
class_list,
train_folder_relative=None,
val_folder_relative=None,
test_folder_relative=None):
"""
Write a YOLOv5 dataset.yaml file to the absolute path [yolo_dataset_file] (should
have a .yaml extension, though it's only a warning if it doesn't).
Args:
yolo_dataset_file (str): the file, typically ending in .yaml or .yml, to write.
Does not have to be within dataset_base_dir.
dataset_base_dir (str): the absolute base path of the YOLO dataset
class_list (list or str): an ordered list of class names (the first item will be class 0,
etc.), or the name of a text file containing an ordered list of class names (one per
line, starting from class zero).
train_folder_relative (str, optional): train folder name, used only to
populate dataset.yaml. Can also be a filename (e.g. a .txt file with image
files).
val_folder_relative (str, optional): val folder name, used only to
populate dataset.yaml. Can also be a filename (e.g. a .txt file with image
files).
test_folder_relative (str, optional): test folder name, used only to
populate dataset.yaml. Can also be a filename (e.g. a .txt file with image
files).
"""
# Read class names
if isinstance(class_list,str):
with open(class_list,'r') as f:
class_lines = f.readlines()
class_lines = [s.strip() for s in class_lines]
class_list = [s for s in class_lines if len(s) > 0]
if not (yolo_dataset_file.endswith('.yml') or yolo_dataset_file.endswith('.yaml')):
print('Warning: writing dataset file to a non-yml/yaml extension:\n{}'.format(
yolo_dataset_file))
# Write dataset.yaml
with open(yolo_dataset_file,'w') as f:
f.write('# Train/val sets\n')
f.write('path: {}\n'.format(dataset_base_dir))
if train_folder_relative is not None:
f.write('train: {}\n'.format(train_folder_relative))
if val_folder_relative is not None:
f.write('val: {}\n'.format(val_folder_relative))
if test_folder_relative is not None:
f.write('test: {}\n'.format(test_folder_relative))
f.write('\n')
f.write('# Classes\n')
f.write('names:\n')
for i_class,class_name in enumerate(class_list):
f.write(' {}: {}\n'.format(i_class,class_name))
# ...def write_yolo_dataset_file(...)
[docs]
def coco_to_yolo(input_image_folder,
output_folder,
input_file,
source_format='coco',
overwrite_images=False,
create_image_and_label_folders=False,
class_file_name='classes.txt',
allow_empty_annotations=False,
clip_boxes=False,
image_id_to_output_image_json_file=None,
images_to_exclude=None,
path_replacement_char='#',
category_names_to_exclude=None,
category_names_to_include=None,
write_output=True,
flatten_paths=False,
empty_image_handling='write_empty'):
"""
Converts a COCO-formatted dataset to a YOLO-formatted dataset, optionally flattening the
dataset to a single folder in the process.
If the input and output folders are the same, writes .txt files to the input folder,
and neither moves nor modifies images.
Currently ignores segmentation masks, and errors if an annotation has a
segmentation polygon but no bbox.
Args:
input_image_folder (str): the folder where images live; filenames in the COCO .json
file [input_file] should be relative to this folder
output_folder (str): the base folder for the YOLO dataset
input_file (str): a .json file in COCO format; can be the same as [input_image_folder], in which case
images are left alone.
source_format (str, optional): can be 'coco' (default) or 'coco_camera_traps'. The only difference
is that when source_format is 'coco_camera_traps', we treat an image with a non-bbox
annotation as a special case, i.e. that's how an empty image is indicated. The original
COCO standard is a little ambiguous on this issue. If source_format is 'coco', we
either treat images as empty or error, depending on the value of [allow_empty_annotations].
[allow_empty_annotations] has no effect if source_format is 'coco_camera_traps'.
overwrite_images (bool, optional): over-write images in the output folder if they exist
create_image_and_label_folders (bool, optional): whether to create separate folders called 'images' and
'labels' in the YOLO output folder. If create_image_and_label_folders is False,
a/b/c/image001.jpg will become a#b#c#image001.jpg, and the corresponding text file will
be a#b#c#image001.txt. If create_image_and_label_folders is True, a/b/c/image001.jpg will become
images/a#b#c#image001.jpg, and the corresponding text file will be
labels/a#b#c#image001.txt.
class_file_name (str, optional): .txt file (relative to the output folder) that we should
populate with a list of classes (or None to omit)
allow_empty_annotations (bool, optional): if this is False and [source_format] is 'coco',
we'll error on annotations that have no 'bbox' field
clip_boxes (bool, optional): whether to clip bounding box coordinates to the range [0,1] before
converting to YOLO xywh format
image_id_to_output_image_json_file (str, optional): an optional *output* file, to which we will write
a mapping from image IDs to output file names
images_to_exclude (list, optional): a list of image files (relative paths in the input folder) that we
should ignore
path_replacement_char (str, optional): only relevant if [flatten_paths] is True; this is used to replace
path separators, e.g. if [path_replacement_char] is '#' and [flatten_paths] is True, a/b/c/d.jpg
becomes a#b#c#d.jpg
category_names_to_exclude (str, optional): category names that should not be represented in the
YOLO output; only impacts annotations, does not prevent copying images. There's almost no reason
you would want to specify this and [category_names_to_include].
category_names_to_include (str, optional): allow-list of category names that should be represented
in the YOLO output; only impacts annotations, does not prevent copying images. There's almost
no reason you would want to specify this and [category_names_to_exclude].
write_output (bool, optional): determines whether we actually copy images and write annotations;
setting this to False mostly puts this function in "dry run" "mode. The class list
file is written regardless of the value of write_output.
flatten_paths (bool, optional): replace /'s in image filenames with [path_replacement_char],
which ensures that the output folder is a single flat folder.
empty_image_handling (str, optional): whether to omit .txt files for images with no
annotations ('omit') or write empty .txt files ('write_empty'). Both are generally considered
valid YOLO.
Returns:
dict: information about the coco --> yolo mapping, containing at least the fields:
- class_list_filename: the filename to which we wrote the flat list of class names required
by the YOLO format.
- source_image_to_dest_image: a dict mapping source images to destination images
- coco_id_to_yolo_id: a dict mapping COCO category IDs to YOLO category IDs
"""
## Validate input
assert empty_image_handling in ('omit','write_empty'), \
'Unrecognized value for empty_image_handling: {}'.format(empty_image_handling)
if category_names_to_include is not None and category_names_to_exclude is not None:
raise ValueError('category_names_to_include and category_names_to_exclude are mutually exclusive')
if output_folder is None:
output_folder = input_image_folder
if images_to_exclude is not None:
images_to_exclude = set(images_to_exclude)
if category_names_to_exclude is None:
category_names_to_exclude = {}
assert os.path.isdir(input_image_folder)
assert os.path.isfile(input_file)
os.makedirs(output_folder,exist_ok=True)
if (output_folder == input_image_folder) and (overwrite_images) and \
(not create_image_and_label_folders) and (not flatten_paths):
print('Warning: output folder and input folder are the same, disabling overwrite_images')
overwrite_images = False
## Read input data
with open(input_file,'r') as f:
data = json.load(f)
## Parse annotations
image_id_to_annotations = defaultdict(list)
# i_ann = 0; ann = data['annotations'][0]
for i_ann,ann in enumerate(data['annotations']):
# Make sure no annotations have *only* segmentation data
if ( \
('segmentation' in ann.keys()) and \
(ann['segmentation'] is not None) and \
(len(ann['segmentation']) > 0) ) \
and \
(('bbox' not in ann.keys()) or (ann['bbox'] is None) or (len(ann['bbox'])==0)):
raise ValueError('Oops: segmentation data present without bbox information, ' + \
'this script isn\'t ready for this dataset')
image_id_to_annotations[ann['image_id']].append(ann)
print('Parsed annotations for {} images'.format(len(image_id_to_annotations)))
# Re-map class IDs to make sure they run from 0...n-classes-1
#
# Note: this allows unused categories in the output data set. This is OK for
# some training pipelines, not for others.
next_category_id = 0
coco_id_to_yolo_id = {}
coco_id_to_name = {}
yolo_id_to_name = {}
coco_category_ids_to_exclude = set()
for category in data['categories']:
coco_id_to_name[category['id']] = category['name']
if (category_names_to_include is not None) and \
(category['name'] not in category_names_to_include):
coco_category_ids_to_exclude.add(category['id'])
continue
elif (category['name'] in category_names_to_exclude):
coco_category_ids_to_exclude.add(category['id'])
continue
assert category['id'] not in coco_id_to_yolo_id
coco_id_to_yolo_id[category['id']] = next_category_id
yolo_id_to_name[next_category_id] = category['name']
next_category_id += 1
## Process images (everything but I/O)
# List of dictionaries with keys 'source_image','dest_image','bboxes','dest_txt'
images_to_copy = []
missing_images = []
excluded_images = []
image_names = set()
typical_image_extensions = set(['.jpg','.jpeg','.png','.gif','.tif','.bmp'])
printed_empty_annotation_warning = False
image_id_to_output_image_name = {}
print('Processing annotations')
n_clipped_boxes = 0
n_total_boxes = 0
# i_image = 0; im = data['images'][i_image]
for i_image,im in tqdm(enumerate(data['images']),total=len(data['images'])):
output_info = {}
source_image = os.path.join(input_image_folder,im['file_name'])
output_info['source_image'] = source_image
if images_to_exclude is not None and im['file_name'] in images_to_exclude:
excluded_images.append(im['file_name'])
continue
tokens = os.path.splitext(im['file_name'])
if tokens[1].lower() not in typical_image_extensions:
print('Warning: unusual image file name {}'.format(im['file_name']))
if flatten_paths:
image_name = tokens[0].replace('\\','/').replace('/',path_replacement_char) + \
'_' + str(i_image).zfill(6)
else:
image_name = tokens[0]
assert image_name not in image_names, 'Image name collision for {}'.format(image_name)
image_names.add(image_name)
assert im['id'] not in image_id_to_output_image_name
image_id_to_output_image_name[im['id']] = image_name
dest_image_relative = image_name + tokens[1]
output_info['dest_image_relative'] = dest_image_relative
dest_txt_relative = image_name + '.txt'
output_info['dest_txt_relative'] = dest_txt_relative
output_info['bboxes'] = []
# assert os.path.isfile(source_image), 'Could not find image {}'.format(source_image)
if not os.path.isfile(source_image):
print('Warning: could not find image {}'.format(source_image))
missing_images.append(im['file_name'])
continue
image_id = im['id']
image_bboxes = []
if image_id in image_id_to_annotations:
for ann in image_id_to_annotations[image_id]:
# If this annotation has no bounding boxes...
if 'bbox' not in ann or ann['bbox'] is None or len(ann['bbox']) == 0:
if source_format == 'coco':
if not allow_empty_annotations:
# This is not entirely clear from the COCO spec, but it seems to be consensus
# that if you want to specify an image with no objects, you don't include any
# annotations for that image.
raise ValueError('If an annotation exists, it should have content')
else:
continue
elif source_format == 'coco_camera_traps':
# We allow empty bbox lists in COCO camera traps files; this is typically a
# negative example in a dataset that has bounding boxes, and 0 is typically
# the empty category, which is typically 0.
if ann['category_id'] != 0:
if not printed_empty_annotation_warning:
printed_empty_annotation_warning = True
print('Warning: non-bbox annotation found with category {}'.format(
ann['category_id']))
continue
else:
raise ValueError('Unrecognized COCO variant: {}'.format(source_format))
# ...if this is an empty annotation
coco_bbox = ann['bbox']
# This category isn't in our category list. This typically corresponds to whole sets
# of images that were excluded from the YOLO set.
if ann['category_id'] in coco_category_ids_to_exclude:
continue
yolo_category_id = coco_id_to_yolo_id[ann['category_id']]
# COCO: [x_min, y_min, width, height] in absolute coordinates
# YOLO: [class, x_center, y_center, width, height] in normalized coordinates
# Convert from COCO coordinates to YOLO coordinates
img_w = im['width']
img_h = im['height']
if source_format in ('coco','coco_camera_traps'):
x_min_absolute = coco_bbox[0]
y_min_absolute = coco_bbox[1]
box_w_absolute = coco_bbox[2]
box_h_absolute = coco_bbox[3]
x_center_absolute = (x_min_absolute + (x_min_absolute + box_w_absolute)) / 2
y_center_absolute = (y_min_absolute + (y_min_absolute + box_h_absolute)) / 2
x_center_relative = x_center_absolute / img_w
y_center_relative = y_center_absolute / img_h
box_w_relative = box_w_absolute / img_w
box_h_relative = box_h_absolute / img_h
else:
raise ValueError('Unrecognized source format {}'.format(source_format))
if clip_boxes:
clipped_box = False
box_right = x_center_relative + (box_w_relative / 2.0)
if box_right > 1.0:
clipped_box = True
overhang = box_right - 1.0
box_w_relative -= overhang
x_center_relative -= (overhang / 2.0)
box_bottom = y_center_relative + (box_h_relative / 2.0)
if box_bottom > 1.0:
clipped_box = True
overhang = box_bottom - 1.0
box_h_relative -= overhang
y_center_relative -= (overhang / 2.0)
box_left = x_center_relative - (box_w_relative / 2.0)
if box_left < 0.0:
clipped_box = True
overhang = abs(box_left)
box_w_relative -= overhang
x_center_relative += (overhang / 2.0)
box_top = y_center_relative - (box_h_relative / 2.0)
if box_top < 0.0:
clipped_box = True
overhang = abs(box_top)
box_h_relative -= overhang
y_center_relative += (overhang / 2.0)
if clipped_box:
n_clipped_boxes += 1
yolo_box = [yolo_category_id,
x_center_relative, y_center_relative,
box_w_relative, box_h_relative]
image_bboxes.append(yolo_box)
n_total_boxes += 1
# ...for each annotation
# ...if this image has annotations
output_info['bboxes'] = image_bboxes
images_to_copy.append(output_info)
# ...for each image
print('\nWriting {} boxes ({} clipped) for {} images'.format(n_total_boxes,
n_clipped_boxes,len(images_to_copy)))
print('{} missing images (of {})'.format(len(missing_images),len(data['images'])))
if images_to_exclude is not None:
print('{} excluded images (of {})'.format(len(excluded_images),len(data['images'])))
## Write output
print('Generating class list')
if class_file_name is not None:
class_list_filename = os.path.join(output_folder,class_file_name)
with open(class_list_filename, 'w') as f:
print('Writing class list to {}'.format(class_list_filename))
for i_class in range(0,len(yolo_id_to_name)):
# Category IDs should range from 0..N-1
assert i_class in yolo_id_to_name
f.write(yolo_id_to_name[i_class] + '\n')
else:
class_list_filename = None
if image_id_to_output_image_json_file is not None:
print('Writing image ID mapping to {}'.format(image_id_to_output_image_json_file))
with open(image_id_to_output_image_json_file,'w') as f:
json.dump(image_id_to_output_image_name,f,indent=1)
if (output_folder == input_image_folder) and (not create_image_and_label_folders):
print('Creating annotation files (not copying images, input and output folder are the same)')
else:
print('Copying images and creating annotation files')
if create_image_and_label_folders:
dest_image_folder = os.path.join(output_folder,'images')
dest_txt_folder = os.path.join(output_folder,'labels')
else:
dest_image_folder = output_folder
dest_txt_folder = output_folder
source_image_to_dest_image = {}
label_files_written = []
n_boxes_written = 0
# TODO: parallelize this loop
#
# output_info = images_to_copy[0]
for output_info in tqdm(images_to_copy):
source_image = output_info['source_image']
dest_image_relative = output_info['dest_image_relative']
dest_txt_relative = output_info['dest_txt_relative']
dest_image = os.path.join(dest_image_folder,dest_image_relative)
dest_txt = os.path.join(dest_txt_folder,dest_txt_relative)
source_image_to_dest_image[source_image] = dest_image
# Copy the image if necessary
if write_output:
os.makedirs(os.path.dirname(dest_image),exist_ok=True)
os.makedirs(os.path.dirname(dest_txt),exist_ok=True)
if not create_image_and_label_folders:
assert os.path.dirname(dest_image) == os.path.dirname(dest_txt)
if (not os.path.isfile(dest_image)) or (overwrite_images):
shutil.copyfile(source_image,dest_image)
bboxes = output_info['bboxes']
# Write the annotation file if necessary
if (len(bboxes) > 0) or (empty_image_handling == 'write_empty'):
n_boxes_written += len(bboxes)
label_files_written.append(dest_txt)
if write_output:
with open(dest_txt,'w') as f:
# bbox = bboxes[0]
for bbox in bboxes:
assert len(bbox) == 5
s = '{} {} {} {} {}'.format(bbox[0],bbox[1],bbox[2],bbox[3],bbox[4])
f.write(s + '\n')
# ...if there are boxes for this image
# ...for each image
coco_to_yolo_info = {}
coco_to_yolo_info['class_list_filename'] = class_list_filename
coco_to_yolo_info['source_image_to_dest_image'] = source_image_to_dest_image
coco_to_yolo_info['coco_id_to_yolo_id'] = coco_id_to_yolo_id
coco_to_yolo_info['label_files_written'] = label_files_written
coco_to_yolo_info['n_boxes_written'] = n_boxes_written
return coco_to_yolo_info
# ...def coco_to_yolo(...)
def create_yolo_symlinks(source_folder,
images_folder,
labels_folder,
class_list_file=None,
class_list_output_name='object.data',
force_lowercase_image_extension=False):
"""
Given a YOLO-formatted folder of images and .txt files, creates a folder
of symlinks to all the images, and a folder of symlinks to all the labels.
Used to support preview/editing tools that assume images and labels are in separate
folders.
Args:
source_folder (str): input folder
images_folder (str): output folder with links to images
labels_folder (str): output folder with links to labels
class_list_file (str, optional): list to classes.txt file
class_list_output_name (str, optional): output file to write with class information
force_lowercase_image_extension (bool, False): create symlinks with, e.g., .jpg, even
if the input image is, e.g., .JPG
:meta private:
"""
assert source_folder != images_folder and source_folder != labels_folder
os.makedirs(images_folder,exist_ok=True)
os.makedirs(labels_folder,exist_ok=True)
image_files_relative = find_images(source_folder,recursive=True,return_relative_paths=True)
# image_fn_relative = image_files_relative[0]=
for image_fn_relative in tqdm(image_files_relative):
source_file_abs = os.path.join(source_folder,image_fn_relative)
target_file_abs = os.path.join(images_folder,image_fn_relative)
if force_lowercase_image_extension:
tokens = os.path.splitext(target_file_abs)
target_file_abs = tokens[0] + tokens[1].lower()
os.makedirs(os.path.dirname(target_file_abs),exist_ok=True)
safe_create_link(source_file_abs,target_file_abs)
source_annotation_file_abs = os.path.splitext(source_file_abs)[0] + '.txt'
if os.path.isfile(source_annotation_file_abs):
target_annotation_file_abs = \
os.path.splitext(os.path.join(labels_folder,image_fn_relative))[0] + '.txt'
os.makedirs(os.path.dirname(target_annotation_file_abs),exist_ok=True)
safe_create_link(source_annotation_file_abs,target_annotation_file_abs)
# ...for each image
if class_list_file is not None:
target_class_list_file = os.path.join(labels_folder,class_list_output_name)
safe_create_link(class_list_file,target_class_list_file)
# ...def create_yolo_symlinks(...)
#%% Interactive driver
if False:
pass
#%% Options
input_file = os.path.expanduser('~/data/md-test-coco.json')
image_folder = os.path.expanduser('~/data/md-test')
output_folder = os.path.expanduser('~/data/md-test-yolo')
create_image_and_label_folders=False
class_file_name='classes.txt'
allow_empty_annotations=False
clip_boxes=False
image_id_to_output_image_json_file=None
images_to_exclude=None
path_replacement_char='#'
category_names_to_exclude=None
#%% Programmatic execution
coco_to_yolo_results = coco_to_yolo(image_folder,output_folder,input_file,
source_format='coco',
overwrite_images=False,
create_image_and_label_folders=create_image_and_label_folders,
class_file_name=class_file_name,
allow_empty_annotations=allow_empty_annotations,
clip_boxes=clip_boxes)
create_yolo_symlinks(source_folder=output_folder,
images_folder=output_folder + '/images',
labels_folder=output_folder + '/labels',
class_list_file=coco_to_yolo_results['class_list_filename'],
class_list_output_name='object.data',
force_lowercase_image_extension=True)
#%% Prepare command-line example
s = 'python coco_to_yolo.py {} {} {} --create_bounding_box_editor_symlinks'.format(
image_folder,output_folder,input_file)
print(s)
import clipboard; clipboard.copy(s)
#%% Command-line driver
def main(): # noqa
parser = argparse.ArgumentParser(
description='Convert COCO-formatted data to YOLO format, flattening the image structure')
# input_image_folder,output_folder,input_file
parser.add_argument(
'input_folder',
type=str,
help='Path to input images')
parser.add_argument(
'output_folder',
type=str,
help='Path to flat, YOLO-formatted dataset')
parser.add_argument(
'input_file',
type=str,
help='Path to COCO dataset file (.json)')
parser.add_argument(
'--create_bounding_box_editor_symlinks',
action='store_true',
help='Prepare symlinks so the whole folder appears to contain "images" and "labels" folderss')
if len(sys.argv[1:]) == 0:
parser.print_help()
parser.exit()
args = parser.parse_args()
coco_to_yolo_results = coco_to_yolo(args.input_folder,args.output_folder,args.input_file)
if args.create_bounding_box_editor_symlinks:
create_yolo_symlinks(source_folder=args.output_folder,
images_folder=args.output_folder + '/images',
labels_folder=args.output_folder + '/labels',
class_list_file=coco_to_yolo_results['class_list_filename'],
class_list_output_name='object.data',
force_lowercase_image_extension=True)
if __name__ == '__main__':
main()