42 lines
1.4 KiB
Python
42 lines
1.4 KiB
Python
|
"""Shared validation and normalization for user-supplied images."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from io import BytesIO
|
||
|
|
|
||
|
|
from PIL import Image, ImageOps, UnidentifiedImageError
|
||
|
|
|
||
|
|
import settings
|
||
|
|
|
||
|
|
|
||
|
|
def normalize_image(data: bytes) -> tuple[bytes, int, int]:
|
||
|
|
"""Return a bounded, orientation-corrected JPEG safe for model input."""
|
||
|
|
|
||
|
|
try:
|
||
|
|
source = Image.open(BytesIO(data))
|
||
|
|
width, height = source.size
|
||
|
|
if width <= 0 or height <= 0:
|
||
|
|
raise ValueError("图片尺寸无效")
|
||
|
|
if width * height > settings.INPUT_IMAGE_MAX_PIXELS:
|
||
|
|
raise ValueError("图片像素尺寸过大")
|
||
|
|
source.load()
|
||
|
|
except (UnidentifiedImageError, OSError) as exc:
|
||
|
|
raise ValueError("文件不是可识别的图片") from exc
|
||
|
|
|
||
|
|
image = ImageOps.exif_transpose(source)
|
||
|
|
image.thumbnail(
|
||
|
|
(settings.INPUT_IMAGE_MAX_EDGE, settings.INPUT_IMAGE_MAX_EDGE),
|
||
|
|
Image.Resampling.LANCZOS,
|
||
|
|
)
|
||
|
|
if image.mode in {"RGBA", "LA"} or "transparency" in image.info:
|
||
|
|
rgba = image.convert("RGBA")
|
||
|
|
background = Image.new("RGB", rgba.size, "white")
|
||
|
|
background.paste(rgba, mask=rgba.getchannel("A"))
|
||
|
|
image = background
|
||
|
|
else:
|
||
|
|
image = image.convert("RGB")
|
||
|
|
|
||
|
|
buffer = BytesIO()
|
||
|
|
image.save(buffer, format="JPEG", quality=85, optimize=True)
|
||
|
|
return buffer.getvalue(), image.width, image.height
|