| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485 |
- # !/usr/bin/env python3
- # -*- coding: UTF-8 -*-
- ################################################################################
- #
- # Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
- #
- ################################################################################
- """
- Author: PaddlePaddle Authors
- """
- import os
- import numpy as np
- from ....utils import logging
- from ...base.predictor.transforms import image_common
- from ...base import BasePredictor
- from .keys import TableRecKeys as K
- from . import transforms as T
- from ..model_list import MODELS
- class TableRecPredictor(BasePredictor):
- """ TableRecPredictor """
- entities = MODELS
- def __init__(self,
- model_dir,
- kernel_option,
- output,
- pre_transforms=None,
- post_transforms=None,
- table_max_len=488):
- super().__init__(
- model_dir=model_dir,
- kernel_option=kernel_option,
- output=output,
- pre_transforms=pre_transforms,
- post_transforms=post_transforms)
- self.table_max_len = table_max_len
- @classmethod
- def get_input_keys(cls):
- """ get input keys """
- return [[K.IMAGE, K.ORI_IM_SIZE], [K.IM_PATH, K.ORI_IM_SIZE]]
- @classmethod
- def get_output_keys(cls):
- """ get output keys """
- return [K.STRUCTURE_PROB, K.LOC_PROB, K.SHAPE_LIST]
- def _run(self, batch_input):
- """ run """
- images = [data[K.IMAGE] for data in batch_input]
- input_ = np.stack(images, axis=0)
- if input_.ndim == 3:
- input_ = input_[:, np.newaxis]
- input_ = input_.astype(dtype=np.float32, copy=False)
- outputs = self._predictor.predict([input_])
- struc_probs = outputs[1]
- bbox_probs = outputs[0]
- for data in batch_input:
- data[K.SHAPE_LIST] = [data[K.ORI_IM_SIZE]]
- # In-place update
- pred = batch_input
- for dict_, struc_prob, bbox_prob in zip(pred, struc_probs, bbox_probs):
- dict_[K.STRUCTURE_PROB] = struc_prob[np.newaxis, :]
- dict_[K.LOC_PROB] = bbox_prob[np.newaxis, :]
- return pred
- def _get_pre_transforms_from_config(self):
- """ _get_pre_transforms_from_config """
- return [
- image_common.ReadImage(),
- image_common.ResizeByLong(target_long_edge=self.table_max_len),
- image_common.Normalize(
- mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
- image_common.Pad(target_size=self.table_max_len, val=0.0),
- image_common.ToCHWImage()
- ]
- def _get_post_transforms_from_config(self):
- """ get postprocess transforms """
- return [T.TableLabelDecode(), T.SaveTableResults(self.output)]
|