From 1bf24676c25730fb620423e2c6c23f49c041a104 Mon Sep 17 00:00:00 2001 From: LDOUBLEV Date: Mon, 23 Aug 2021 09:13:22 +0000 Subject: [PATCH] fix rec trt --- tools/infer/predict_rec.py | 2 +- tools/infer/utility.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tools/infer/predict_rec.py b/tools/infer/predict_rec.py index bb4a3170..5a9d4fda 100755 --- a/tools/infer/predict_rec.py +++ b/tools/infer/predict_rec.py @@ -277,7 +277,7 @@ def main(args): if args.warmup: img = np.random.uniform(0, 255, [32, 320, 3]).astype(np.uint8) for i in range(2): - res = text_recognizer([img]) + res = text_recognizer([img] * int(args.rec_batch_num)) for image_file in image_file_list: img, flag = check_and_read_gif(image_file) diff --git a/tools/infer/utility.py b/tools/infer/utility.py index 93b54505..1c822800 100755 --- a/tools/infer/utility.py +++ b/tools/infer/utility.py @@ -231,13 +231,13 @@ def create_predictor(args, mode, logger): max_input_shape.update(max_pact_shape) opt_input_shape.update(opt_pact_shape) elif mode == "rec": - min_input_shape = {"x": [1, 3, 32, 10]} + min_input_shape = {"x": [args.rec_batch_num, 3, 32, 10]} max_input_shape = {"x": [args.rec_batch_num, 3, 32, 2000]} - opt_input_shape = {"x": [512, 3, 32, 320]} + opt_input_shape = {"x": [args.rec_batch_num, 3, 32, 320]} elif mode == "cls": - min_input_shape = {"x": [1, 3, 48, 10]} + min_input_shape = {"x": [args.rec_batch_num, 3, 48, 10]} max_input_shape = {"x": [args.rec_batch_num, 3, 48, 2000]} - opt_input_shape = {"x": [512, 3, 48, 320]} + opt_input_shape = {"x": [args.rec_batch_num, 3, 48, 320]} else: min_input_shape = {"x": [1, 3, 10, 10]} max_input_shape = {"x": [1, 3, 1000, 1000]} -- GitLab