fix py infer

tools/infer/py_infer.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 import utils
 import argparse
 import numpy as np
@@ -26,8 +27,6 @@ def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("-i", "--image_file", type=str)
     parser.add_argument("-d", "--model_dir", type=str)
-    parser.add_argument("-m", "--model_file", type=str)
-    parser.add_argument("-p", "--params_file", type=str)
     parser.add_argument("--use_gpu", type=str2bool, default=True)
 
     return parser.parse_args()
@@ -41,10 +40,7 @@ def create_predictor(args):
 
     exe = fluid.Executor(place)
     [program, feed_names, fetch_lists] = fluid.io.load_inference_model(
-        args.model_dir,
-        exe,
-        model_filename=args.model_file,
-        params_filename=args.params_file)
+        args.model_dir, exe, model_filename="model", params_filename="params")
     compiled_program = fluid.compiler.CompiledProgram(program)
 
     return exe, compiled_program, feed_names, fetch_lists
@@ -70,7 +66,6 @@ def preprocess(fname, ops):
     data = open(fname, 'rb').read()
     for op in ops:
         data = op(data)
-
     return data
 
 
@@ -81,21 +76,40 @@ def postprocess(outputs, topk=5):
     return zip(index, prob[index])
 
 
+def get_image_list(img_file):
+    imgs_lists = []
+    if img_file is None or not os.path.exists(img_file):
+        raise Exception("not found any img file in {}".format(img_file))
+
+    img_end = ['jpg', 'png', 'jpeg', 'JPEG', 'JPG', 'bmp']
+    if os.path.isfile(img_file) and img_file.split('.')[-1] in img_end:
+        imgs_lists.append(img_file)
+    elif os.path.isdir(img_file):
+        for single_file in os.listdir(img_file):
+            if single_file.split('.')[-1] in img_end:
+                imgs_lists.append(os.path.join(img_file, single_file))
+    if len(imgs_lists) == 0:
+        raise Exception("not found any img file in {}".format(img_file))
+    return imgs_lists
+
+
 def main():
     args = parse_args()
     operators = create_operators()
     exe, program, feed_names, fetch_lists = create_predictor(args)
 
-    data = preprocess(args.image_file, operators)
-    data = np.expand_dims(data, axis=0)
-    outputs = exe.run(program,
-                      feed={feed_names[0]: data},
-                      fetch_list=fetch_lists,
-                      return_numpy=False)
-    probs = postprocess(outputs)
-
-    for idx, prob in probs:
-        print("class id: {:d}, probability: {:.4f}".format(idx, prob))
+    image_list = get_image_list(args.image_file)
+    for idx, filename in enumerate(image_list):
+        data = preprocess(filename, operators)
+        data = np.expand_dims(data, axis=0)
+        outputs = exe.run(program,
+                          feed={feed_names[0]: data},
+                          fetch_list=fetch_lists,
+                          return_numpy=False)
+        probs = postprocess(outputs)
+        print("Current image file: {}".format(filename))
+        for idx, prob in probs:
+            print("\tclass id: {:d}, probability: {:.4f}".format(idx, prob))
 
 
 if __name__ == "__main__":