Spaces:
Sleeping
Sleeping
File size: 4,367 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 82 83 84 85 86 87 88 89 90 91 92 93 | """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
|