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 ...@@ -12,7 +12,7 @@ from shutil import copyfile, move
import argparse 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) parent_list = os.listdir(path)
for child in parent_list: for child in parent_list:
...@@ -31,7 +31,7 @@ def create_small_dataset(path, new_path, dataset_size, move_it=True): ...@@ -31,7 +31,7 @@ def create_small_dataset(path, new_path, dataset_size, move_it=True):
subdir_list = os.listdir(subdir_path) subdir_list = os.listdir(subdir_path)
count = 0 count = 0
for file in subdir_list: 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) file_path = os.path.join(subdir_path, file)
new_file_path = os.path.join(new_subdir_path, file) new_file_path = os.path.join(new_subdir_path, file)
if move_it: if move_it:
...@@ -53,6 +53,6 @@ if __name__ == "__main__": ...@@ -53,6 +53,6 @@ if __name__ == "__main__":
parser.add_argument('--source', required=True, help="input directory") parser.add_argument('--source', required=True, help="input directory")
parser.add_argument('--target', required=True, help="output 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('--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() cli_args = parser.parse_args()
create_small_dataset(cli_args.source, cli_args.target, cli_args.size, cli_args.move) 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