fix vl

tools/export_model.py

@@ -78,6 +78,12 @@ def export_single_model(model,
                 shape=[None, 3, 64, 512], dtype="float32"),
         ]
         model = to_static(model, input_spec=other_shape)
+    elif arch_config["model_type"] == "sr":
+        other_shape = [
+            paddle.static.InputSpec(
+                shape=[None, 3, 16, 64], dtype="float32")
+        ]
+        model = to_static(model, input_spec=other_shape)
     elif arch_config["algorithm"] == "ViTSTR":
         other_shape = [
             paddle.static.InputSpec(
@@ -116,7 +122,7 @@ def export_single_model(model,
             paddle.static.InputSpec(
                 shape=[None, 3, 224, 224], dtype="int64"),  # image
         ]
-        if arch_config["algorithm"] == "LayoutLM":
+        if model.backbone.use_visual_backbone is False:
             input_spec.pop(4)
         model = to_static(model, input_spec=[input_spec])
     else:
@@ -195,6 +201,9 @@ def main():
         else:  # base rec model
             config["Architecture"]["Head"]["out_channels"] = char_num
 
+    # for sr algorithm
+    if config["Architecture"]["model_type"] == "sr":
+        config['Architecture']["Transform"]['infer_mode'] = True
     model = build_model(config["Architecture"])
     load_model(config, model, model_type=config['Architecture']["model_type"])
     model.eval()

tools/infer/predict_rec.py

@@ -342,7 +342,7 @@ class TextRecognizer(object):
         for beg_img_no in range(0, img_num, batch_num):
             end_img_no = min(img_num, beg_img_no + batch_num)
             norm_img_batch = []
-            imgC, imgH, imgW = self.rec_image_shape
+            imgC, imgH, imgW = self.rec_image_shape[:3]
             max_wh_ratio = imgW / imgH
             # max_wh_ratio = 0
             for ino in range(beg_img_no, end_img_no):

tools/infer/predict_sr.py

@@ -0,0 +1,155 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import sys
+from PIL import Image
+__dir__ = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, __dir__)
+sys.path.insert(0, os.path.abspath(os.path.join(__dir__, '../..')))
+
+os.environ["FLAGS_allocator_strategy"] = 'auto_growth'
+
+import cv2
+import numpy as np
+import math
+import time
+import traceback
+import paddle
+
+import tools.infer.utility as utility
+from ppocr.postprocess import build_post_process
+from ppocr.utils.logging import get_logger
+from ppocr.utils.utility import get_image_file_list, check_and_read_gif
+
+logger = get_logger()
+
+
+class TextSR(object):
+    def __init__(self, args):
+        self.sr_image_shape = [int(v) for v in args.sr_image_shape.split(",")]
+        self.sr_batch_num = args.sr_batch_num
+
+        self.predictor, self.input_tensor, self.output_tensors, self.config = \
+            utility.create_predictor(args, 'sr', logger)
+        self.benchmark = args.benchmark
+        if args.benchmark:
+            import auto_log
+            pid = os.getpid()
+            gpu_id = utility.get_infer_gpuid()
+            self.autolog = auto_log.AutoLogger(
+                model_name="sr",
+                model_precision=args.precision,
+                batch_size=args.sr_batch_num,
+                data_shape="dynamic",
+                save_path=None,  #args.save_log_path,
+                inference_config=self.config,
+                pids=pid,
+                process_name=None,
+                gpu_ids=gpu_id if args.use_gpu else None,
+                time_keys=[
+                    'preprocess_time', 'inference_time', 'postprocess_time'
+                ],
+                warmup=0,
+                logger=logger)
+
+    def resize_norm_img(self, img):
+        imgC, imgH, imgW = self.sr_image_shape
+        img = img.resize((imgW // 2, imgH // 2), Image.BICUBIC)
+        img_numpy = np.array(img).astype("float32")
+        img_numpy = img_numpy.transpose((2, 0, 1)) / 255
+        return img_numpy
+
+    def __call__(self, img_list):
+        img_num = len(img_list)
+        batch_num = self.sr_batch_num
+        st = time.time()
+        st = time.time()
+        all_result = [] * img_num
+        if self.benchmark:
+            self.autolog.times.start()
+        for beg_img_no in range(0, img_num, batch_num):
+            end_img_no = min(img_num, beg_img_no + batch_num)
+            norm_img_batch = []
+            imgC, imgH, imgW = self.sr_image_shape
+            for ino in range(beg_img_no, end_img_no):
+                norm_img = self.resize_norm_img(img_list[ino])
+                norm_img = norm_img[np.newaxis, :]
+                norm_img_batch.append(norm_img)
+
+            norm_img_batch = np.concatenate(norm_img_batch)
+            norm_img_batch = norm_img_batch.copy()
+            if self.benchmark:
+                self.autolog.times.stamp()
+            self.input_tensor.copy_from_cpu(norm_img_batch)
+            self.predictor.run()
+            outputs = []
+            for output_tensor in self.output_tensors:
+                output = output_tensor.copy_to_cpu()
+                outputs.append(output)
+            if len(outputs) != 1:
+                preds = outputs
+            else:
+                preds = outputs[0]
+            all_result.append(outputs)
+        if self.benchmark:
+            self.autolog.times.end(stamp=True)
+        return all_result, time.time() - st
+
+
+def main(args):
+    image_file_list = get_image_file_list(args.image_dir)
+    text_recognizer = TextSR(args)
+    valid_image_file_list = []
+    img_list = []
+
+    # warmup 2 times
+    if args.warmup:
+        img = np.random.uniform(0, 255, [16, 64, 3]).astype(np.uint8)
+        for i in range(2):
+            res = text_recognizer([img] * int(args.sr_batch_num))
+
+    for image_file in image_file_list:
+        img, flag = check_and_read_gif(image_file)
+        if not flag:
+            img = Image.open(image_file).convert("RGB")
+        if img is None:
+            logger.info("error in loading image:{}".format(image_file))
+            continue
+        valid_image_file_list.append(image_file)
+        img_list.append(img)
+    try:
+        preds, _ = text_recognizer(img_list)
+        for beg_no in range(len(preds)):
+            sr_img = preds[beg_no][1]
+            lr_img = preds[beg_no][0]
+            for i in (range(sr_img.shape[0])):
+                fm_sr = (sr_img[i] * 255).transpose(1, 2, 0).astype(np.uint8)
+                fm_lr = (lr_img[i] * 255).transpose(1, 2, 0).astype(np.uint8)
+                img_name_pure = os.path.split(valid_image_file_list[
+                    beg_no * args.sr_batch_num + i])[-1]
+                cv2.imwrite("infer_result/sr_{}".format(img_name_pure),
+                            fm_sr[:, :, ::-1])
+                logger.info("The visualized image saved in infer_result/sr_{}".
+                            format(img_name_pure))
+
+    except Exception as E:
+        logger.info(traceback.format_exc())
+        logger.info(E)
+        exit()
+    if args.benchmark:
+        text_recognizer.autolog.report()
+
+
+if __name__ == "__main__":
+    main(utility.parse_args())

tools/infer/utility.py

@@ -121,6 +121,11 @@ def init_args():
     parser.add_argument("--use_pdserving", type=str2bool, default=False)
     parser.add_argument("--warmup", type=str2bool, default=False)
 
+    # SR parmas
+    parser.add_argument("--sr_model_dir", type=str)
+    parser.add_argument("--sr_image_shape", type=str, default="3, 32, 128")
+    parser.add_argument("--sr_batch_num", type=int, default=1)
+
     #
     parser.add_argument(
         "--draw_img_save_dir", type=str, default="./inference_results")
@@ -156,6 +161,8 @@ def create_predictor(args, mode, logger):
         model_dir = args.table_model_dir
     elif mode == 'ser':
         model_dir = args.ser_model_dir
+    elif mode == "sr":
+        model_dir = args.sr_model_dir
     else:
         model_dir = args.e2e_model_dir
 
@@ -205,18 +212,24 @@ def create_predictor(args, mode, logger):
                     workspace_size=1 << 30,
                     precision_mode=precision,
                     max_batch_size=args.max_batch_size,
-                    min_subgraph_size=args.min_subgraph_size, # skip the minmum trt subgraph
+                    min_subgraph_size=args.
+                    min_subgraph_size,  # skip the minmum trt subgraph
                     use_calib_mode=False)
-            
+
             # collect shape
             if args.shape_info_filename is not None:
                 if not os.path.exists(args.shape_info_filename):
                     config.collect_shape_range_info(args.shape_info_filename)
-                    logger.info(f"collect dynamic shape info into : {args.shape_info_filename}")
+                    logger.info(
+                        f"collect dynamic shape info into : {args.shape_info_filename}"
+                    )
                 else:
-                    logger.info(f"dynamic shape info file( {args.shape_info_filename} ) already exists, not need to generate again.")
-                config.enable_tuned_tensorrt_dynamic_shape(args.shape_info_filename, True)
-            
+                    logger.info(
+                        f"dynamic shape info file( {args.shape_info_filename} ) already exists, not need to generate again."
+                    )
+                config.enable_tuned_tensorrt_dynamic_shape(
+                    args.shape_info_filename, True)
+
             use_dynamic_shape = True
             if mode == "det":
                 min_input_shape = {

tools/infer_sr.py

@@ -0,0 +1,100 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+import os
+import sys
+import json
+from PIL import Image
+import cv2
+
+__dir__ = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, __dir__)
+sys.path.insert(0, os.path.abspath(os.path.join(__dir__, '..')))
+
+os.environ["FLAGS_allocator_strategy"] = 'auto_growth'
+
+import paddle
+
+from ppocr.data import create_operators, transform
+from ppocr.modeling.architectures import build_model
+from ppocr.postprocess import build_post_process
+from ppocr.utils.save_load import load_model
+from ppocr.utils.utility import get_image_file_list
+import tools.program as program
+
+
+def main():
+    global_config = config['Global']
+
+    # build post process
+    post_process_class = build_post_process(config['PostProcess'],
+                                            global_config)
+
+    # sr transform
+    config['Architecture']["Transform"]['infer_mode'] = True
+
+    model = build_model(config['Architecture'])
+
+    load_model(config, model)
+
+    # create data ops
+    transforms = []
+    for op in config['Eval']['dataset']['transforms']:
+        op_name = list(op)[0]
+        if 'Label' in op_name:
+            continue
+        elif op_name in ['SRResize']:
+            op[op_name]['infer_mode'] = True
+        elif op_name == 'KeepKeys':
+            op[op_name]['keep_keys'] = ['imge_lr']
+        transforms.append(op)
+    global_config['infer_mode'] = True
+    ops = create_operators(transforms, global_config)
+
+    save_res_path = config['Global'].get('save_res_path', "./infer_result")
+    if not os.path.exists(os.path.dirname(save_res_path)):
+        os.makedirs(os.path.dirname(save_res_path))
+
+    model.eval()
+    for file in get_image_file_list(config['Global']['infer_img']):
+        logger.info("infer_img: {}".format(file))
+        img = Image.open(file).convert("RGB")
+        data = {'image_lr': img}
+        batch = transform(data, ops)
+        images = np.expand_dims(batch[0], axis=0)
+        images = paddle.to_tensor(images)
+
+        preds = model(images)
+        sr_img = preds["sr_img"][0]
+        lr_img = preds["lr_img"][0]
+        fm_sr = (sr_img.numpy() * 255).transpose(1, 2, 0).astype(np.uint8)
+        fm_lr = (lr_img.numpy() * 255).transpose(1, 2, 0).astype(np.uint8)
+        img_name_pure = os.path.split(file)[-1]
+        cv2.imwrite("infer_result/sr_{}".format(img_name_pure),
+                    fm_sr[:, :, ::-1])
+        logger.info("The visualized image saved in infer_result/sr_{}".format(
+            img_name_pure))
+
+    logger.info("success!")
+
+
+if __name__ == '__main__':
+    config, device, logger, vdl_writer = program.preprocess()
+    main()

tools/infer_vqa_token_ser.py

@@ -104,8 +104,6 @@ class SerPredictor(object):
         batch = transform(data, self.ops)
         batch = to_tensor(batch)
         preds = self.model(batch)
-        if self.algorithm in ['LayoutLMv2', 'LayoutXLM']:
-            preds = preds[0]
 
         post_result = self.post_process_class(
             preds, segment_offset_ids=batch[6], ocr_infos=batch[7])

tools/program.py

@@ -25,6 +25,8 @@ import datetime
 import paddle
 import paddle.distributed as dist
 from tqdm import tqdm
+import cv2
+import numpy as np
 from argparse import ArgumentParser, RawDescriptionHelpFormatter
 
 from ppocr.utils.stats import TrainingStats
@@ -262,6 +264,7 @@ def train(config,
                 config, 'Train', device, logger, seed=epoch)
             max_iter = len(train_dataloader) - 1 if platform.system(
             ) == "Windows" else len(train_dataloader)
+
         for idx, batch in enumerate(train_dataloader):
             profiler.add_profiler_step(profiler_options)
             train_reader_cost += time.time() - reader_start
@@ -289,7 +292,7 @@ def train(config,
             else:
                 if model_type == 'table' or extra_input:
                     preds = model(images, data=batch[1:])
-                elif model_type in ["kie", 'vqa']:
+                elif model_type in ["kie", 'vqa', 'sr']:
                     preds = model(batch)
                 else:
                     preds = model(images)
@@ -297,11 +300,12 @@ def train(config,
                 avg_loss = loss['loss']
                 avg_loss.backward()
                 optimizer.step()
+
             optimizer.clear_grad()
 
             if cal_metric_during_train and epoch % calc_epoch_interval == 0:  # only rec and cls need
                 batch = [item.numpy() for item in batch]
-                if model_type in ['kie']:
+                if model_type in ['kie', 'sr']:
                     eval_class(preds, batch)
                 elif model_type in ['table']:
                     post_result = post_process_class(preds, batch)
@@ -347,8 +351,8 @@ def train(config,
                     len(train_dataloader) - idx - 1) * eta_meter.avg
                 eta_sec_format = str(datetime.timedelta(seconds=int(eta_sec)))
                 strs = 'epoch: [{}/{}], global_step: {}, {}, avg_reader_cost: ' \
-                       '{:.5f} s, avg_batch_cost: {:.5f} s, avg_samples: {}, ' \
-                       'ips: {:.5f} samples/s, eta: {}'.format(
+                    '{:.5f} s, avg_batch_cost: {:.5f} s, avg_samples: {}, ' \
+                    'ips: {:.5f} samples/s, eta: {}'.format(
                     epoch, epoch_num, global_step, logs,
                     train_reader_cost / print_batch_step,
                     train_batch_cost / print_batch_step,
@@ -376,7 +380,8 @@ def train(config,
                     post_process_class,
                     eval_class,
                     model_type,
-                    extra_input=extra_input)
+                    extra_input=extra_input,
+                    scaler=scaler)
                 cur_metric_str = 'cur metric, {}'.format(', '.join(
                     ['{}: {}'.format(k, v) for k, v in cur_metric.items()]))
                 logger.info(cur_metric_str)
@@ -466,7 +471,8 @@ def eval(model,
          post_process_class,
          eval_class,
          model_type=None,
-         extra_input=False):
+         extra_input=False,
+         scaler=None):
     model.eval()
     with paddle.no_grad():
         total_frame = 0.0
@@ -478,17 +484,58 @@ def eval(model,
             leave=True)
         max_iter = len(valid_dataloader) - 1 if platform.system(
         ) == "Windows" else len(valid_dataloader)
+        sum_images = 0
         for idx, batch in enumerate(valid_dataloader):
             if idx >= max_iter:
                 break
             images = batch[0]
             start = time.time()
-            if model_type == 'table' or extra_input:
-                preds = model(images, data=batch[1:])
-            elif model_type in ["kie", 'vqa']:
-                preds = model(batch)
+
+            # use amp
+            if scaler:
+                with paddle.amp.auto_cast(level='O2'):
+                    if model_type == 'table' or extra_input:
+                        preds = model(images, data=batch[1:])
+                    elif model_type in ["kie", 'vqa']:
+                        preds = model(batch)
+                    elif model_type in ['sr']:
+                        preds = model(batch)
+                        sr_img = preds["sr_img"]
+                        lr_img = preds["lr_img"]
+
+                        for i in (range(sr_img.shape[0])):
+                            fm_sr = (sr_img[i].numpy() * 255).transpose(
+                                1, 2, 0).astype(np.uint8)
+                            fm_lr = (lr_img[i].numpy() * 255).transpose(
+                                1, 2, 0).astype(np.uint8)
+                            cv2.imwrite("output/images/{}_{}_sr.jpg".format(
+                                sum_images, i), fm_sr)
+                            cv2.imwrite("output/images/{}_{}_lr.jpg".format(
+                                sum_images, i), fm_lr)
+                    else:
+                        preds = model(images)
             else:
-                preds = model(images)
+                if model_type == 'table' or extra_input:
+                    preds = model(images, data=batch[1:])
+                elif model_type in ["kie", 'vqa']:
+                    preds = model(batch)
+                elif model_type in ['sr']:
+                    preds = model(batch)
+                    sr_img = preds["sr_img"]
+                    lr_img = preds["lr_img"]
+
+                    for i in (range(sr_img.shape[0])):
+                        fm_sr = (sr_img[i].numpy() * 255).transpose(
+                            1, 2, 0).astype(np.uint8)
+                        fm_lr = (lr_img[i].numpy() * 255).transpose(
+                            1, 2, 0).astype(np.uint8)
+                        cv2.imwrite("output/images/{}_{}_sr.jpg".format(
+                            sum_images, i), fm_sr)
+                        cv2.imwrite("output/images/{}_{}_lr.jpg".format(
+                            sum_images, i), fm_lr)
+                else:
+                    preds = model(images)
+
             batch_numpy = []
             for item in batch:
                 if isinstance(item, paddle.Tensor):
@@ -503,12 +550,15 @@ def eval(model,
             elif model_type in ['table', 'vqa']:
                 post_result = post_process_class(preds, batch_numpy)
                 eval_class(post_result, batch_numpy)
+            elif model_type in ['sr']:
+                eval_class(preds, batch_numpy)
             else:
                 post_result = post_process_class(preds, batch_numpy[1])
                 eval_class(post_result, batch_numpy)
 
             pbar.update(1)
             total_frame += len(images)
+            sum_images += 1
         # Get final metric，eg. acc or hmean
         metric = eval_class.get_metric()
 
@@ -602,7 +652,8 @@ def preprocess(is_train=False):
         'EAST', 'DB', 'SAST', 'Rosetta', 'CRNN', 'STARNet', 'RARE', 'SRN',
         'CLS', 'PGNet', 'Distillation', 'NRTR', 'TableAttn', 'SAR', 'PSE',
         'SEED', 'SDMGR', 'LayoutXLM', 'LayoutLM', 'LayoutLMv2', 'PREN', 'FCE',
-        'SVTR', 'ViTSTR', 'ABINet', 'DB++', 'TableMaster', 'SPIN', 'VisionLAN'
+        'SVTR', 'ViTSTR', 'ABINet', 'DB++', 'TableMaster', 'SPIN', 'VisionLAN',
+        'Gestalt'
     ]
 
     if use_xpu:

tools/train.py

@@ -119,6 +119,7 @@ def main(config, device, logger, vdl_writer):
             config['Loss']['ignore_index'] = char_num - 1
 
     model = build_model(config['Architecture'])
+
     model = apply_to_static(model, config, logger)
 
     # build loss