diff options
Diffstat (limited to 'gantools')
| -rw-r--r-- | gantools/cli.py | 10 | ||||
| -rw-r--r-- | gantools/image_utils.py | 7 |
2 files changed, 10 insertions, 7 deletions
diff --git a/gantools/cli.py b/gantools/cli.py index 73213ed..0ad3a8d 100644 --- a/gantools/cli.py +++ b/gantools/cli.py @@ -16,7 +16,8 @@ def main(): parser.add_argument('-b', '--nbatch', metavar='N', type=int, help='Number of frames in each \'batch\' \ (note: the truncation value can only change once per batch. Don\'t fuck with this unless you know \ what it does.).', default=1) - parser.add_argument('-f', '--pathprefix', help='Directory path and file prefix for output images.') + parser.add_argument('-o', '--output-dir', help='Directory path for output images.') + parser.add_argument('--prefix', help='File prefix for output images.') args = parser.parse_args() # validate args if args.keys and not (args.username and args.password): @@ -38,8 +39,9 @@ def main(): ims = gan.sample(z_seq, label_seq, truncation_seq, args.nbatch) # save images to file - pathprefix = '' if args.pathprefix == None else str(args.pathprefix) - print('Saving image files: '+pathprefix) - image_utils.save_images(ims, pathprefix) + path = '' if args.output_dir == None else str(args.output_dir) + prefix = '' if args.prefix == None else str(args.prefix) + print('Saving image files: '+path + prefix) + image_utils.save_images(ims, output_dir=output_dir, prefix=prefix) print('Done.') diff --git a/gantools/image_utils.py b/gantools/image_utils.py index 016891f..72130b0 100644 --- a/gantools/image_utils.py +++ b/gantools/image_utils.py @@ -1,10 +1,11 @@ +import os import PIL.Image def save_image(arr, fp, format='JPEG'): image = PIL.Image.fromarray(arr) image.save(fp, format=format, quality=90) -def save_images(ims, path_prefix='', format='JPEG'): +def save_images(ims, output_dir='', prefix='', format='JPEG'): for i, im in enumerate(ims): - path = str(path_prefix)+str(i).zfill(4)+'.'+str(format).lower() - save_image(im, path) + full_path = os.path.join(output_dir, prefix + str(i).zfill(4) + '.' + format.lower()) + save_image(im, full_path, format) |