bvandeusen
7e74fa767c
Trusted local library, not an upload surface, so a legitimately large image (90–95M px, operator-flagged) must load. PIL only WARNS at the 89M-px default but RAISES DecompressionBombError at ~179M px, which would fail those jobs. Set Image.MAX_IMAGE_PIXELS = None. (The agent works off individual extracted files — curator's archive_extractor unpacks zip/cbz/rar/7z at import — so this is about big single images, not archives.) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa
76 lines
3.1 KiB
Python
76 lines
3.1 KiB
Python
"""Image + video handling. Stills load directly; videos are sampled into frames
|
||
(ffmpeg) at the cadence FC sends — so a video becomes a bag of per-frame
|
||
instances, each with a timestamp."""
|
||
import io
|
||
import os
|
||
import subprocess
|
||
import tempfile
|
||
|
||
from PIL import Image, ImageFile
|
||
|
||
# Load slightly-truncated images (a few missing trailing bytes) instead of
|
||
# raising — matches the server embedder. These are common in scraped libraries
|
||
# and would otherwise fail the job 3× then error (operator-flagged 2026-06-30).
|
||
ImageFile.LOAD_TRUNCATED_IMAGES = True
|
||
|
||
# Disable PIL's decompression-bomb guard: this is a TRUSTED local library, not an
|
||
# untrusted upload surface, so a legitimately huge image (high-res scans/prints,
|
||
# 90M+ pixels) must load. The default 89M-pixel limit only WARNS, but PIL raises
|
||
# DecompressionBombError at 2× (~179M px) — which would fail those jobs outright
|
||
# (operator-flagged 2026-06-30, images of 90–95M px).
|
||
Image.MAX_IMAGE_PIXELS = None
|
||
|
||
|
||
def is_video(mime: str) -> bool:
|
||
return bool(mime) and (mime.startswith("video/") or mime in {"image/gif"})
|
||
|
||
|
||
def to_rgb(img: Image.Image) -> Image.Image:
|
||
"""RGB, flattening any transparency onto white first. A naive convert('RGB')
|
||
on a palette-with-transparency image (common for character PNGs on a clear
|
||
background) lets PIL guess the transparent pixels — usually black artifacts
|
||
that bleed into the crop + the embedding (and the "should be converted to
|
||
RGBA" warning). Compositing over white gives a clean, consistent background."""
|
||
if img.mode in ("RGBA", "LA", "PA") or (
|
||
img.mode == "P" and "transparency" in img.info
|
||
):
|
||
img = img.convert("RGBA")
|
||
bg = Image.new("RGBA", img.size, (255, 255, 255, 255))
|
||
return Image.alpha_composite(bg, img).convert("RGB")
|
||
return img.convert("RGB")
|
||
|
||
|
||
def load_image(data: bytes) -> Image.Image:
|
||
return to_rgb(Image.open(io.BytesIO(data)))
|
||
|
||
|
||
def sample_frames(
|
||
data: bytes, interval_seconds: float, max_frames: int
|
||
) -> list[tuple[float, Image.Image]]:
|
||
"""Extract up to max_frames frames at one-every-interval_seconds via ffmpeg.
|
||
Returns [(timestamp_seconds, frame)]. Empty on failure (caller falls back)."""
|
||
interval = max(0.5, float(interval_seconds or 4.0))
|
||
cap = max(1, int(max_frames or 64))
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
src = os.path.join(tmp, "in")
|
||
with open(src, "wb") as fh:
|
||
fh.write(data)
|
||
pattern = os.path.join(tmp, "f_%05d.jpg")
|
||
try:
|
||
subprocess.run(
|
||
[
|
||
"ffmpeg", "-nostdin", "-loglevel", "error", "-i", src,
|
||
"-vf", f"fps=1/{interval}", "-frames:v", str(cap),
|
||
"-q:v", "3", pattern,
|
||
],
|
||
check=True, timeout=600,
|
||
)
|
||
except (subprocess.SubprocessError, FileNotFoundError):
|
||
return []
|
||
out: list[tuple[float, Image.Image]] = []
|
||
names = sorted(n for n in os.listdir(tmp) if n.startswith("f_"))
|
||
for i, name in enumerate(names[:cap]):
|
||
with Image.open(os.path.join(tmp, name)) as im:
|
||
out.append((round(i * interval, 2), to_rgb(im)))
|
||
return out
|