import glob
import logging
import os
import shutil
import sys

"""script to divide a folder with generated/training data into a train and val folder
    - val folder contains 500 Samples if not changed in source code
    - DOES NOT work if images structured in subfolders, see below
    - if there is no dir in the given folder -> split this folder
    - if there are dir/s in the folder -> perform split on each folder
    - split on sorted list -> repeated runs should give the same result
    """


def main(args):
    foldername = args[1]
    print("CWD: {}".format(os.getcwd()))
    print("foldername: {}".format(foldername))

    dirs = os.walk(foldername).next()[1]
    dirs = [os.path.join(foldername, x) for x in dirs]
    print(dirs)
    if len(dirs) == 0:
        print("no subdirs found -> run directly on {}".format(foldername))
        dirs = [foldername]

    for dir in dirs:
        print("perform split on {}".format(dir))
        dir_path = dir
        # image_list = sorted(glob.glob1(os.path.join(foldername, dir_path), "*.jpg"))
        image_list = sorted(glob.glob1(dir_path, "*.jpg"))
        # image_list = sorted(glob.glob1(dir_path , "*.png"))
        if len(image_list) == 0:
            logging.error("Could not find any '*.jpg' in {}".format(dir_path))
            exit(1)
        else:
            print("  found {} images".format(len(image_list)))

        # val_len = int(len(image_list) * 0.1)
        val_len = int(500)

        val_list = image_list[:val_len]
        train_list = image_list[val_len:]
        #  save first 10%/500 of list to val list

        for subdir, part_list in zip(["val", "train"], [val_list, train_list]):
            os.makedirs(os.path.join(dir_path, subdir))
            print("  move files in {}...".format(subdir))
            for image_file in part_list:
                shutil.move(os.path.join(dir_path, image_file), os.path.join(dir_path, subdir, image_file))
                try:
                    shutil.move(os.path.join(dir_path, image_file + ".txt"),
                                os.path.join(dir_path, subdir, image_file + ".txt"))
                except IOError as ex:
                    print(ex)
                try:
                    shutil.move(os.path.join(dir_path, image_file + ".info"),
                                os.path.join(dir_path, subdir, image_file + ".info"))
                except IOError as ex:
                    pass

            print("  write list: {}...".format(os.path.join(dir_path, "{}_{}.lst".format(dir_path, subdir))))
            with open(os.path.join(foldername, "{}_{}.lst".format(os.path.basename(dir_path), subdir)), "w") as fobj:
                fobj.writelines([os.path.join(dir_path, subdir, x) + "\n" for x in part_list])


if __name__ == '__main__':
    main(sys.argv)

[((22, 12, 22, 39), 'os.path.join', 'os.path.join', ({(22, 25, 22, 35): 'foldername', (22, 37, 22, 38): 'x'}, {}), '(foldername, x)', False, 'import os\n'), ((18, 27, 18, 38), 'os.getcwd', 'os.getcwd', ({}, {}), '()', False, 'import os\n'), ((32, 28, 32, 57), 'glob.glob1', 'glob.glob1', ({(32, 39, 32, 47): 'dir_path', (32, 49, 32, 56): '"""*.jpg"""'}, {}), "(dir_path, '*.jpg')", False, 'import glob\n'), ((21, 11, 21, 30), 'os.walk', 'os.walk', ({(21, 19, 21, 29): 'foldername'}, {}), '(foldername)', False, 'import os\n'), ((48, 24, 48, 54), 'os.path.join', 'os.path.join', ({(48, 37, 48, 45): 'dir_path', (48, 47, 48, 53): 'subdir'}, {}), '(dir_path, subdir)', False, 'import os\n'), ((51, 28, 51, 62), 'os.path.join', 'os.path.join', ({(51, 41, 51, 49): 'dir_path', (51, 51, 51, 61): 'image_file'}, {}), '(dir_path, image_file)', False, 'import os\n'), ((51, 64, 51, 106), 'os.path.join', 'os.path.join', ({(51, 77, 51, 85): 'dir_path', (51, 87, 51, 93): 'subdir', (51, 95, 51, 105): 'image_file'}, {}), '(dir_path, subdir, image_file)', False, 'import os\n'), ((53, 32, 53, 75), 'os.path.join', 'os.path.join', ({(53, 45, 53, 53): 'dir_path', (53, 55, 53, 74): "(image_file + '.txt')"}, {}), "(dir_path, image_file + '.txt')", False, 'import os\n'), ((54, 32, 54, 83), 'os.path.join', 'os.path.join', ({(54, 45, 54, 53): 'dir_path', (54, 55, 54, 61): 'subdir', (54, 63, 54, 82): "(image_file + '.txt')"}, {}), "(dir_path, subdir, image_file + '.txt')", False, 'import os\n'), ((58, 32, 58, 76), 'os.path.join', 'os.path.join', ({(58, 45, 58, 53): 'dir_path', (58, 55, 58, 75): "(image_file + '.info')"}, {}), "(dir_path, image_file + '.info')", False, 'import os\n'), ((59, 32, 59, 84), 'os.path.join', 'os.path.join', ({(59, 45, 59, 53): 'dir_path', (59, 55, 59, 61): 'subdir', (59, 63, 59, 83): "(image_file + '.info')"}, {}), "(dir_path, subdir, image_file + '.info')", False, 'import os\n'), ((64, 66, 64, 92), 'os.path.basename', 'os.path.basename', ({(64, 83, 64, 91): 'dir_path'}, {}), '(dir_path)', False, 'import os\n'), ((65, 33, 65, 66), 'os.path.join', 'os.path.join', ({(65, 46, 65, 54): 'dir_path', (65, 56, 65, 62): 'subdir', (65, 64, 65, 65): 'x'}, {}), '(dir_path, subdir, x)', False, 'import os\n')]