def csv_from_pascal(
file_path_csv: str,
images_directory: str,
pascal_directory: str,
):
bbox_file_ids = []
with open(file_path_csv, "w") as csv_file:
csv_file.write(
"ImageID,Source,LabelName,Confidence,XMin,XMax,YMin,YMax,"
"IsOccluded,IsTruncated,IsGroupOf,IsDepiction,IsInside,id,ClassName\n",
)
matching_file_ids = matching_ids(pascal_directory, images_directory, ".xml", ".jpg")
for file_id in matching_file_ids:
bboxes = bounding_boxes_pascal(
os.path.join(pascal_directory, file_id + ".xml"),
)
if len(bboxes):
bbox_file_ids.append(file_id)
for bbox in bboxes:
csv_file.write(
f"{file_id},,,,{bbox['xmin']},{bbox['xmax']},"
f"{bbox['ymin']},{bbox['ymax']},,,,,,,{bbox['label']}\n",
)
# return list of image file IDs that are included in the CSV
return set(bbox_file_ids)
def purge_non_matching(
images_dir: str,
annotations_dir: str,
annotation_format: str,
problems_dir: str = None,
) -> Set[str]:
"""
TODO
:param images_dir:
:param annotations_dir:
:param annotation_format:
:param problems_dir:
:return:
"""
# determine the file extensions
if annotation_format not in ["darknet", "kitti", "pascal"]:
raise ValueError(f"Unsupported annotation format: {annotation_format}")
else:
annotation_ext = FORMAT_EXTENSIONS[annotation_format]
image_ext = ".jpg"
# make the problem files directory if necessary, in case it doesn't already exist
if problems_dir is not None:
os.makedirs(problems_dir, exist_ok=True)
# remove files that aren't matches
matching_file_ids = matching_ids(annotations_dir, images_dir,
annotation_ext, image_ext)
for directory in [annotations_dir, images_dir]:
for file_name in os.listdir(directory):
# only filter out image and Darknet annotation files (this is
# needed in case a subdirectory exists in the directory)
# and skip the file named "labels.txt"
if file_name != "labels.txt" and \
(file_name.endswith(annotation_ext) or file_name.endswith(image_ext)):
if os.path.splitext(file_name)[0] not in matching_file_ids:
unmatched_file = os.path.join(directory, file_name)
if problems_dir is not None:
shutil.move(unmatched_file,
os.path.join(problems_dir, file_name))
else:
os.remove(unmatched_file)
return matching_file_ids
def clean_darknet(
darknet_dir: str,
images_dir: str,
label_replacements: Dict,
label_removals: List[str] = None,
label_keep: List[str] = None,
problems_dir: str = None,
):
"""
TODO
:param darknet_dir:
:param images_dir:
:param label_replacements:
:param label_removals:
:param problems_dir:
:return:
"""
_logger.info("Cleaning dataset with Darknet annotations")
# convert all PNG images to JPG, and remove the original PNG file
for file_id in matching_ids(darknet_dir, images_dir, ".txt", ".png"):
png_file_path = os.path.join(images_dir, file_id + ".png")
png_to_jpg(png_file_path, remove_png=True)
# get the set of file IDs of the Darknet-format annotations and corresponding images
file_ids = purge_non_matching(images_dir, darknet_dir, "darknet",
problems_dir)
# loop over all the matching files and clean the Darknet annotations
for file_id in tqdm(file_ids):
# update the Darknet annotation file
src_annotation_file_path = os.path.join(darknet_dir, file_id + ".txt")
for line in fileinput.input(src_annotation_file_path, inplace=True):
# get the bounding box label
parts = line.split()
label = parts[0]
# skip rewriting this line if it's a label we want removed
if (label_removals is not None) and (label in label_removals):
continue
# skip rewriting this line if it's a label we do not want to keep
if (label_keep is not None) and (label not in label_keep):
continue
# get the bounding box coordinates
center_x = float(parts[1])
center_y = float(parts[2])
bbox_width = float(parts[3])
bbox_height = float(parts[4])
if (label_replacements is not None) and (label
in label_replacements):
# update the label
label = label_replacements[label]
# make sure we don't have wonky bounding box values
# and if so we'll skip them
if (center_x > 1.0) or (center_x < 0.0):
# report the issue via log message
_logger.warning(
"Bounding box center X is out of valid range -- skipping "
f"in Darknet annotation file {src_annotation_file_path}", )
continue
if (center_y > 1.0) or (center_y < 0.0):
# report the issue via log message
_logger.warning(
"Bounding box center Y is out of valid range -- skipping "
f"in Darknet annotation file {src_annotation_file_path}", )
continue
if (bbox_width > 1.0) or (bbox_width < 0.0):
# report the issue via log message
_logger.warning(
"Bounding box width is out of valid range -- skipping "
f"in Darknet annotation file {src_annotation_file_path}", )
continue
if (bbox_height > 1.0) or (bbox_height < 0.0):
# report the issue via log message
_logger.warning(
"Bounding box height is out of valid range -- skipping "
f"in Darknet annotation file {src_annotation_file_path}", )
continue
# write the line back into the file in-place
darknet_parts = [
label,
f'{center_x:.4f}',
f'{center_y:.4f}',
f'{bbox_width:.4f}',
f'{bbox_height:.4f}',
]
print(" ".join(darknet_parts))
def clean_pascal(
pascal_dir: str,
images_dir: str,
label_replacements: Dict = None,
label_removals: List[str] = None,
label_keep: List[str] = None,
problems_dir: str = None,
):
"""
TODO
:param pascal_dir:
:param images_dir:
:param label_replacements:
:param label_removals:
:param problems_dir:
:return:
"""
_logger.info("Cleaning dataset with PASCAL annotations")
# convert all PNG images to JPG, and remove the original PNG file
for file_id in matching_ids(pascal_dir, images_dir, ".xml", ".png"):
png_file_path = os.path.join(images_dir, file_id + ".png")
png_to_jpg(png_file_path, remove_png=True)
# get the set of file IDs of the Darknet-format annotations and corresponding images
file_ids = purge_non_matching(images_dir, pascal_dir, "pascal",
problems_dir)
# loop over all the matching files and clean the PASCAL annotations
for i, file_id in tqdm(enumerate(file_ids)):
# get the image width and height
jpg_file_name = file_id + ".jpg"
image_file_path = os.path.join(images_dir, jpg_file_name)
image = Image.open(image_file_path)
img_width, img_height = image.size
# update the image file name in the PASCAL file
src_annotation_file_path = os.path.join(pascal_dir, file_id + ".xml")
if os.path.exists(src_annotation_file_path):
tree = etree.parse(src_annotation_file_path)
root = tree.getroot()
size = tree.find("size")
width = int(size.find("width").text)
height = int(size.find("height").text)
if (width != img_width) or (height != img_height):
# something's amiss that we can't reasonably fix, remove files
if problems_dir is not None:
for file_path in [
src_annotation_file_path, image_file_path
]:
dest_file_path = os.path.join(
problems_dir,
os.path.split(file_path)[1])
shutil.move(file_path, dest_file_path)
else:
os.remove(src_annotation_file_path)
os.remove(image_file_path)
continue
# update the image file name
file_name = root.find("filename")
if (file_name is not None) and (file_name.text != jpg_file_name):
file_name.text = jpg_file_name
# loop over all bounding boxes
for obj in root.iter("object"):
# replace all bounding box labels if specified in the replacement dictionary
name = obj.find("name")
if (name is None) or ((label_removals is not None) and
(name.text in label_removals)):
# drop the bounding box
parent = obj.getparent()
parent.remove(obj)
# move on, nothing more to do for this box
continue
# skip rewriting this line if it's a label we do not want to keep
elif (name is None) or ((label_keep is not None) and
(name.text not in label_keep)):
# drop the bounding box
parent = obj.getparent()
parent.remove(obj)
# move on, nothing more to do for this box
continue
elif (label_replacements
is not None) and (name.text in label_replacements):
# update the label
name.text = label_replacements[name.text]
# for each bounding box make sure we have max
# values that are one less than the width/height
bbox = obj.find("bndbox")
bbox_min_x = int(float(bbox.find("xmin").text))
bbox_min_y = int(float(bbox.find("ymin").text))
bbox_max_x = int(float(bbox.find("xmax").text))
bbox_max_y = int(float(bbox.find("ymax").text))
# make sure we don't have wonky values with mins > maxs
if (bbox_min_x >= bbox_max_x) or (bbox_min_y >= bbox_max_y):
# drop the bounding box
_logger.warning(
"Dropping bounding box for object in file "
f"{src_annotation_file_path} due to invalid "
"min/max values", )
parent = obj.getparent()
parent.remove(obj)
else:
# make sure the max values don't go past the edge
if bbox_max_x >= img_width:
bbox.find("xmax").text = str(img_width - 1)
if bbox_max_y >= img_height:
bbox.find("ymax").text = str(img_height - 1)
# drop the image path, it's not reliable
path = root.find("path")
if path is not None:
parent = path.getparent()
parent.remove(path)
# drop the image folder, it's not reliable
folder = root.find("folder")
if folder is not None:
parent = folder.getparent()
parent.remove(folder)
# write the tree back to file
tree.write(src_annotation_file_path)
def clean_kitti(
kitti_dir: str,
images_dir: str,
label_replacements: Dict = None,
label_removals: List[str] = None,
label_keep: List[str] = None,
problems_dir: str = None,
):
"""
TODO
:param kitti_dir:
:param images_dir:
:param label_replacements:
:param label_removals:
:param problems_dir:
:return:
"""
_logger.info("Cleaning dataset with KITTI annotations")
# convert all PNG images to JPG, and remove the original PNG file
for file_id in matching_ids(kitti_dir, images_dir, ".txt", ".png"):
png_file_path = os.path.join(images_dir, file_id + ".png")
png_to_jpg(png_file_path, remove_png=True)
# get the set of file IDs of the Darknet-format annotations and corresponding images
file_ids = purge_non_matching(images_dir, kitti_dir, "kitti", problems_dir)
# loop over all the matching files and clean the KITTI annotations
for file_id in tqdm(file_ids):
# get the image width and height
jpg_file_name = file_id + ".jpg"
image_file_path = os.path.join(images_dir, jpg_file_name)
image = Image.open(image_file_path)
img_width, img_height = image.size
# update the image file name in the KITTI annotation file
src_annotation_file_path = os.path.join(kitti_dir, file_id + ".txt")
for line in fileinput.input(src_annotation_file_path, inplace=True):
parts = line.split()
label = parts[0]
# skip rewriting this line if it's a label we want removed
if (label_removals is not None) and (label in label_removals):
continue
# skip rewriting this line if it's a label we do not want to keep
if (label_keep is not None) and (label not in label_keep):
continue
truncated = parts[1]
occluded = parts[2]
alpha = parts[3]
bbox_min_x = int(float(parts[4]))
bbox_min_y = int(float(parts[5]))
bbox_max_x = int(float(parts[6]))
bbox_max_y = int(float(parts[7]))
dim_x = parts[8]
dim_y = parts[9]
dim_z = parts[10]
loc_x = parts[11]
loc_y = parts[12]
loc_z = parts[13]
rotation_y = parts[14]
# not all KITTI-formatted files have a score field
if len(parts) == 16:
score = parts[15]
else:
score = " "
if (label_replacements is not None) and (label
in label_replacements):
# update the label
label = label_replacements[label]
# make sure we don't have wonky bounding box values
# with mins > maxs, and if so we'll reverse them
if bbox_min_x > bbox_max_x:
# report the issue via log message
_logger.warning(
"Bounding box minimum X is greater than the maximum X "
f"in KITTI annotation file {src_annotation_file_path}", )
tmp_holder = bbox_min_x
bbox_min_x = bbox_max_x
bbox_max_x = tmp_holder
if bbox_min_y > bbox_max_y:
# report the issue via log message
_logger.warning(
"Bounding box minimum Y is greater than the maximum Y "
f"in KITTI annotation file {src_annotation_file_path}", )
tmp_holder = bbox_min_y
bbox_min_y = bbox_max_y
bbox_max_y = tmp_holder
# perform sanity checks on max values
if bbox_max_x >= img_width:
# report the issue via log message
_logger.warning(
"Bounding box maximum X is greater than width in KITTI "
f"annotation file {src_annotation_file_path}", )
# fix the issue
bbox_max_x = img_width - 1
if bbox_max_y >= img_height:
# report the issue via log message
_logger.warning(
"Bounding box maximum Y is greater than height in KITTI "
f"annotation file {src_annotation_file_path}", )
# fix the issue
bbox_max_y = img_height - 1
# write the line back into the file in-place
kitti_parts = [
label,
truncated,
occluded,
alpha,
f'{bbox_min_x:.1f}',
f'{bbox_min_y:.1f}',
f'{bbox_max_x:.1f}',
f'{bbox_max_y:.1f}',
dim_x,
dim_y,
dim_z,
loc_x,
loc_y,
loc_z,
rotation_y,
]
if len(parts) == 16:
kitti_parts.append(score)
print(" ".join(kitti_parts))
def masked_dataset_to_tfrecords(
images_dir: str,
masks_dir: str,
tfrecord_dir: str,
num_shards: int = 1,
dataset_base_name: str = "tfrecord",
train_pct: float = 1.0,
):
"""
Creates TFRecord files corresponding to a dataset of JPG images with
corresponding set PNG masks.
:param images_dir: directory containing image files
:param masks_dir: directory containing mask files corresponding to the images
:param tfrecord_dir: directory where the output TFRecord files will be written
:param num_shards: number of shards
:param dataset_base_name: base name of the TFRecord files to be produced
:param train_pct: the percentage of images/masks to use for training, with
(1.0 minus this value as the validation percentage), if this value is 1.0
then no split will occur
"""
masks_ext = ".png"
images_ext = ".jpg"
file_ids = list(matching_ids(masks_dir, images_dir, masks_ext, images_ext))
random.shuffle(file_ids)
# create a mapping of base file names and subsets of file IDs
if train_pct < 1.0:
# get the correct file name prefix for the TFRecord files
# based on the presence of a specified file base name
tfrecord_file_prefix_train = "train"
tfrecord_file_prefix_valid = "valid"
if dataset_base_name != "":
tfrecord_file_prefix_train = tfrecord_file_prefix_train + "_" + dataset_base_name
tfrecord_file_prefix_valid = tfrecord_file_prefix_valid + "_" + dataset_base_name
# get the split index to use for splitting into train/valid sets
split_index = int(len(file_ids) * train_pct)
# map the file prefixes to the sets of file IDs for the split sections
split_names_to_ids = {
tfrecord_file_prefix_train: file_ids[:split_index],
tfrecord_file_prefix_valid: file_ids[split_index:],
}
# report the number of samples in each split section
_logger.info(f"TFRecord dataset contains {len(file_ids[:split_index])} training samples")
_logger.info(f"TFRecord dataset contains {len(file_ids[split_index:])} validation samples")
else:
# we'll just have one base file name mapped to all file IDs
if "" == dataset_base_name:
tfrecord_file_prefix = "tfrecord"
else:
tfrecord_file_prefix = dataset_base_name
# map the file prefixes to the set of file IDs
split_names_to_ids = {
tfrecord_file_prefix: file_ids,
}
# report the number of samples
_logger.info(f"TFRecord dataset contains {len(file_ids)} samples (no train/valid split)")
# create an iterable of arguments that will be mapped to concurrent future processes
args_iterable = []
for base_name, file_ids in split_names_to_ids.items():
num_images = len(file_ids)
num_per_shard = int(math.ceil(num_images / num_shards))
for shard_id in range(num_shards):
output_filename = os.path.join(
tfrecord_dir,
f'{base_name}-{str(shard_id).zfill(5)}-of-{str(num_shards).zfill(5)}.tfrecord',
)
tfrecord_writing_args = {
"output_path": output_filename,
"shard_id": shard_id,
"num_per_shard": num_per_shard,
"num_images": num_images,
"file_ids": file_ids,
"images_dir": images_dir,
"masks_dir": masks_dir,
}
args_iterable.append(tfrecord_writing_args)
# use a ProcessPoolExecutor to facilitate creating the TFRecords in parallel
with concurrent.futures.ProcessPoolExecutor() as executor:
# map the TFRecord creation function to the iterable of arguments
_logger.info(f"Building TFRecords in directory {tfrecord_dir} ")
executor.map(_build_write_tfrecord, args_iterable)
def resize_dataset(
input_images_dir: str,
input_annotations_dir: str,
output_images_dir: str,
output_annotations_dir: str,
new_width: int,
new_height: int,
annotation_format: str,
):
"""
Resizes all images and corresponding annotations located within the
specified directories.
:param input_images_dir: directory where image files are located
:param input_annotations_dir: directory where annotation files are located
:param output_images_dir: directory where resized image files should be written
:param output_annotations_dir: directory where resized annotation files
should be written
:param new_width: new width to which the image should be resized
:param new_height: new height to which the image should be resized
:param annotation_format: "coco", "darknet", "kitti", or "pascal"
:return: the number of resized image/annotation files
"""
# only allow for KITTI and PASCAL annotation formats
if annotation_format not in ["kitti", "pascal"]:
raise ValueError(f"Unsupported annotation format: {annotation_format}")
else:
annotation_ext = cvdata.common.FORMAT_EXTENSIONS[annotation_format]
# create the destination directories in case they don't already exist
os.makedirs(output_annotations_dir, exist_ok=True)
os.makedirs(output_images_dir, exist_ok=True)
resize_arguments_list = []
for image_ext in (".jpg", ".png"):
# find matching annotation and image files (i.e. same base file name)
file_ids = matching_ids(
input_annotations_dir,
input_images_dir,
annotation_ext,
image_ext,
)
# loop over all image files and perform scaling/padding on each
for file_id in file_ids:
resize_arguments = {
"file_id": file_id,
"image_ext": image_ext,
"annotation_format": annotation_format,
"annotation_ext": annotation_ext,
"input_images_dir": input_images_dir,
"input_annotations_dir": input_annotations_dir,
"output_images_dir": output_images_dir,
"output_annotations_dir": output_annotations_dir,
"new_width": new_width,
"new_height": new_height,
}
resize_arguments_list.append(resize_arguments)
# use a ProcessPoolExecutor to download the images in parallel
with concurrent.futures.ProcessPoolExecutor() as executor:
_logger.info("Resizing files")
# use the executor to map the download function to the iterable of arguments
list(tqdm(executor.map(_resize_image_label, resize_arguments_list),
total=len(resize_arguments_list)))
def filter_class_boxes(
src_images_dir: str,
src_annotations_dir: str,
dest_images_dir: str,
dest_annotations_dir: str,
class_label_counts: Dict,
annotation_format: str,
darknet_labels_path: str = None,
):
"""
TODO
:param src_images_dir:
:param src_annotations_dir:
:param dest_images_dir:
:param dest_annotations_dir:
:param class_label_counts:
:param annotation_format:
:param darknet_labels_path: path to the labels file corresponding
to a Darknet-annotated dataset
"""
_logger.info(
f"Filtering dataset into annotations directory {dest_annotations_dir} "
f"and images directory {dest_images_dir}")
# make sure we don't have the same directories for src and dest
if src_images_dir == dest_images_dir:
raise ValueError(
"Source and destination image directories are "
"the same, must be different", )
if src_annotations_dir == dest_annotations_dir:
raise ValueError(
"Source and destination annotation directories are "
"the same, must be different", )
# determine the file extension to be used for annotations
if annotation_format not in ["darknet", "kitti"]:
raise ValueError(f"Unsupported annotation format: {annotation_format}")
else:
annotation_ext = FORMAT_EXTENSIONS[annotation_format]
image_ext = ".jpg"
# make the destination directories, in case they don't already exist
os.makedirs(dest_images_dir, exist_ok=True)
os.makedirs(dest_annotations_dir, exist_ok=True)
# keep a count of the boxes we've processed for each image class type
processed_class_counts = {k: 0 for k in class_label_counts.keys()}
# get all file IDs for image/annotation file matches
file_ids = \
matching_ids(
src_annotations_dir,
src_images_dir,
annotation_ext,
image_ext,
)
# only include the labels specified in the class counts dictionary
valid_labels = set(class_label_counts.keys())
# if we're processing Darknet annotations then read the labels file to get
# a mapping of indices to labels used with the Darknet annotation files
darknet_valid_indices = None
darknet_index_labels = None
if annotation_format == "darknet":
# read the Darknet labels into a dictionary mapping label to label index
darknet_index_labels = darknet_indices_to_labels(darknet_labels_path)
# get the set of valid indices, i.e. all Darknet indices
# corresponding to the labels to be included in the filtered dataset
darknet_valid_indices = set()
for index, darknet_label in darknet_index_labels.items():
if darknet_label in valid_labels:
darknet_valid_indices.add(index)
# loop over all the possible image/annotation file pairs
for file_id in tqdm(file_ids):
# only include the file if we find a box for one of the specified labels
include_file = False
# if any labels are found in the annotation file that aren't included
# in the list of image classes to filter then we'll want to remove the
# boxes from the annotation file before writing to the destination directory
irrelevant_labels_found = False
# get the count(s) of boxes per class label
annotation_file_name = file_id + annotation_ext
src_annotation_path = os.path.join(src_annotations_dir,
annotation_file_name)
box_counts = _count_boxes(src_annotation_path, annotation_format,
darknet_index_labels)
for class_label in box_counts.keys():
if class_label in class_label_counts:
processed_class_count = processed_class_counts[class_label]
if processed_class_counts[class_label] < class_label_counts[
class_label]:
include_file = True
processed_class_counts[
class_label] = processed_class_count + box_counts[
class_label]
else:
irrelevant_labels_found = True
if include_file:
dest_annotation_path = os.path.join(dest_annotations_dir,
annotation_file_name)
if irrelevant_labels_found:
# remove the unnecessary labels or indices from the
# annotation and write it into the destination directory
_write_with_removed_labels(
src_annotation_path,
dest_annotation_path,
annotation_format,
valid_labels,
darknet_valid_indices,
)
else:
# copy te annotation file into the destination directory as-is
shutil.copy(src_annotation_path, dest_annotation_path)
# copy the source image file into the destination images directory
image_file_name = file_id + image_ext
src_image_path = os.path.join(src_images_dir, image_file_name)
dest_image_path = os.path.join(dest_images_dir, image_file_name)
shutil.copy(src_image_path, dest_image_path)
def pascal_to_kitti(
pascal_dir: str,
images_dir: str,
kitti_data_dir: str,
kitti_ids_file_name: str = None,
move_image_files: bool = False,
) -> int:
"""
Builds KITTI annotation files from PASCAL VOC annotation XML files.
:param pascal_dir: directory containing input PASCAL VOC annotation
XML files, all XML files in this directory matching to corresponding JPG
files in the images directory will be converted to KITTI format
:param images_dir: directory containing image files corresponding to the
PASCAL VOC annotation files to be converted to KITTI format, all files
matching to PASCAL VOC annotation files will be either copied (default)
or moved (if move_image_files is True) to <kitti_data_dir>/image_2
:param kitti_data_dir: directory under which images will be copied or moved
into a subdirectory named "image_2" and KITTI annotation files will be
written into a subdirectory named "label_2"
:param kitti_ids_file_name: file name to contain all file IDs, to be written
into the parent directory above the <kitti_data_dir>
:param move_image_files: whether or not to move image files to
<kitti_data_dir>/image_2 (default is False and image files are copied
instead)
:return:
"""
_logger.info(f"Converting from PASCAL to KITTI for images in directory {images_dir}")
# create the KITTI directories in case they don't already exist
kitti_images_dir = os.path.join(kitti_data_dir, "image_2")
kitti_labels_dir = os.path.join(kitti_data_dir, "label_2")
for data_dir_name in (kitti_images_dir, kitti_labels_dir):
os.makedirs(os.path.join(kitti_data_dir, data_dir_name), exist_ok=True)
# assumed file extensions
img_ext = ".jpg"
pascal_ext = ".xml"
# get list of file IDs of the PASCAL VOC annotations and corresponding images
file_ids = matching_ids(pascal_dir, images_dir, pascal_ext, img_ext)
# write the KITTI IDs file in the KITTI directory's parent directory
if kitti_ids_file_name is not None:
kitti_ids_file_path = os.path.join(Path(kitti_data_dir).parent, kitti_ids_file_name)
with open(kitti_ids_file_path, "w") as kitti_ids_file:
for file_id in file_ids:
kitti_ids_file.write(f"{file_id}\n")
# build KITTI annotations from PASCAL and copy or
# move the image files into KITTI images directory
conversion_arguments_list = []
for file_id in file_ids:
conversion_arguments = {
"file_id": file_id,
"pascal_ext": pascal_ext,
"img_ext": img_ext,
"pascal_dir": pascal_dir,
"images_dir": images_dir,
"kitti_labels_dir": kitti_labels_dir,
"kitti_images_dir": kitti_images_dir,
"move_image_files": move_image_files,
}
conversion_arguments_list.append(conversion_arguments)
# use a ProcessPoolExecutor to download the images in parallel
with concurrent.futures.ProcessPoolExecutor() as executor:
# use the executor to map the download function to the iterable of arguments
_logger.info(f"Building KITTI labels in directory {kitti_labels_dir} ")
list(tqdm(executor.map(single_pascal_to_kitti, conversion_arguments_list),
total=len(conversion_arguments_list)))
# return the number of annotations converted
return len(file_ids)
def kitti_to_darknet(
images_dir: str,
kitti_dir: str,
darknet_dir: str,
darknet_labels: str,
):
"""
Creates equivalent Darknet (YOLO) annotation files corresponding to a dataset
with KITTI annotations.
:param images_dir: directory containing the dataset's images
:param kitti_dir: directory containing the dataset's KITTI annotation files
:param darknet_dir: directory where the equivalent Darknet annotation files
will be written
:param darknet_labels: labels file corresponding to the label indices used
in the Darknet annotation files, will be written into the specified
Darknet annotations directory
"""
_logger.info("Converting annotations in KITTI format to Darknet format equivalents")
# create the Darknet annotations directory in case it doesn't yet exist
os.makedirs(darknet_dir, exist_ok=True)
# get list of file IDs of the KITTI annotations and corresponding images
annotation_ext = ".txt"
image_ext = ".jpg"
file_ids = matching_ids(kitti_dir, images_dir, annotation_ext, image_ext)
# dictionary of labels to indices
label_indices = {}
# build Darknet annotations from KITTI
for file_id in tqdm(file_ids):
# get the image's dimensions
image_file_name = file_id + image_ext
width, height, _ = image_dimensions(os.path.join(images_dir, image_file_name))
# loop over all annotation lines in the KITTI file and compute Darknet equivalents
annotation_file_name = file_id + annotation_ext
with open(os.path.join(kitti_dir, annotation_file_name), "r") as kitti_file:
darknet_bboxes = []
for line in kitti_file:
parts = line.split()
label = parts[0]
if label in label_indices:
label_index = label_indices[label]
else:
label_index = len(label_indices)
label_indices[label] = label_index
box_width_pixels = float(parts[6]) - float(parts[4]) + 1
box_height_pixels = float(parts[7]) - float(parts[5]) + 1
darknet_bbox = {
"label_index": label_index,
"center_x": ((box_width_pixels / 2) + float(parts[4])) / width,
"center_y": ((box_height_pixels / 2) + float(parts[5])) / height,
"box_width": box_width_pixels / width,
"box_height": box_height_pixels / height,
}
darknet_bboxes.append(darknet_bbox)
# write the Darknet annotation boxes into a Darknet annotation file
with open(os.path.join(darknet_dir, annotation_file_name), "w") as darknet_file:
for darknet_bbox in darknet_bboxes:
darknet_file.write(
f"{darknet_bbox['label_index']} {darknet_bbox['center_x']} "
f"{darknet_bbox['center_y']} {darknet_bbox['box_width']} "
f"{darknet_bbox['box_height']}\n"
)
# write the Darknet labels into a text file, one label per line,
# in order according to the indices used in the annotation files
with open(os.path.join(darknet_dir, darknet_labels), "w") as darknet_labels_file:
index_labels = {v: k for k, v in label_indices.items()}
for i in range(len(index_labels)):
darknet_labels_file.write(f"{index_labels[i]}\n")
def _dataset_bbox_examples(
images_dir: str,
annotations_dir: str,
annotation_format: str,
darknet_labels: str = None,
) -> pd.DataFrame:
"""
:param images_dir: directory containing the dataset's *.jpg image files
:param annotations_dir: directory containing the dataset's annotation files
:param annotation_format: currently supported: "darknet", "kitti", and "pascal"
:param darknet_labels: path to the class labels file corresponding to Darknet
(YOLO) annotation files, only necessary if using "darknet" annotation format
:return: pandas DataFrame with rows corresponding to the dataset's bounding boxes
"""
# we expect all images to use the *.jpg extension
image_ext = ".jpg"
# list of bounding box annotations we'll eventually write to CSV
bboxes = []
if annotation_format == "pascal":
# get the file IDs for all matching image/PASCAL pairs (i.e. the dataset)
annotation_ext = ".xml"
for file_id in matching_ids(
annotations_dir,
images_dir,
annotation_ext,
image_ext,
):
# add all bounding boxes from the PASCAL file to the list of boxes
pascal_path = os.path.join(annotations_dir, file_id + annotation_ext)
tree = ElementTree.parse(pascal_path)
root = tree.getroot()
for member in root.findall('object'):
bbox_values = (
root.find('filename').text,
int(root.find('size')[0].text),
int(root.find('size')[1].text),
member[0].text,
int(member[4][0].text),
int(member[4][1].text),
int(member[4][2].text),
int(member[4][3].text),
)
bboxes.append(bbox_values)
elif annotation_format == "kitti":
# get the file IDs for all matching image/KITTI pairs (i.e. the dataset)
annotation_ext = ".txt"
for file_id in matching_ids(
annotations_dir,
images_dir,
annotation_ext,
image_ext,
):
# get the image dimensions from the image file since this
# info is not present in the corresponding KITTI annotation
image_file_name = file_id + image_ext
image_path = os.path.join(images_dir, image_file_name)
width, height, _ = image_dimensions(image_path)
# add all bounding boxes from the KITTI file to the list of boxes
kitti_path = os.path.join(annotations_dir, file_id + annotation_ext)
with open(kitti_path, "r") as kitti_file:
for line in kitti_file:
darknet_box = line.split()
bbox_values = (
image_file_name,
width,
height,
darknet_box[0],
darknet_box[4],
darknet_box[5],
darknet_box[6],
darknet_box[7],
)
bboxes.append(bbox_values)
elif annotation_format == "darknet":
# read class labels into index/label dictionary
darknet_index_labels = darknet_indices_to_labels(darknet_labels)
# get the file IDs for all matching image/Darknet pairs (i.e. the dataset)
annotation_ext = ".txt"
file_ids = matching_ids(
annotations_dir,
images_dir,
annotation_ext,
image_ext,
)
# get the bounding boxes from the annotation files
_logger.info("Extracting bounding box info from Darknet annotations...")
for file_id in tqdm(file_ids):
# get the image dimensions from the image file since this
# info is not present in the corresponding KITTI annotation
image_file_name = file_id + image_ext
image_path = os.path.join(images_dir, image_file_name)
width, height, _ = image_dimensions(image_path)
# add all bounding boxes from the Darknet file to the list of boxes
darknet_path = os.path.join(annotations_dir, file_id + annotation_ext)
with open(darknet_path, "r") as darknet_file:
for line in darknet_file:
darknet_box = line.split()
label_index = int(darknet_box[0])
# only use annotations corresponding to the specified labels
if label_index not in darknet_index_labels:
# skip this annotation line
continue
center_x = float(darknet_box[1]) * width
center_y = float(darknet_box[2]) * height
box_width = float(darknet_box[3]) * width
box_height = float(darknet_box[4]) * height
bbox_values = (
image_file_name,
width,
height,
darknet_index_labels[label_index],
int(center_x - (box_width / 2)),
int(center_y - (box_height / 2)),
int(center_x + (box_width / 2)),
int(center_y + (box_height / 2)),
)
bboxes.append(bbox_values)
else:
raise ValueError(f"Unsupported annotation format: {annotation_format}")
# stuff the bounding boxes into a pandas DataFrame
column_names = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax']
return pd.DataFrame(bboxes, columns=column_names)
def main():
# parse the command line arguments
args_parser = argparse.ArgumentParser()
args_parser.add_argument(
"--annotations",
required=True,
type=str,
help="annotations directory path",
)
args_parser.add_argument(
"--images",
required=False,
type=str,
help="images directory path",
)
args_parser.add_argument(
"--format",
type=str,
required=True,
choices=cvdata.common.FORMAT_CHOICES,
help="annotation format",
)
args_parser.add_argument(
"--file_ids",
type=str,
required=False,
help="directory in which to write text files with file IDs for each "
"image/annotation in the dataset that contains a label, with "
"one file per label",
)
args = vars(args_parser.parse_args())
if args["format"] == "tfrecord":
# count and report the examples in the collection of TFRecord files
examples_count = count_tfrecord_examples(args["annotations"])
print(f"Total number of examples: {examples_count}")
else:
# the two dictionaries we'll build for final reporting
label_counts = {}
label_file_ids = {}
if args["format"] == "openimages":
# read the OpenImages CSV into a pandas DataFrame
df_annotations = pd.read_csv(args["annotations"])
df_annotations = df_annotations[['ImageID', 'LabelName']]
# TODO get another dataframe from the class descriptions and get the
# readable label names from there to map to the LabelName column
# whittle it down to only the rows that match to image IDs
file_ids = [os.path.splitext(file_name)[0] for file_name in os.listdir(args["images"])]
df_annotations = df_annotations[df_annotations["ImageID"].isin(file_ids)]
# TODO populate the label counts and label file IDs dictionaries
else:
annotation_ext = cvdata.common.FORMAT_EXTENSIONS[args["format"]]
# only annotations matching to the images are considered to be valid
file_ids = matching_ids(args["annotations"], args["images"], annotation_ext, ".jpg")
for file_id in file_ids:
annotation_file_path = \
os.path.join(args["annotations"], file_id + annotation_ext)
# get the images per label count
for label, count in count_labels(annotation_file_path, args["format"]).items():
if label in label_counts:
label_counts[label] += 1
else:
label_counts[label] = 1
# for each label found in the annotation file add this file ID
# to the set of file IDs corresponding to the label
if args["file_ids"]:
if label in label_file_ids:
# add this file ID to the existing set for the label
label_file_ids[label].add(file_id)
else:
# first file ID seen for this label so create new set
label_file_ids[label] = {file_id}
# write the images per label counts
for label, count in label_counts.items():
print(f"Label: {label}\t\tCount: {count}")
# write the label ID files, if requested
if args["file_ids"]:
for label, file_ids_for_label in label_file_ids.items():
label_file_ids_path = os.path.join(args["file_ids"], label + ".txt")
with open(label_file_ids_path, "w") as label_file_ids_file:
for file_id in file_ids_for_label:
label_file_ids_file.write(f"{file_id}\n")