CVAT zip annotation creator

create_cvat_yolo_zip.py

from os import listdir as ls, makedirs, path as osp
from shutil import copy2 as cp, make_archive
from tempfile import TemporaryDirectory


def create_cvat_yolo_zip(path_source, path_classes=None, path_target=None, copy_images=False):
    assert osp.isdir(path_source), f"'path_source' must be a directory with [flat] YOLO annotations!"
    path_classes = path_classes or osp.join(path_source, 'classes.txt')
    assert osp.isfile(path_classes), f"either 'path_classes' or 'classes.txt' in 'source_path' must exist!"
    path_target = path_target or '.'
    assert osp.isdir(path_target), f"'path_target' must be a valid directory!"
    annotations = [fn for fn in ls(path_source) if fn.endswith('.txt') and fn != 'classes.txt']
    names = set(map(lambda a: a.rstrip('.txt'), annotations))
    images = [fn for fn in ls(path_source) if fn.rsplit('.', 1)[0] in names and not fn in annotations]
    train_txt = []
    with open(path_classes) as pc:
        classes = pc.read().splitlines()
        num_classes = len(classes)
    obj_data = f"classes = {num_classes}\ntrain = data/train.txt\nnames = data/obj.names\nbackup = backup/\n"
    with TemporaryDirectory() as tmp:
        # Create 'obj_train_data'
        obj_train_data = osp.join(tmp, 'obj_train_data')
        makedirs(obj_train_data)
        for image in images:
            # Copy annotations (optionally: images) <-- TODO
            train_txt.append(f"data/obj_train_data/{image}\n")
            annotation = osp.join(path_source, f"{osp.splitext(image)[0]}.txt")
            cp(annotation, obj_train_data)
            if copy_images:
                # cp(osp.join(path_source, image), obj_train_data)
                ...
        # Create 'obj.names'
        cp(path_classes, osp.join(tmp, 'obj.names'))
        # Create 'obj.data'
        with open(osp.join(tmp, 'obj.data'), 'w') as output:
            output.write(obj_data)
        # Create 'train.txt'
        with open(osp.join(tmp, 'train.txt'), 'w') as output:
            output.writelines(train_txt)
        make_archive(osp.join(path_target, osp.basename(path_source)), 'zip', tmp)
    print(f"Done!")