添加对一个目录的文件的支持 (#427)

tools/infer/predict.py

@@ -30,6 +30,7 @@ def parse_args():
 
     parser = argparse.ArgumentParser()
     parser.add_argument("-i", "--image_file", type=str)
+    parser.add_argument("-d", "--image_dir", type=str)
     parser.add_argument("-m", "--model_file", type=str)
     parser.add_argument("-p", "--params_file", type=str)
     parser.add_argument("-b", "--batch_size", type=int, default=1)
@@ -93,6 +94,8 @@ def preprocess(fname, ops):
 
 
 def main():
+    import os
+
     args = parse_args()
 
     if not args.enable_benchmark:
@@ -102,6 +105,8 @@ def main():
         assert args.use_gpu is True
         assert args.model_name is not None
         assert args.use_tensorrt is True
+        assert args.image_file is not None
+
     # HALF precission predict only work when using tensorrt
     if args.use_fp16 is True:
         assert args.use_tensorrt is True
@@ -118,7 +123,16 @@ def main():
     test_num = 500
     test_time = 0.0
     if not args.enable_benchmark:
-        inputs = preprocess(args.image_file, operators)
+        image_files = []
+        if args.image_file is not None:
+            image_files = [args.image_file]
+        elif args.image_dir is not None:
+            supported_exts = ('.jpg', 'jpeg', '.png', '.gif', '.bmp')
+            for root, _, files in os.walk(args.image_dir, topdown=False):
+                image_files += [os.path.join(root, f) for f in files
+                                if os.path.splitext(f)[-1].lower() in supported_exts]
+        for image_file in image_files:
+            inputs = preprocess(image_file, operators)
             inputs = np.expand_dims(
                 inputs, axis=0).repeat(
                     args.batch_size, axis=0).copy()
@@ -130,6 +144,7 @@ def main():
             output = output.flatten()
             cls = np.argmax(output)
             score = output[cls]
+            logger.info("image file: {0}".format(image_file))
             logger.info("class: {0}".format(cls))
             logger.info("score: {0}".format(score))
     else: