SquirrelDetector/scripts/convert_labels.py

import argparse
import os

# These labels (zero-indexed) correspond to:
# 14: bird
# 15: cat
# 16: dog
# 17: horse
# 18: sheep
# 19: cow
# 21: bear
LABELS = {'14', '15', '16', '17', '18', '19', '21'}

def convert_label(filename, inpath, outpath):
    infile = os.path.join(inpath, filename)
    outfile = os.path.join(outpath, filename)

    with open(infile) as f, open(outfile, mode='w') as of:
        for line in f:
            parts = line.split()
            if parts[0] in LABELS:
                parts[0] = '0'
                newline = ' '.join(parts[:5])
                of.write(f'{newline}\n')
        

def main():
    parser = argparse.ArgumentParser('Converts COCO labels for most animals to a single class within YOLO-style txt files')
    parser.add_argument('--input', help='Input directory with YOLO-style txt files, using COCO labels')
    parser.add_argument('--output', help='Output directory to write the YOLO-style txt files with the converted labels')

    args = parser.parse_args()

    os.makedirs(args.output, exist_ok=True)

    for filename in os.listdir(args.input):
        ext = os.path.splitext(filename)[-1].lower()
        if ext != '.txt':
            continue      

        convert_label(filename, args.input, args.output)


if __name__ == '__main__':
    main()