layout_parsing.py 7.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205
  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. from typing import Final, List, Literal, Optional, Tuple
  16. import numpy as np
  17. from fastapi import FastAPI, HTTPException
  18. from numpy.typing import ArrayLike
  19. from pydantic import BaseModel, Field
  20. from typing_extensions import Annotated, TypeAlias
  21. from .....utils import logging
  22. from ...layout_parsing import LayoutParsingPipeline
  23. from ..storage import SupportsGetURL, Storage, create_storage
  24. from .. import utils as serving_utils
  25. from ..app import AppConfig, create_app
  26. from ..models import Response, ResultResponse
  27. _DEFAULT_MAX_IMG_SIZE: Final[Tuple[int, int]] = (2000, 2000)
  28. _DEFAULT_MAX_NUM_IMGS: Final[int] = 10
  29. FileType: TypeAlias = Literal[0, 1]
  30. class InferenceParams(BaseModel):
  31. maxLongSide: Optional[Annotated[int, Field(gt=0)]] = None
  32. class InferRequest(BaseModel):
  33. file: str
  34. fileType: Optional[FileType] = None
  35. useImgOrientationCls: bool = True
  36. useImgUnwrapping: bool = True
  37. useSealTextDet: bool = True
  38. inferenceParams: Optional[InferenceParams] = None
  39. BoundingBox: TypeAlias = Annotated[List[float], Field(min_length=4, max_length=4)]
  40. class LayoutElement(BaseModel):
  41. bbox: BoundingBox
  42. label: str
  43. text: str
  44. layoutType: Literal["single", "double"]
  45. image: Optional[str] = None
  46. class LayoutParsingResult(BaseModel):
  47. layoutElements: List[LayoutElement]
  48. class InferResult(BaseModel):
  49. layoutParsingResults: List[LayoutParsingResult]
  50. def _postprocess_image(
  51. img: ArrayLike,
  52. request_id: str,
  53. filename: str,
  54. file_storage: Optional[Storage],
  55. ) -> str:
  56. key = f"{request_id}/{filename}"
  57. ext = os.path.splitext(filename)[1]
  58. img = np.asarray(img)
  59. img_bytes = serving_utils.image_array_to_bytes(img, ext=ext)
  60. if file_storage is not None:
  61. file_storage.set(key, img_bytes)
  62. if isinstance(file_storage, SupportsGetURL):
  63. return file_storage.get_url(key)
  64. return serving_utils.base64_encode(img_bytes)
  65. def create_pipeline_app(
  66. pipeline: LayoutParsingPipeline, app_config: AppConfig
  67. ) -> FastAPI:
  68. app, ctx = create_app(
  69. pipeline=pipeline, app_config=app_config, app_aiohttp_session=True
  70. )
  71. if ctx.config.extra and "file_storage" in ctx.config.extra:
  72. ctx.extra["file_storage"] = create_storage(ctx.config.extra["file_storage"])
  73. else:
  74. ctx.extra["file_storage"] = None
  75. ctx.extra.setdefault("max_img_size", _DEFAULT_MAX_IMG_SIZE)
  76. ctx.extra.setdefault("max_num_imgs", _DEFAULT_MAX_NUM_IMGS)
  77. @app.post(
  78. "/layout-parsing",
  79. operation_id="infer",
  80. responses={422: {"model": Response}},
  81. response_model_exclude_none=True,
  82. )
  83. async def _infer(
  84. request: InferRequest,
  85. ) -> ResultResponse[InferResult]:
  86. pipeline = ctx.pipeline
  87. aiohttp_session = ctx.aiohttp_session
  88. request_id = serving_utils.generate_request_id()
  89. if request.fileType is None:
  90. if serving_utils.is_url(request.file):
  91. try:
  92. file_type = serving_utils.infer_file_type(request.file)
  93. except Exception as e:
  94. logging.exception(e)
  95. raise HTTPException(
  96. status_code=422,
  97. detail="The file type cannot be inferred from the URL. Please specify the file type explicitly.",
  98. )
  99. else:
  100. raise HTTPException(status_code=422, detail="Unknown file type")
  101. else:
  102. file_type = "PDF" if request.fileType == 0 else "IMAGE"
  103. if request.inferenceParams:
  104. max_long_side = request.inferenceParams.maxLongSide
  105. if max_long_side:
  106. raise HTTPException(
  107. status_code=422,
  108. detail="`max_long_side` is currently not supported.",
  109. )
  110. try:
  111. file_bytes = await serving_utils.get_raw_bytes(
  112. request.file, aiohttp_session
  113. )
  114. images = await serving_utils.call_async(
  115. serving_utils.file_to_images,
  116. file_bytes,
  117. file_type,
  118. max_img_size=ctx.extra["max_img_size"],
  119. max_num_imgs=ctx.extra["max_num_imgs"],
  120. )
  121. result = await pipeline.infer(
  122. images,
  123. use_doc_image_ori_cls_model=request.useImgOrientationCls,
  124. use_doc_image_unwarp_model=request.useImgUnwrapping,
  125. use_seal_text_det_model=request.useSealTextDet,
  126. )
  127. layout_parsing_results: List[LayoutParsingResult] = []
  128. for i, item in enumerate(result):
  129. layout_elements: List[LayoutElement] = []
  130. for j, subitem in enumerate(
  131. item["layout_parsing_result"]["parsing_result"]
  132. ):
  133. dyn_keys = subitem.keys() - {"input_path", "layout_bbox", "layout"}
  134. if len(dyn_keys) != 1:
  135. raise RuntimeError(f"Unexpected result: {subitem}")
  136. label = next(iter(dyn_keys))
  137. if label in ("image", "figure", "img", "fig"):
  138. image_ = await serving_utils.call_async(
  139. _postprocess_image,
  140. subitem[label]["img"],
  141. request_id=request_id,
  142. filename=f"image_{i}_{j}.jpg",
  143. file_storage=ctx.extra["file_storage"],
  144. )
  145. text = subitem[label]["image_text"]
  146. else:
  147. image_ = None
  148. text = subitem[label]
  149. layout_elements.append(
  150. LayoutElement(
  151. bbox=subitem["layout_bbox"],
  152. label=label,
  153. text=text,
  154. layoutType=subitem["layout"],
  155. image=image_,
  156. )
  157. )
  158. layout_parsing_results.append(
  159. LayoutParsingResult(layoutElements=layout_elements)
  160. )
  161. return ResultResponse(
  162. logId=serving_utils.generate_log_id(),
  163. errorCode=0,
  164. errorMsg="Success",
  165. result=InferResult(
  166. layoutParsingResults=layout_parsing_results,
  167. ),
  168. )
  169. except Exception as e:
  170. logging.exception(e)
  171. raise HTTPException(status_code=500, detail="Internal server error")
  172. return app