layout_parsing.py 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107
  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, HTTPException
  16. from .....utils import logging
  17. from ...infra import utils as serving_utils
  18. from ...infra.config import AppConfig
  19. from ...infra.models import ResultResponse
  20. from ...schemas.layout_parsing import INFER_ENDPOINT, InferRequest, InferResult
  21. from .._app import create_app, primary_operation
  22. from ._common import image as image_common
  23. from ._common import ocr as ocr_common
  24. def create_pipeline_app(pipeline: Any, app_config: AppConfig) -> FastAPI:
  25. app, ctx = create_app(
  26. pipeline=pipeline, app_config=app_config, app_aiohttp_session=True
  27. )
  28. ocr_common.update_app_context(ctx)
  29. @primary_operation(
  30. app,
  31. INFER_ENDPOINT,
  32. "infer",
  33. )
  34. async def _infer(
  35. request: InferRequest,
  36. ) -> ResultResponse[InferResult]:
  37. pipeline = ctx.pipeline
  38. log_id = serving_utils.generate_log_id()
  39. images, data_info = await ocr_common.get_images(request, ctx)
  40. result = await pipeline.infer(
  41. images,
  42. use_doc_image_ori_cls_model=request.useImgOrientationCls,
  43. use_doc_image_unwarp_model=request.useImgUnwarping,
  44. use_seal_text_det_model=request.useSealTextDet,
  45. )
  46. layout_parsing_results: List[Dict[str, Any]] = []
  47. for i, item in enumerate(result):
  48. layout_elements: List[Dict[str, Any]] = []
  49. for j, subitem in enumerate(
  50. item["layout_parsing_result"]["parsing_result"]
  51. ):
  52. dyn_keys = subitem.keys() - {"input_path", "layout_bbox", "layout"}
  53. if len(dyn_keys) != 1:
  54. logging.error("Unexpected result: %s", subitem)
  55. raise HTTPException(
  56. status_code=500,
  57. detail="Internal server error",
  58. )
  59. label = next(iter(dyn_keys))
  60. if label in ("image", "figure", "img", "fig"):
  61. text = subitem[label]["image_text"]
  62. if ctx.config.visualize:
  63. image = await serving_utils.call_async(
  64. image_common.postprocess_image,
  65. subitem[label]["img"],
  66. log_id=log_id,
  67. filename=f"image_{i}_{j}.jpg",
  68. file_storage=ctx.extra["file_storage"],
  69. return_url=ctx.extra["return_img_urls"],
  70. max_img_size=ctx.extra["max_output_img_size"],
  71. )
  72. else:
  73. image = None
  74. else:
  75. text = subitem[label]
  76. image = None
  77. layout_elements.append(
  78. dict(
  79. bbox=subitem["layout_bbox"],
  80. label=label,
  81. text=text,
  82. layoutType=subitem["layout"],
  83. image=image,
  84. )
  85. )
  86. layout_parsing_results.append(dict(layoutElements=layout_elements))
  87. return ResultResponse[InferResult](
  88. logId=log_id,
  89. result=InferResult(
  90. layoutParsingResults=layout_parsing_results,
  91. dataInfo=data_info,
  92. ),
  93. )
  94. return app