Add batch-size to ImageHandler CLI
@@ -216,11 +216,18 @@ class ImageHandler(BentoHandler):
parser = argparse.ArgumentParser()
parser.add_argument("--input", required=True, nargs='+')
parser.add_argument("-o", "--output", default="str", choices=["str", "json"])
+ parser.add_argument("--batch-size", default=None, type=int)
parsed_args = parser.parse_args(args)
file_paths = parsed_args.input
+ batch_size = (
+ parsed_args.batch_size if parsed_args.batch_size else len(file_paths)
+ )
+
+ for i in range(0, len(file_paths), batch_size):
+ step_file_paths = file_paths[i : i + batch_size]
image_arrays = []
- for file_path in file_paths:
+ for file_path in step_file_paths:
verify_image_format_or_raise(file_path, self.accept_image_formats)
if not os.path.isabs(file_path):
file_path = os.path.abspath(file_path)
@@ -228,6 +235,7 @@ class ImageHandler(BentoHandler):
image_arrays.append(self.imread(file_path, pilmode=self.pilmode))
results = func(image_arrays)
+
for result in results:
if parsed_args.output == "json":
result = api_func_result_to_json(result)
