utils.py 6.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217
  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 asyncio
  15. import base64
  16. import io
  17. import os
  18. import re
  19. import uuid
  20. from functools import partial
  21. from typing import (
  22. Awaitable,
  23. Callable,
  24. List,
  25. Literal,
  26. Optional,
  27. TypeVar,
  28. Final,
  29. Tuple,
  30. overload,
  31. Union,
  32. )
  33. from urllib.parse import parse_qs, urlparse
  34. import aiohttp
  35. import cv2
  36. import fitz
  37. import numpy as np
  38. import pandas as pd
  39. import yarl
  40. from PIL import Image
  41. from typing_extensions import ParamSpec, assert_never
  42. from .models import ImageInfo, PDFInfo, PDFPageInfo
  43. FileType = Literal["IMAGE", "PDF"]
  44. _P = ParamSpec("_P")
  45. _R = TypeVar("_R")
  46. def generate_log_id() -> str:
  47. return str(uuid.uuid4())
  48. def is_url(s: str) -> bool:
  49. if not (s.startswith("http://") or s.startswith("https://")):
  50. # Quick rejection
  51. return False
  52. result = urlparse(s)
  53. return all([result.scheme, result.netloc]) and result.scheme in ("http", "https")
  54. def infer_file_type(url: str) -> FileType:
  55. # Is it more reliable to guess the file type based on the response headers?
  56. SUPPORTED_IMG_EXTS: Final[List[str]] = [".jpg", ".jpeg", ".png"]
  57. url_parts = urlparse(url)
  58. ext = os.path.splitext(url_parts.path)[1]
  59. # HACK: The support for BOS URLs with query params is implementation-based,
  60. # not interface-based.
  61. is_bos_url = (
  62. re.fullmatch(r"(?:bj|bd|su|gz|cd|hkg|fwh|fsh)\.bcebos\.com", url_parts.netloc)
  63. is not None
  64. )
  65. if is_bos_url and url_parts.query:
  66. params = parse_qs(url_parts.query)
  67. if (
  68. "responseContentDisposition" not in params
  69. or len(params["responseContentDisposition"]) != 1
  70. ):
  71. raise ValueError("`responseContentDisposition` not found")
  72. match_ = re.match(
  73. r"attachment;filename=(.*)", params["responseContentDisposition"][0]
  74. )
  75. if not match_ or not match_.groups()[0] is not None:
  76. raise ValueError(
  77. "Failed to extract the filename from `responseContentDisposition`"
  78. )
  79. ext = os.path.splitext(match_.groups()[0])[1]
  80. ext = ext.lower()
  81. if ext == ".pdf":
  82. return "PDF"
  83. elif ext in SUPPORTED_IMG_EXTS:
  84. return "IMAGE"
  85. else:
  86. raise ValueError("Unsupported file type")
  87. async def get_raw_bytes(file: str, session: aiohttp.ClientSession) -> bytes:
  88. if is_url(file):
  89. async with session.get(yarl.URL(file, encoded=True)) as resp:
  90. return await resp.read()
  91. else:
  92. return base64.b64decode(file)
  93. def image_bytes_to_array(data: bytes) -> np.ndarray:
  94. return cv2.imdecode(np.frombuffer(data, np.uint8), cv2.IMREAD_COLOR)
  95. def image_bytes_to_image(data: bytes) -> Image.Image:
  96. return Image.open(io.BytesIO(data))
  97. def image_to_bytes(image: Image.Image, format: str = "JPEG") -> bytes:
  98. with io.BytesIO() as f:
  99. image.save(f, format=format)
  100. img_bytes = f.getvalue()
  101. return img_bytes
  102. def image_array_to_bytes(image: np.ndarray, ext: str = ".jpg") -> bytes:
  103. image = cv2.imencode(ext, image)[1]
  104. return image.tobytes()
  105. def csv_bytes_to_data_frame(data: bytes) -> pd.DataFrame:
  106. with io.StringIO(data.decode("utf-8")) as f:
  107. df = pd.read_csv(f)
  108. return df
  109. def data_frame_to_bytes(df: str) -> str:
  110. return df.to_csv().encode("utf-8")
  111. def base64_encode(data: bytes) -> str:
  112. return base64.b64encode(data).decode("ascii")
  113. def read_pdf(
  114. bytes_: bytes, max_num_imgs: Optional[int] = None
  115. ) -> Tuple[List[np.ndarray], PDFInfo]:
  116. images: List[np.ndarray] = []
  117. page_info_list: List[PDFPageInfo] = []
  118. with fitz.open("pdf", bytes_) as doc:
  119. for page in doc:
  120. if max_num_imgs is not None and len(images) >= max_num_imgs:
  121. break
  122. # TODO: Do not always use zoom=2.0
  123. zoom = 2.0
  124. deg = 0
  125. mat = fitz.Matrix(zoom, zoom).prerotate(deg)
  126. pixmap = page.get_pixmap(matrix=mat, alpha=False)
  127. image = np.frombuffer(pixmap.samples, dtype=np.uint8).reshape(
  128. pixmap.h, pixmap.w, pixmap.n
  129. )
  130. image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
  131. images.append(image)
  132. page_info = PDFPageInfo(
  133. width=pixmap.w,
  134. height=pixmap.h,
  135. )
  136. page_info_list.append(page_info)
  137. pdf_info = PDFInfo(
  138. numPages=len(page_info_list),
  139. pages=page_info_list,
  140. )
  141. return images, pdf_info
  142. @overload
  143. def file_to_images(
  144. file_bytes: bytes,
  145. file_type: Literal["IMAGE"],
  146. *,
  147. max_num_imgs: Optional[int] = ...,
  148. ) -> Tuple[List[np.ndarray], ImageInfo]: ...
  149. @overload
  150. def file_to_images(
  151. file_bytes: bytes,
  152. file_type: Literal["PDF"],
  153. *,
  154. max_num_imgs: Optional[int] = ...,
  155. ) -> Tuple[List[np.ndarray], PDFInfo]: ...
  156. def file_to_images(
  157. file_bytes: bytes,
  158. file_type: Literal["IMAGE", "PDF"],
  159. *,
  160. max_num_imgs: Optional[int] = None,
  161. ) -> Tuple[List[np.ndarray], Union[ImageInfo, PDFInfo]]:
  162. if file_type == "IMAGE":
  163. images = [image_bytes_to_array(file_bytes)]
  164. data_info = get_image_info(images[0])
  165. elif file_type == "PDF":
  166. images, data_info = read_pdf(file_bytes, max_num_imgs=max_num_imgs)
  167. else:
  168. assert_never(file_type)
  169. return images, data_info
  170. def get_image_info(image: np.ndarray) -> ImageInfo:
  171. return ImageInfo(width=image.shape[1], height=image.shape[0])
  172. def call_async(
  173. func: Callable[_P, _R], /, *args: _P.args, **kwargs: _P.kwargs
  174. ) -> Awaitable[_R]:
  175. return asyncio.get_running_loop().run_in_executor(
  176. None, partial(func, *args, **kwargs)
  177. )