Spaces:
Build error
Build error
File size: 2,398 Bytes
3c2cd23 | 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 | 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
@dataclass(slots=True)
class LoadedImage:
rgba: np.ndarray
icc_profile: bytes | None
source_format: str
@property
def width(self) -> int:
return int(self.rgba.shape[1])
@property
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)
|