predictor.py 2.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980
  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. from operator import le
  12. import os
  13. import numpy as np
  14. from . import transforms as T
  15. from ....utils import logging
  16. from ...base import BasePredictor
  17. from ...base.predictor.transforms import image_common
  18. from .keys import TextDetKeys as K
  19. from ..model_list import MODELS
  20. class TextDetPredictor(BasePredictor):
  21. """ TextDetPredictor """
  22. entities = MODELS
  23. @classmethod
  24. def get_input_keys(cls):
  25. """ get input keys """
  26. return [[K.IMAGE], [K.IM_PATH]]
  27. @classmethod
  28. def get_output_keys(cls):
  29. """ get output keys """
  30. return [K.PROB_MAP, K.SHAPE]
  31. def _run(self, batch_input):
  32. """ _run """
  33. if len(batch_input) != 1:
  34. raise ValueError(
  35. f"For `{self.__class__.__name__}`, batch size can only be set to 1."
  36. )
  37. images = [data[K.IMAGE] for data in batch_input]
  38. input_ = np.stack(images, axis=0)
  39. if input_.ndim == 3:
  40. input_ = input_[:, np.newaxis]
  41. input_ = input_.astype(dtype=np.float32, copy=False)
  42. outputs = self._predictor.predict([input_])
  43. pred = batch_input
  44. pred[0][K.PROB_MAP] = outputs
  45. return pred
  46. def _get_pre_transforms_from_config(self):
  47. """ get preprocess transforms """
  48. return [
  49. image_common.ReadImage(), T.DetResizeForTest(
  50. limit_side_len=960, limit_type="max"), T.NormalizeImage(
  51. mean=[0.485, 0.456, 0.406],
  52. std=[0.229, 0.224, 0.225],
  53. scale=1. / 255,
  54. order='hwc'), image_common.ToCHWImage()
  55. ]
  56. def _get_post_transforms_from_config(self):
  57. """ get postprocess transforms """
  58. post_transforms = [
  59. T.DBPostProcess(
  60. thresh=0.3,
  61. box_thresh=0.6,
  62. max_candidates=1000,
  63. unclip_ratio=1.5,
  64. use_dilation=False,
  65. score_mode='fast',
  66. box_type='quad'), T.SaveTextDetResults(self.output),
  67. T.PrintResult()
  68. ]
  69. return post_transforms