fix infer for dir
# 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
