To generate train.txt/test.txt for YOLO object detection

traingen.py

'''
Usage: python3 classgen.py --input path >> train.txt
'''

import os
import argparse

def argsparser():
    '''
    To use arguments from command line.
    '''
    parser = argparse.ArgumentParser(description='Renamer script')
    parser.add_argument('--input', dest='input', help='Input Directory', default='path', type=str)
    return parser.parse_args()

if __name__ == '__main__':
    args = argsparser()
    files = os.listdir(args.input)
    
    for f in files:
        if 'jpg' in f:
            print('data/obj/' + f)

Tidy little script, just change the usage comment to traingen.py
Cheers!