SquirrelDetector/scripts/copy_labels.py

import argparse
import os
import shutil

def copy_labels(image_path, from_labels, to_labels):
    for img in os.listdir(image_path):
        label = f'{os.path.splitext(img)[0]}.txt'
        copy_from = os.path.join(from_labels, label)
        if not os.path.exists(copy_from):
            print(f"    WARNING: no label file found for '{copy_from}'. This may just mean that no objects were detected.")
            continue
        shutil.copy(os.path.join(from_labels, label), to_labels)

def main():
    parser = argparse.ArgumentParser('Copies label files from a cental path to the appropriate YOLO-organized dataset (train, val, test) based on the image files present')
    parser.add_argument('--labels', required=True, help='Path to directory with label files')
    parser.add_argument('--dataset', required=True, help='Path to YOLO organized dataset')
    
    args = parser.parse_args()

    for dataset in ['train', 'val', 'test']:
        image_path = os.path.join(args.dataset, 'images', dataset)
        if not os.path.exists(image_path):
            print(f"'{image_path}' does not exist. Skipping...")
            continue
        to_labels = os.path.join(args.dataset, 'labels', dataset) 
        os.makedirs(to_labels, exist_ok=True)
        copy_labels(image_path, args.labels, to_labels)


if __name__ == '__main__':
    main()