fix py infer
# 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)
+ 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("class id: {:d}, probability: {:.4f}".format(idx, prob))
+ print("\tclass id: {:d}, probability: {:.4f}".format(idx, prob))
if __name__ == "__main__":
