Spaces:
Sleeping
Sleeping
| """Bounded uploads and in-memory image decoding.""" | |
| import io | |
| import warnings | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image, ImageOps, UnidentifiedImageError | |
| from fastapi import UploadFile | |
| from app.core.config import settings | |
| from app.core.exceptions import InvalidImageError, UnsupportedFormatError, PayloadTooLargeError | |
| from app.domain.schemas.core import ImageMetadata | |
| from app.services.image_processing.validator import detect_image_magic_bytes | |
| MAX_IMAGE_DIMENSION = 10_000 | |
| MAX_IMAGE_PIXELS = 50_000_000 | |
| Image.MAX_IMAGE_PIXELS = MAX_IMAGE_PIXELS | |
| async def read_upload(file: UploadFile) -> bytes: | |
| limit = settings.MAX_UPLOAD_SIZE_MB * 1024 * 1024 | |
| chunks, total = [], 0 | |
| try: | |
| while total <= limit: | |
| chunk = await file.read(min(1024 * 1024, limit + 1 - total)) | |
| if not chunk: | |
| break | |
| total += len(chunk) | |
| if total > limit: | |
| raise PayloadTooLargeError(f"File exceeds {settings.MAX_UPLOAD_SIZE_MB}MB upload limit.") | |
| chunks.append(chunk) | |
| finally: | |
| await file.seek(0) | |
| if not total: | |
| raise InvalidImageError("Empty file uploaded.") | |
| return b"".join(chunks) | |
| def validate_content(contents: bytes, declared_type: str | None, allow_pdf=False) -> str: | |
| mime = "application/pdf" if allow_pdf and contents.startswith(b"%PDF-") else detect_image_magic_bytes(contents[:32]) | |
| allowed = list(settings.ALLOWED_IMAGE_TYPES) + (["application/pdf"] if allow_pdf else []) | |
| if mime not in allowed: | |
| raise UnsupportedFormatError(f"Unsupported file contents. Allowed: {allowed}") | |
| declared_type = {"image/jpg": "image/jpeg", "image/x-ms-bmp": "image/bmp"}.get(declared_type, declared_type) | |
| if declared_type and declared_type not in ("application/octet-stream", mime): | |
| raise UnsupportedFormatError(f"Declared content type '{declared_type}' does not match '{mime}'.") | |
| return mime | |
| def check_dimensions(width, height): | |
| if width < 1 or height < 1 or max(width, height) > MAX_IMAGE_DIMENSION or width * height > MAX_IMAGE_PIXELS: | |
| raise InvalidImageError("Image exceeds supported dimensions or pixel limit.") | |
| def pil_to_bgr(image: Image.Image) -> np.ndarray: | |
| # Check headers before convert()/np.array() allocate decoded pixel buffers. | |
| check_dimensions(*image.size) | |
| transposed = ImageOps.exif_transpose(image) | |
| try: | |
| # Removing alpha without compositing can hide black text on transparency. | |
| if "A" in transposed.getbands() or "transparency" in transposed.info: | |
| rgba = transposed.convert("RGBA") | |
| try: | |
| background = Image.new("RGBA", rgba.size, "white") | |
| try: | |
| background.alpha_composite(rgba) | |
| rgb = np.array(background.convert("RGB")) | |
| finally: | |
| background.close() | |
| finally: | |
| rgba.close() | |
| else: | |
| rgb = np.array(transposed.convert("RGB")) | |
| return cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR) | |
| finally: | |
| transposed.close() | |
| async def load_image(file: UploadFile) -> tuple[np.ndarray, ImageMetadata]: | |
| contents = await read_upload(file) | |
| mime = validate_content(contents, file.content_type) | |
| from starlette.concurrency import run_in_threadpool | |
| image = await run_in_threadpool(decode_image, contents) | |
| height, width = image.shape[:2] | |
| return image, ImageMetadata(filename=file.filename or "unknown", width=width, height=height, | |
| content_type=mime, size_bytes=len(contents)) | |
| def decode_image(contents: bytes) -> np.ndarray: | |
| try: | |
| with warnings.catch_warnings(): | |
| warnings.simplefilter("error", Image.DecompressionBombWarning) | |
| with Image.open(io.BytesIO(contents)) as image: | |
| if getattr(image, "n_frames", 1) != 1: | |
| raise InvalidImageError("Multi-page images require /ocr/document; no pages were silently discarded.") | |
| return pil_to_bgr(image) | |
| except (Image.DecompressionBombWarning, Image.DecompressionBombError) as exc: | |
| raise InvalidImageError("Image exceeds safe decoding pixel limit.") from exc | |
| except (UnidentifiedImageError, OSError, ValueError) as exc: | |
| raise InvalidImageError("File is corrupted or cannot be decoded as an image.") from exc | |