seal_recognition.py 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104
  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 numpy as np
  15. from .utils import *
  16. from ..base import BasePipeline
  17. from ..ocr import OCRPipeline
  18. from ....utils import logging
  19. from ...components import CropByBoxes
  20. from ...results import SealResult
  21. class SealTextRecPipeline(BasePipeline):
  22. """Seal Recognition Pipeline"""
  23. entities = "seal_recognition"
  24. def __init__(
  25. self,
  26. layout_model,
  27. text_det_model,
  28. text_rec_model,
  29. layout_batch_size=1,
  30. text_det_batch_size=1,
  31. text_rec_batch_size=1,
  32. predictor_kwargs=None,
  33. ):
  34. self.layout_model = layout_model
  35. self.text_det_model = text_det_model
  36. self.text_rec_model = text_rec_model
  37. self.layout_batch_size = layout_batch_size
  38. self.text_det_batch_size = text_det_batch_size
  39. self.text_rec_batch_size = text_rec_batch_size
  40. self.predictor_kwargs = predictor_kwargs
  41. super().__init__(predictor_kwargs=predictor_kwargs)
  42. self._build_predictor()
  43. def _build_predictor(
  44. self,
  45. ):
  46. self.layout_predictor = self._create_model(model=self.layout_model)
  47. self.ocr_pipeline = OCRPipeline(
  48. text_det_model=self.text_det_model,
  49. text_rec_model=self.text_rec_model,
  50. text_det_batch_size=self.text_det_batch_size,
  51. text_rec_batch_size=self.text_rec_batch_size,
  52. predictor_kwargs=self.predictor_kwargs,
  53. )
  54. self._crop_by_boxes = CropByBoxes()
  55. self.layout_predictor.set_predictor(batch_size=self.layout_batch_size)
  56. self.ocr_pipeline.text_rec_model.set_predictor(
  57. batch_size=self.text_rec_batch_size
  58. )
  59. def set_predictor(
  60. self,
  61. layout_batch_size=None,
  62. text_det_batch_size=None,
  63. text_rec_batch_size=None,
  64. ):
  65. if text_det_batch_size and text_det_batch_size > 1:
  66. logging.warning(
  67. f"text det model only support batch_size=1 now,the setting of text_det_batch_size={text_det_batch_size} will not using! "
  68. )
  69. if layout_batch_size:
  70. self.layout_predictor.set_predictor(batch_size=layout_batch_size)
  71. if text_rec_batch_size:
  72. self.ocr_pipeline.text_rec_model.set_predictor(
  73. batch_size=text_rec_batch_size
  74. )
  75. def predict(self, x):
  76. for layout_pred in self.layout_predictor(x):
  77. single_img_res = {
  78. "input_path": "",
  79. "layout_result": {},
  80. "ocr_result": {},
  81. }
  82. # update layout result
  83. single_img_res["input_path"] = layout_pred["input_path"]
  84. single_img_res["layout_result"] = layout_pred
  85. seal_subs = []
  86. if len(layout_pred["boxes"]) > 0:
  87. subs_of_img = list(self._crop_by_boxes(layout_pred))
  88. # get cropped images with label "seal"
  89. for sub in subs_of_img:
  90. box = sub["box"]
  91. if sub["label"].lower() == "seal":
  92. seal_subs.append(sub)
  93. all_seal_ocr_res = get_ocr_res(self.ocr_pipeline, seal_subs)
  94. single_img_res["ocr_result"] = all_seal_ocr_res
  95. yield SealResult(single_img_res)