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)