fix infer for dir

tools/infer/infer.py

@@ -12,13 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import utils
-import argparse
-import numpy as np
-
-import paddle.fluid as fluid
-from ppcls.modeling import architectures
 from ppcls.utils.save_load import load_dygraph_pretrain
+from ppcls.modeling import architectures
+import paddle.fluid as fluid
+import numpy as np
+import argparse
+import utils
+import os
+import sys
+__dir__ = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(__dir__)
+sys.path.append(os.path.abspath(os.path.join(__dir__, '../..')))
 
 
 def parse_args():
@@ -66,6 +70,23 @@ 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()
@@ -78,11 +99,13 @@ def main():
 
     with fluid.dygraph.guard(place):
         net = architectures.__dict__[args.model]()
-        data = preprocess(args.image_file, operators)
-        data = np.expand_dims(data, axis=0)
-        data = fluid.dygraph.to_variable(data)
         load_dygraph_pretrain(net, args.pretrained_model,
                               args.load_static_weights)
+        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)
+            data = fluid.dygraph.to_variable(data)
             net.eval()
             outputs = net(data)
             outputs = fluid.layers.softmax(outputs)
@@ -90,9 +113,10 @@ def main():
 
             probs = postprocess(outputs)
             rank = 1
+            print("Current image file: {}".format(filename))
             for idx, prob in probs:
-        print("top{:d}, class id: {:d}, probability: {:.4f}".format(rank, idx,
-                                                                    prob))
+                print("\ttop{:d}, class id: {:d}, probability: {:.4f}".format(
+                    rank, idx, prob))
                 rank += 1
     return