Commit d5d71668 authored by Paul Bethge's avatar Paul Bethge
Browse files

fix count bug

parent 6f1c1db3
......@@ -12,7 +12,7 @@ from shutil import copyfile, move
import argparse
def create_small_dataset(path, new_path, dataset_size, move_it=True):
def create_small_dataset(path, new_path, dataset_size, move_it):
parent_list = os.listdir(path)
for child in parent_list:
......@@ -31,7 +31,7 @@ def create_small_dataset(path, new_path, dataset_size, move_it=True):
subdir_list = os.listdir(subdir_path)
count = 0
for file in subdir_list:
if dataset_size != -1 or count < dataset_size:
if dataset_size == -1 or count < dataset_size:
file_path = os.path.join(subdir_path, file)
new_file_path = os.path.join(new_subdir_path, file)
if move_it:
......@@ -53,6 +53,6 @@ if __name__ == "__main__":
parser.add_argument('--source', required=True, help="input directory")
parser.add_argument('--target', required=True, help="output directory")
parser.add_argument('--size', type=int, default=10, help="put -1 for all")
parser.add_argument('--move', type=bool, default=True, help="whether to move or just copy")
parser.add_argument('--move', type=bool, default=False, help="whether to move or just copy")
cli_args = parser.parse_args()
create_small_dataset(cli_args.source, cli_args.target, cli_args.size, cli_args.move)
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment