pipeline.py 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152
  1. # copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import os
  15. import cv2
  16. from ..base import BasePipeline
  17. from ...modules.text_detection.model_list import MODELS as text_det_models
  18. from ...modules.text_recognition.model_list import MODELS as text_rec_models
  19. from ...modules import create_model, PaddleInferenceOption
  20. from ...modules.text_detection import transforms as text_det_T
  21. from .utils import draw_ocr_box_txt
  22. class OCRPipeline(BasePipeline):
  23. """OCR Pipeline"""
  24. entities = "OCR"
  25. def __init__(
  26. self,
  27. text_det_model_name=None,
  28. text_rec_model_name=None,
  29. text_det_model_dir=None,
  30. text_rec_model_dir=None,
  31. text_det_kernel_option=None,
  32. text_rec_kernel_option=None,
  33. output="./",
  34. device="gpu",
  35. **kwargs,
  36. ):
  37. self.text_det_model_name = text_det_model_name
  38. self.text_rec_model_name = text_rec_model_name
  39. self.text_det_model_dir = text_det_model_dir
  40. self.text_rec_model_dir = text_rec_model_dir
  41. self.output = output
  42. self.device = device
  43. self.text_det_kernel_option = text_det_kernel_option
  44. self.text_rec_kernel_option = text_rec_kernel_option
  45. if (
  46. self.text_det_model_name is not None
  47. and self.text_rec_model_name is not None
  48. ):
  49. self.load_model()
  50. def check_model_name(self):
  51. """check that model name is valid"""
  52. assert (
  53. self.text_det_model_name in text_det_models
  54. ), f"The model name({self.text_det_model_name}) error. \
  55. Only support: {text_det_models}."
  56. assert (
  57. self.text_rec_model_name in text_rec_models
  58. ), f"The model name({self.text_rec_model_name}) error. \
  59. Only support: {text_rec_models}."
  60. def load_model(self):
  61. """load model predictor"""
  62. self.check_model_name()
  63. text_det_kernel_option = (
  64. self.get_kernel_option()
  65. if self.text_det_kernel_option is None
  66. else self.text_det_kernel_option
  67. )
  68. text_rec_kernel_option = (
  69. self.get_kernel_option()
  70. if self.text_rec_kernel_option is None
  71. else self.text_rec_kernel_option
  72. )
  73. text_det_post_transforms = [
  74. text_det_T.DBPostProcess(
  75. thresh=0.3,
  76. box_thresh=0.6,
  77. max_candidates=1000,
  78. unclip_ratio=1.5,
  79. use_dilation=False,
  80. score_mode="fast",
  81. box_type="quad",
  82. ),
  83. # TODO
  84. text_det_T.CropByPolys(det_box_type="foo"),
  85. ]
  86. self.text_det_model = create_model(
  87. self.text_det_model_name,
  88. self.text_det_model_dir,
  89. kernel_option=text_det_kernel_option,
  90. post_transforms=text_det_post_transforms,
  91. )
  92. self.text_rec_model = create_model(
  93. self.text_rec_model_name,
  94. self.text_rec_model_dir,
  95. kernel_option=text_rec_kernel_option,
  96. )
  97. def predict(self, input):
  98. """predict"""
  99. result = self.text_det_model.predict(input)
  100. all_rec_result = []
  101. for i, img in enumerate(result["sub_imgs"]):
  102. rec_result = self.text_rec_model.predict({"image": img})
  103. all_rec_result.append(rec_result["rec_text"][0])
  104. result["rec_text"] = all_rec_result
  105. if self.output is not None:
  106. draw_img = draw_ocr_box_txt(
  107. result["original_image"], result["dt_polys"], result["rec_text"]
  108. )
  109. fn = os.path.basename(result["input_path"])
  110. cv2.imwrite(
  111. os.path.join(self.output, fn),
  112. draw_img[:, :, ::-1],
  113. )
  114. return result
  115. def update_model(self, model_name_list, model_dir_list):
  116. """update model
  117. Args:
  118. model_name_list (list): list of model name.
  119. model_dir_list (list): list of model directory.
  120. """
  121. assert len(model_name_list) == 2
  122. self.text_det_model_name = model_name_list[0]
  123. self.text_rec_model_name = model_name_list[1]
  124. if model_dir_list:
  125. assert len(model_dir_list) == 2
  126. self.text_det_model_dir = model_dir_list[0]
  127. self.text_rec_model_dir = model_dir_list[1]
  128. def get_kernel_option(self):
  129. """get kernel option"""
  130. kernel_option = PaddleInferenceOption()
  131. kernel_option.set_device(self.device)
  132. return kernel_option
  133. def get_input_keys(self):
  134. """get dict keys of input argument input"""
  135. return self.text_det_model.get_input_keys()