Spaces:
Sleeping
Sleeping
File size: 3,678 Bytes
a39757f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | """Sequential bounded PDF/TIFF/image page decoding. No external OCR service."""
import io
import math
import threading
import warnings
from PIL import Image, UnidentifiedImageError
from app.core.config import settings
from app.core.exceptions import InvalidImageError, PayloadTooLargeError, ConfigurationError
from app.services.image_processing.loader import pil_to_bgr, check_dimensions
# PDFium is not thread safe, including calls on different documents.
_pdf_lock = threading.Lock()
MAX_DOCUMENT_PIXELS = 100_000_000
MAX_PDF_PAGE_PIXELS = 12_000_000
def _check_pages(count):
if count < 1:
raise InvalidImageError("Document has no pages.")
if count > settings.DOCUMENT_MAX_PAGES:
raise PayloadTooLargeError(f"Document exceeds {settings.DOCUMENT_MAX_PAGES} page limit.")
def document_pages(contents, mime):
if mime == "application/pdf":
yield from _pdf_pages(contents)
return
try:
with warnings.catch_warnings():
warnings.simplefilter("error", Image.DecompressionBombWarning)
with Image.open(io.BytesIO(contents)) as document:
count = getattr(document, "n_frames", 1)
_check_pages(count)
pixels = 0
for index in range(count):
document.seek(index)
check_dimensions(*document.size)
pixels += document.width * document.height
if pixels > MAX_DOCUMENT_PIXELS:
raise PayloadTooLargeError("Document exceeds total pixel limit.")
yield pil_to_bgr(document)
except (Image.DecompressionBombWarning, Image.DecompressionBombError, UnidentifiedImageError, OSError, ValueError) as exc:
raise InvalidImageError("Image document is corrupted or exceeds safe decoding limits.") from exc
def _pdf_pages(contents):
try:
import pypdfium2 as pdfium
except ImportError as exc:
raise ConfigurationError("PDF support requires pypdfium2; install project requirements.") from exc
with _pdf_lock:
try:
with pdfium.PdfDocument(contents) as document:
_check_pages(len(document))
total_pixels = 0
for index in range(len(document)):
page = document[index]
try:
width, height = page.get_size()
if not all(math.isfinite(n) and n > 0 for n in (width, height)):
raise InvalidImageError("PDF has invalid page dimensions.")
scale = min(settings.DOCUMENT_DPI / 72, settings.OCR_CANVAS_SIZE / max(width, height),
math.sqrt(MAX_PDF_PAGE_PIXELS / (width * height)))
total_pixels += math.ceil(width * scale) * math.ceil(height * scale)
if total_pixels > MAX_DOCUMENT_PIXELS:
raise PayloadTooLargeError("PDF exceeds total rendered pixel limit.")
bitmap = page.render(scale=scale)
try:
pil_image = bitmap.to_pil()
try:
image = pil_to_bgr(pil_image)
finally:
pil_image.close()
finally:
bitmap.close()
finally:
page.close()
yield image
except pdfium.PdfiumError as exc:
raise InvalidImageError("PDF is corrupted, encrypted or cannot be rendered.") from exc
|