| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177 |
- # copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
- #
- # Licensed under the Apache License, Version 2.0 (the "License");
- # you may not use this file except in compliance with the License.
- # You may obtain a copy of the License at
- #
- # http://www.apache.org/licenses/LICENSE-2.0
- #
- # Unless required by applicable law or agreed to in writing, software
- # distributed under the License is distributed on an "AS IS" BASIS,
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- # See the License for the specific language governing permissions and
- # limitations under the License.
- import asyncio
- import base64
- import io
- import os
- import re
- import uuid
- from functools import partial
- from typing import Awaitable, Callable, List, Literal, Optional, TypeVar, Final, Tuple
- from urllib.parse import parse_qs, urlparse
- import aiohttp
- import cv2
- import fitz
- import numpy as np
- import pandas as pd
- import yarl
- from PIL import Image
- from typing_extensions import ParamSpec, assert_never
- FileType = Literal["IMAGE", "PDF"]
- _P = ParamSpec("_P")
- _R = TypeVar("_R")
- def generate_log_id() -> str:
- return str(uuid.uuid4())
- def generate_request_id() -> str:
- return str(uuid.uuid4())
- def is_url(s: str) -> bool:
- if not (s.startswith("http://") or s.startswith("https://")):
- # Quick rejection
- return False
- result = urlparse(s)
- return all([result.scheme, result.netloc]) and result.scheme in ("http", "https")
- def infer_file_type(url: str) -> FileType:
- # Is it more reliable to guess the file type based on the response headers?
- SUPPORTED_IMG_EXTS: Final[List[str]] = [".jpg", ".jpeg", ".png"]
- url_parts = urlparse(url)
- ext = os.path.splitext(url_parts.path)[1]
- # HACK: The support for BOS URLs with query params is implementation-based,
- # not interface-based.
- is_bos_url = (
- re.fullmatch(r"(?:bj|bd|su|gz|cd|hkg|fwh|fsh)\.bcebos\.com", url_parts.netloc)
- is not None
- )
- if is_bos_url and url_parts.query:
- params = parse_qs(url_parts.query)
- if (
- "responseContentDisposition" not in params
- or len(params["responseContentDisposition"]) != 1
- ):
- raise ValueError("`responseContentDisposition` not found")
- match_ = re.match(
- r"attachment;filename=(.*)", params["responseContentDisposition"][0]
- )
- if not match_ or not match_.groups()[0] is not None:
- raise ValueError(
- "Failed to extract the filename from `responseContentDisposition`"
- )
- ext = os.path.splitext(match_.groups()[0])[1]
- ext = ext.lower()
- if ext == ".pdf":
- return "PDF"
- elif ext in SUPPORTED_IMG_EXTS:
- return "IMAGE"
- else:
- raise ValueError("Unsupported file type")
- async def get_raw_bytes(file: str, session: aiohttp.ClientSession) -> bytes:
- if is_url(file):
- async with session.get(yarl.URL(file, encoded=True)) as resp:
- return await resp.read()
- else:
- return base64.b64decode(file)
- def image_bytes_to_array(data: bytes) -> np.ndarray:
- return cv2.imdecode(np.frombuffer(data, np.uint8), cv2.IMREAD_COLOR)
- def image_to_base64(image: Image.Image) -> str:
- with io.BytesIO() as f:
- image.save(f, format="JPEG")
- image_base64 = base64.b64encode(f.getvalue()).decode("ascii")
- return image_base64
- def csv_bytes_to_data_frame(data: bytes) -> pd.DataFrame:
- with io.StringIO(data.decode("utf-8")) as f:
- df = pd.read_csv(f)
- return df
- def data_frame_to_base64(df: str) -> str:
- return base64.b64encode(df.to_csv().encode("utf-8")).decode("ascii")
- def read_pdf(
- bytes_: bytes, resize: bool = False, max_num_imgs: Optional[int] = None
- ) -> List[np.ndarray]:
- images: List[np.ndarray] = []
- img_size = None
- with fitz.open("pdf", bytes_) as doc:
- for page in doc:
- if max_num_imgs is not None and len(images) >= max_num_imgs:
- break
- # TODO: Do not always use zoom=2.0
- zoom = 2.0
- deg = 0
- mat = fitz.Matrix(zoom, zoom).prerotate(deg)
- pixmap = page.get_pixmap(matrix=mat, alpha=False)
- image = np.frombuffer(pixmap.samples, dtype=np.uint8).reshape(
- pixmap.h, pixmap.w, pixmap.n
- )
- image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
- if resize:
- if img_size is None:
- img_size = (image.shape[1], image.shape[0])
- else:
- if (image.shape[1], image.shape[0]) != img_size:
- image = cv2.resize(image, img_size)
- images.append(image)
- return images
- def file_to_images(
- file_bytes: bytes,
- file_type: Literal["IMAGE", "PDF"],
- *,
- max_img_size: Tuple[int, int],
- max_num_imgs: int,
- ) -> List[np.ndarray]:
- if file_type == "IMAGE":
- images = [image_bytes_to_array(file_bytes)]
- elif file_type == "PDF":
- images = read_pdf(file_bytes, resize=True, max_num_imgs=max_num_imgs)
- else:
- assert_never(file_type)
- h, w = images[0].shape[0:2]
- if w > max_img_size[1] or h > max_img_size[0]:
- if w / h > max_img_size[0] / max_img_size[1]:
- factor = max_img_size[0] / w
- else:
- factor = max_img_size[1] / h
- images = [cv2.resize(img, (int(factor * w), int(factor * h))) for img in images]
- return images
- def call_async(
- func: Callable[_P, _R], /, *args: _P.args, **kwargs: _P.kwargs
- ) -> Awaitable[_R]:
- return asyncio.get_running_loop().run_in_executor(
- None, partial(func, *args, **kwargs)
- )
|