Spaces:
Build error
Build error
| from __future__ import annotations | |
| from dataclasses import dataclass | |
| from io import BytesIO | |
| from pathlib import Path | |
| import numpy as np | |
| from PIL import Image, ImageOps, UnidentifiedImageError | |
| ALLOWED_FORMATS = {"PNG", "JPEG", "WEBP"} | |
| MAX_FILE_BYTES = 50 * 1024 * 1024 | |
| class InvalidImage(ValueError): | |
| pass | |
| class LoadedImage: | |
| rgba: np.ndarray | |
| icc_profile: bytes | None | |
| source_format: str | |
| def width(self) -> int: | |
| return int(self.rgba.shape[1]) | |
| def height(self) -> int: | |
| return int(self.rgba.shape[0]) | |
| def load_image_bytes(data: bytes, max_pixels: int = 40_000_000) -> LoadedImage: | |
| if not data: | |
| raise InvalidImage("The image is empty") | |
| if len(data) > MAX_FILE_BYTES: | |
| raise InvalidImage("The image exceeds the 50 MB limit") | |
| try: | |
| with Image.open(BytesIO(data)) as source: | |
| if source.format not in ALLOWED_FORMATS: | |
| raise InvalidImage("Only PNG, JPEG, and WebP images are supported") | |
| if getattr(source, "is_animated", False): | |
| raise InvalidImage("Animated images are not supported") | |
| source.load() | |
| image = ImageOps.exif_transpose(source) | |
| if image.width * image.height > max_pixels: | |
| raise InvalidImage(f"The decoded image exceeds {max_pixels:,} pixels") | |
| rgba = np.asarray(image.convert("RGBA"), dtype=np.uint8).copy() | |
| return LoadedImage( | |
| rgba=rgba, | |
| icc_profile=source.info.get("icc_profile"), | |
| source_format=source.format, | |
| ) | |
| except (UnidentifiedImageError, OSError) as exc: | |
| raise InvalidImage("The uploaded file is not a readable image") from exc | |
| def load_image_path(path: Path, max_pixels: int = 40_000_000) -> LoadedImage: | |
| try: | |
| return load_image_bytes(path.read_bytes(), max_pixels=max_pixels) | |
| except OSError as exc: | |
| raise InvalidImage(f"Could not read {path}") from exc | |
| def save_png(image: LoadedImage | np.ndarray, path: Path, icc_profile: bytes | None = None) -> None: | |
| rgba = image.rgba if isinstance(image, LoadedImage) else image | |
| if isinstance(image, LoadedImage): | |
| icc_profile = image.icc_profile | |
| target = Image.fromarray(rgba, "RGBA") | |
| target.save(path, format="PNG", icc_profile=icc_profile, optimize=True) | |