yolov4-baby-yoda/prepare.py

import os
import random

import argh

val_ratio = 0.2


def prepare(images_dir, yolo_labels_dir, output_dir):
    """
    Prepare data for training
    """
    yolo_labels_dir = os.path.join(yolo_labels_dir, "labels")
    image_list = os.listdir(images_dir)
    train_lines = []
    val_lines = []
    for txt_file in os.listdir(yolo_labels_dir):
        # Check text file
        if not txt_file.endswith(".txt"):
            continue

        # Check image exists
        image_name, _ = os.path.splitext(txt_file)
        try:
            image_name_with_ext = next(image for image in image_list if image.startswith(image_name))
        except StopIteration:
            print("Image {} not found".format(image_name))
            continue

        # Extract labels
        with open(os.path.join(yolo_labels_dir, txt_file), "r") as f:
            image_labels = []
            for label_data in f.read().splitlines():
                obj_class, x, y, width, height = label_data.split(" ")
                x, y, width, height = float(x), float(y), float(width), float(height)
                x2, y2 = x + width, y + height
                size = 608
                x, y, x2, y2 = str(x * size), str(y * size), str(x2 * size), str(y2 * size)
                train_format_label = ",".join([x, y, x2, y2, obj_class])
                image_labels.append(train_format_label)

            if random.uniform(0, 1) < val_ratio:
                val_lines.append("{} {}".format(image_name_with_ext, " ".join(image_labels)))
            else:
                train_lines.append("{} {}".format(image_name_with_ext, " ".join(image_labels)))

    os.makedirs(output_dir, exist_ok=True)
    with open(os.path.join(output_dir, "train.txt"), "w") as f:
        f.write("\n".join(train_lines))
    with open(os.path.join(output_dir, "val.txt"), "w") as f:
        f.write("\n".join(val_lines))



# assembling:
parser = argh.ArghParser()
parser.add_commands([
    prepare,
])

if __name__ == "__main__":
    parser.dispatch()