layout_parsing.py 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113
  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. from typing import Any, Dict, List
  15. from fastapi import FastAPI
  16. from ...infra import utils as serving_utils
  17. from ...infra.config import AppConfig
  18. from ...infra.models import ResultResponse
  19. from ...schemas.layout_parsing import INFER_ENDPOINT, InferRequest, InferResult
  20. from .._app import create_app, primary_operation
  21. from ._common import common
  22. from ._common import ocr as ocr_common
  23. def create_pipeline_app(pipeline: Any, app_config: AppConfig) -> FastAPI:
  24. app, ctx = create_app(
  25. pipeline=pipeline, app_config=app_config, app_aiohttp_session=True
  26. )
  27. ocr_common.update_app_context(ctx)
  28. @primary_operation(
  29. app,
  30. INFER_ENDPOINT,
  31. "infer",
  32. )
  33. async def _infer(
  34. request: InferRequest,
  35. ) -> ResultResponse[InferResult]:
  36. pipeline = ctx.pipeline
  37. log_id = serving_utils.generate_log_id()
  38. images, data_info = await ocr_common.get_images(request, ctx)
  39. result = await pipeline.infer(
  40. images,
  41. use_doc_orientation_classify=request.useDocOrientationClassify,
  42. use_doc_unwarping=request.useDocUnwarping,
  43. use_general_ocr=request.useGeneralOcr,
  44. use_seal_recognition=request.useSealRecognition,
  45. use_table_recognition=request.useTableRecognition,
  46. use_formula_recognition=request.useFormulaRecognition,
  47. text_det_limit_side_len=request.textDetLimitSideLen,
  48. text_det_limit_type=request.textDetLimitType,
  49. text_det_thresh=request.textDetThresh,
  50. text_det_box_thresh=request.textDetBoxThresh,
  51. text_det_unclip_ratio=request.textDetUnclipRatio,
  52. text_rec_score_thresh=request.textRecScoreThresh,
  53. seal_det_limit_side_len=request.sealDetLimitSideLen,
  54. seal_det_limit_type=request.sealDetLimitType,
  55. seal_det_thresh=request.sealDetThresh,
  56. seal_det_box_thresh=request.sealDetBoxThresh,
  57. seal_det_unclip_ratio=request.sealDetUnclipRatio,
  58. seal_rec_score_thresh=request.sealRecScoreThresh,
  59. layout_nms=request.layoutNms,
  60. layout_unclip_ratio=request.layoutUnclipRatio,
  61. layout_merge_bboxes_mode=request.layoutMergeBboxesMode,
  62. )
  63. layout_parsing_results: List[Dict[str, Any]] = []
  64. for i, (img, item) in enumerate(zip(images, result)):
  65. pruned_res = common.prune_result(item.json["res"])
  66. if ctx.config.visualize:
  67. imgs = {
  68. "input_img": img,
  69. **item.img,
  70. }
  71. imgs = await serving_utils.call_async(
  72. common.postprocess_images,
  73. imgs,
  74. log_id,
  75. filename_template=f"{{key}}_{i}.jpg",
  76. file_storage=ctx.extra["file_storage"],
  77. return_urls=ctx.extra["return_img_urls"],
  78. max_img_size=ctx.extra["max_output_img_size"],
  79. )
  80. else:
  81. imgs = {}
  82. layout_parsing_results.append(
  83. dict(
  84. prunedResult=pruned_res,
  85. outputImages=(
  86. {k: v for k, v in imgs.items() if k != "input_img"}
  87. if imgs
  88. else None
  89. ),
  90. inputImage=imgs.get("input_img"),
  91. )
  92. )
  93. return ResultResponse[InferResult](
  94. logId=log_id,
  95. result=InferResult(
  96. layoutParsingResults=layout_parsing_results,
  97. dataInfo=data_info,
  98. ),
  99. )
  100. return app