predictor.py 2.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485
  1. # !/usr/bin/env python3
  2. # -*- coding: UTF-8 -*-
  3. ################################################################################
  4. #
  5. # Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
  6. #
  7. ################################################################################
  8. """
  9. Author: PaddlePaddle Authors
  10. """
  11. import os
  12. import numpy as np
  13. from ....utils import logging
  14. from ...base.predictor.transforms import image_common
  15. from ...base import BasePredictor
  16. from .keys import TableRecKeys as K
  17. from . import transforms as T
  18. from ..model_list import MODELS
  19. class TableRecPredictor(BasePredictor):
  20. """ TableRecPredictor """
  21. entities = MODELS
  22. def __init__(self,
  23. model_dir,
  24. kernel_option,
  25. output,
  26. pre_transforms=None,
  27. post_transforms=None,
  28. table_max_len=488):
  29. super().__init__(
  30. model_dir=model_dir,
  31. kernel_option=kernel_option,
  32. output=output,
  33. pre_transforms=pre_transforms,
  34. post_transforms=post_transforms)
  35. self.table_max_len = table_max_len
  36. @classmethod
  37. def get_input_keys(cls):
  38. """ get input keys """
  39. return [[K.IMAGE, K.ORI_IM_SIZE], [K.IM_PATH, K.ORI_IM_SIZE]]
  40. @classmethod
  41. def get_output_keys(cls):
  42. """ get output keys """
  43. return [K.STRUCTURE_PROB, K.LOC_PROB, K.SHAPE_LIST]
  44. def _run(self, batch_input):
  45. """ run """
  46. images = [data[K.IMAGE] for data in batch_input]
  47. input_ = np.stack(images, axis=0)
  48. if input_.ndim == 3:
  49. input_ = input_[:, np.newaxis]
  50. input_ = input_.astype(dtype=np.float32, copy=False)
  51. outputs = self._predictor.predict([input_])
  52. struc_probs = outputs[1]
  53. bbox_probs = outputs[0]
  54. for data in batch_input:
  55. data[K.SHAPE_LIST] = [data[K.ORI_IM_SIZE]]
  56. # In-place update
  57. pred = batch_input
  58. for dict_, struc_prob, bbox_prob in zip(pred, struc_probs, bbox_probs):
  59. dict_[K.STRUCTURE_PROB] = struc_prob[np.newaxis, :]
  60. dict_[K.LOC_PROB] = bbox_prob[np.newaxis, :]
  61. return pred
  62. def _get_pre_transforms_from_config(self):
  63. """ _get_pre_transforms_from_config """
  64. return [
  65. image_common.ReadImage(),
  66. image_common.ResizeByLong(target_long_edge=self.table_max_len),
  67. image_common.Normalize(
  68. mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
  69. image_common.Pad(target_size=self.table_max_len, val=0.0),
  70. image_common.ToCHWImage()
  71. ]
  72. def _get_post_transforms_from_config(self):
  73. """ get postprocess transforms """
  74. return [T.TableLabelDecode(), T.SaveTableResults(self.output)]