import os
import shutil
from argparse import ArgumentParser
from pathlib import Path
from xml.etree import ElementTree
import yaml
from tqdm import tqdm
from kitcar_ml.utils.data import image_folder
from kitcar_ml.utils.data.labeled_dataset import LabeledDataset
[docs]def import_kitcar_xml_labels(input_dir: str, output_dir: str):
"""Convert a XML based Dataset to our LabeledDataset Format.
See: https://doc.kitcar-team.de/kitcar-machine-learning/tutorials/datasets.html
"""
images_and_annotations = []
for image_path in image_folder.find_images(input_dir):
base_dir = os.path.dirname(image_path)
# As the annotations folder naming isn't consistent, we need to search for it
sub_dirs = next(os.walk(base_dir))[1]
if len(sub_dirs) != 1:
# Can't identify a annotations directory
continue
annotation = os.path.join(base_dir, sub_dirs[0], f"{Path(image_path).stem}.xml")
if not os.path.exists(annotation):
continue
images_and_annotations.append((image_path, annotation))
num_elements = len(images_and_annotations)
assert num_elements > 0, "Input directory has no images with corresponding annotations."
# As the images in different sub directories could have the same name,
# we need to rename them.
# Otherwise the keys in the LabeledDataset wouldn't we distinct.
img_name_counter = 0
class_id_counter = 0
zeros_for_img_name = len(str(num_elements))
# Create Labeled Dataset
dataset = LabeledDataset()
dataset.attributes = ["x1", "y1", "x2", "y2", "class_id"]
# Create output directory
os.makedirs(output_dir, exist_ok=True)
assert (
len(os.listdir(output_dir)) == 0
), "Output directory is not empty, please empty it manually"
with open(os.path.join(os.path.dirname(__file__), "label_conversion.yaml")) as file:
name_conversion = yaml.load(file, Loader=yaml.SafeLoader)
print("Import xml labels ...")
for image_path, annotation_path in tqdm(images_and_annotations):
xml_document = ElementTree.parse(annotation_path)
root = xml_document.getroot()
# Set new image name and copy to output directory
new_image_name = (
str(img_name_counter).zfill(zeros_for_img_name)
+ os.path.splitext(image_path)[1]
)
img_name_counter += 1
shutil.copyfile(image_path, os.path.join(output_dir, new_image_name))
# Create empty labels array for image
dataset.labels[new_image_name] = []
for annotation in root.findall("object"): # Iterate over all annotations
# Get the new class name
old_name = str(annotation.find("name").text)
if old_name not in name_conversion:
raise ValueError("No conversion existing for: " + old_name)
name = name_conversion[old_name]
try:
# Already a class with the same name, use it's id.
class_id = next(
(
class_id
for class_id, class_name in dataset.classes.items()
if class_name == name
)
)
except StopIteration:
# New class
class_id = class_id_counter
dataset.classes[class_id] = name
class_id_counter += 1
bounding_box = [
int(annotation.find("bndbox/xmin").text),
int(annotation.find("bndbox/ymin").text),
int(annotation.find("bndbox/xmax").text),
int(annotation.find("bndbox/ymax").text),
class_id,
]
dataset.append_label(new_image_name, bounding_box)
# Save dataset as yaml file
dataset.save_as_yaml(os.path.join(output_dir, "labels.yaml"))
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument(
"--input-dir",
type=str,
required=True,
help="The root dir of the whole dataset.",
)
parser.add_argument(
"--output-dir",
type=str,
required=True,
help="The output dir of the whole dataset.",
)
args = parser.parse_args()
import_kitcar_xml_labels(args.input_dir, args.output_dir)