From 0c280c709a9a1c0b7e9454aed18e63274b398f7a Mon Sep 17 00:00:00 2001 From: andyjpaddle Date: Wed, 27 Apr 2022 05:19:55 +0000 Subject: [PATCH] fix key for dis and cls resize --- ppocr/data/imaug/rec_img_aug.py | 2 +- tools/eval.py | 6 ++++-- tools/program.py | 6 ++++-- 3 files changed, 9 insertions(+), 5 deletions(-) diff --git a/ppocr/data/imaug/rec_img_aug.py b/ppocr/data/imaug/rec_img_aug.py index 960a11be..fa3bc5ad 100644 --- a/ppocr/data/imaug/rec_img_aug.py +++ b/ppocr/data/imaug/rec_img_aug.py @@ -81,7 +81,7 @@ class ClsResizeImg(object): def __call__(self, data): img = data['image'] - norm_img = resize_norm_img(img, self.image_shape) + norm_img, _ = resize_norm_img(img, self.image_shape) data['image'] = norm_img return data diff --git a/tools/eval.py b/tools/eval.py index 1038090a..7fd4fa7a 100755 --- a/tools/eval.py +++ b/tools/eval.py @@ -74,9 +74,11 @@ def main(): model = build_model(config['Architecture']) extra_input_models = ["SRN", "NRTR", "SAR", "SEED", "SVTR"] + extra_input = False if config['Architecture']['algorithm'] == 'Distillation': - extra_input = config['Architecture']['Models']['Teacher'][ - 'algorithm'] in extra_input_models + for key in config['Architecture']["Models"]: + extra_input = extra_input or config['Architecture']['Models'][key][ + 'algorithm'] in extra_input_models else: extra_input = config['Architecture']['algorithm'] in extra_input_models if "model_type" in config['Architecture'].keys(): diff --git a/tools/program.py b/tools/program.py index 1742f6c9..90fd309a 100755 --- a/tools/program.py +++ b/tools/program.py @@ -202,9 +202,11 @@ def train(config, use_srn = config['Architecture']['algorithm'] == "SRN" extra_input_models = ["SRN", "NRTR", "SAR", "SEED", "SVTR"] + extra_input = False if config['Architecture']['algorithm'] == 'Distillation': - extra_input = config['Architecture']['Models']['Teacher'][ - 'algorithm'] in extra_input_models + for key in config['Architecture']["Models"]: + extra_input = extra_input or config['Architecture']['Models'][key][ + 'algorithm'] in extra_input_models else: extra_input = config['Architecture']['algorithm'] in extra_input_models try: -- GitLab