import os
import argparse

def prepare_coco(args):
    train_file = open(os.path.join(args.saved_dir, args.saved_txt), 'w')
    dirs = os.listdir(args.raw_dir)

    for file in dirs:
        # Write training files
        train_file.write('%s\n' % (file))

    print('Data Preparation Finished.')

if __name__ == '__main__':
    arg_parser = argparse.ArgumentParser(description="coco prepareing.")
    arg_parser.add_argument('--dataset', type=str, default='coco',
                             help='')
    arg_parser.add_argument('--raw_dir', type=str, default='',
                             help='')
    arg_parser.add_argument('--saved_dir', type=str, default='',
                             help='')
    arg_parser.add_argument('--saved_txt', type=str, default='train2017.txt',
                             help='')
    args = arg_parser.parse_args() 

    prepare_coco(args)