314 lines
12 KiB
Python
314 lines
12 KiB
Python
"""Image format detection, normalization and region cropping for vision tools.
|
|
|
|
Everything here runs BEFORE an image is base64-embedded. A vision tool result
|
|
is baked into immutable conversation history and re-sent every turn, so an
|
|
unsupported media type or corrupt bytes would wedge the session with a
|
|
non-retryable 400 on every resume — normalization must happen up front.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import uuid
|
|
from io import BytesIO
|
|
from pathlib import Path
|
|
from typing import Any, Optional
|
|
|
|
from hermes_constants import get_hermes_dir
|
|
|
|
logger = logging.getLogger("tools.vision_tools")
|
|
|
|
_EXTENSION_MIME_TYPES = {
|
|
".jpg": "image/jpeg",
|
|
".jpeg": "image/jpeg",
|
|
".png": "image/png",
|
|
".gif": "image/gif",
|
|
".bmp": "image/bmp",
|
|
".webp": "image/webp",
|
|
".svg": "image/svg+xml",
|
|
}
|
|
|
|
# Media types the major vision providers (Anthropic in particular) accept
|
|
# inline. SVG/BMP/TIFF are rejected with a non-retryable 400.
|
|
_ANTHROPIC_SUPPORTED_MEDIA_TYPES = frozenset(
|
|
{"image/jpeg", "image/png", "image/gif", "image/webp"}
|
|
)
|
|
|
|
|
|
def _determine_mime_type(image_path: Path) -> str:
|
|
"""MIME type from file extension (defaults to image/jpeg)."""
|
|
return _EXTENSION_MIME_TYPES.get(image_path.suffix.lower(), "image/jpeg")
|
|
|
|
|
|
def _detect_image_mime_type_from_bytes(data: bytes) -> Optional[str]:
|
|
"""Magic-byte MIME sniff (authoritative; no extension trust).
|
|
|
|
Returns ``None`` for anything without a recognized header — including SVG,
|
|
which has no magic bytes (the resolver sniffs ``<svg`` and passes it
|
|
through for rasterization).
|
|
"""
|
|
header = data[:64]
|
|
if header.startswith(b"\x89PNG\r\n\x1a\n"):
|
|
# Magic bytes alone are insufficient: reject corrupt PNGs before they
|
|
# can be embedded. Pillow is optional — without it fall back to
|
|
# header-only sniffing; only an actual failed verify() rejects.
|
|
try:
|
|
from PIL import Image
|
|
except ImportError:
|
|
return "image/png"
|
|
try:
|
|
with Image.open(BytesIO(data)) as image:
|
|
image.verify()
|
|
except Exception:
|
|
return None
|
|
return "image/png"
|
|
if header.startswith(b"\xff\xd8\xff"):
|
|
return "image/jpeg"
|
|
if header.startswith((b"GIF87a", b"GIF89a")):
|
|
return "image/gif"
|
|
if header.startswith(b"BM"):
|
|
return "image/bmp"
|
|
if len(header) >= 12 and header[:4] == b"RIFF" and header[8:12] == b"WEBP":
|
|
return "image/webp"
|
|
return None
|
|
|
|
|
|
def _supported_media_types() -> frozenset:
|
|
"""Formats the ACTIVE main model's server can decode.
|
|
|
|
The managed llama-server decodes with stb_image — no WebP — and an
|
|
undecodable image part fails SILENTLY (the model confabulates), so the set
|
|
is narrowed there and normalization converts those formats to PNG.
|
|
"""
|
|
try:
|
|
from agent.auxiliary_client import _runtime_main_value
|
|
from hermes_cli.local_runtime.capabilities import (
|
|
ACCEPTED_IMAGE_MIMES,
|
|
is_managed_provider,
|
|
)
|
|
|
|
if is_managed_provider(
|
|
str(_runtime_main_value("provider") or ""),
|
|
str(_runtime_main_value("base_url") or "")):
|
|
return ACCEPTED_IMAGE_MIMES
|
|
except Exception: # noqa: BLE001 — best-effort narrowing only
|
|
pass
|
|
return _ANTHROPIC_SUPPORTED_MEDIA_TYPES
|
|
|
|
|
|
def _rasterize_svg_to_png(svg_path: Path, out_path: Path) -> bool:
|
|
"""Best-effort SVG → PNG via cairosvg, svglib+reportlab, rsvg-convert, inkscape (all soft deps)."""
|
|
def _ok() -> bool:
|
|
return out_path.exists() and out_path.stat().st_size > 0
|
|
|
|
try:
|
|
import cairosvg # type: ignore
|
|
cairosvg.svg2png(url=str(svg_path), write_to=str(out_path))
|
|
return _ok()
|
|
except Exception:
|
|
pass
|
|
try:
|
|
from svglib.svglib import svg2rlg # type: ignore
|
|
from reportlab.graphics import renderPM # type: ignore
|
|
drawing = svg2rlg(str(svg_path))
|
|
if drawing is not None:
|
|
renderPM.drawToFile(drawing, str(out_path), fmt="PNG")
|
|
return _ok()
|
|
except Exception:
|
|
pass
|
|
import shutil
|
|
import subprocess
|
|
for cmd in (
|
|
["rsvg-convert", "-o", str(out_path), str(svg_path)],
|
|
["inkscape", str(svg_path), "--export-type=png",
|
|
f"--export-filename={out_path}"],
|
|
):
|
|
if shutil.which(cmd[0]):
|
|
try:
|
|
subprocess.run(
|
|
cmd, check=True, capture_output=True, timeout=30,
|
|
stdin=subprocess.DEVNULL,
|
|
)
|
|
if _ok():
|
|
return True
|
|
except Exception:
|
|
continue
|
|
return False
|
|
|
|
|
|
def _normalize_to_supported_image(
|
|
image_path: Path, detected_mime: str
|
|
) -> tuple[Optional[Path], Optional[str], Optional[str]]:
|
|
"""Ensure an image is in a provider-supported format.
|
|
|
|
Returns ``(path, mime, error)``: the input unchanged when already
|
|
supported; ``(new_png_path, "image/png", None)`` after conversion — a temp
|
|
file the CALLER must clean up; ``(None, None, message)`` when impossible.
|
|
SVG is rasterized; other Pillow-readable rasters (BMP, TIFF) re-encode to PNG.
|
|
"""
|
|
if detected_mime in _supported_media_types():
|
|
return image_path, detected_mime, None
|
|
|
|
out_dir = get_hermes_dir("cache/vision", "temp_vision_images")
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
out_path = out_dir / f"converted_{uuid.uuid4()}.png"
|
|
|
|
if detected_mime == "image/svg+xml":
|
|
if _rasterize_svg_to_png(image_path, out_path):
|
|
return out_path, "image/png", None
|
|
return (
|
|
None,
|
|
None,
|
|
"This is an SVG, which vision models cannot read directly, and no "
|
|
"SVG rasterizer is installed (tried cairosvg, svglib, rsvg-convert, "
|
|
"inkscape). Convert the SVG to PNG first — e.g. open it in a browser "
|
|
"and screenshot it, or install a rasterizer "
|
|
"(`pip install cairosvg`) — then re-run vision_analyze on the PNG.",
|
|
)
|
|
|
|
try:
|
|
from PIL import Image as _PILImage
|
|
with _PILImage.open(image_path) as _img:
|
|
if _img.mode not in ("RGB", "RGBA", "L"):
|
|
_img = _img.convert("RGBA")
|
|
_img.save(out_path, format="PNG")
|
|
if out_path.exists() and out_path.stat().st_size > 0:
|
|
return out_path, "image/png", None
|
|
except Exception as _exc:
|
|
logger.warning("Failed to normalize %s image to PNG: %s",
|
|
detected_mime, _exc)
|
|
return (
|
|
None,
|
|
None,
|
|
f"Image format {detected_mime!r} is not supported by the vision API "
|
|
f"and could not be converted to PNG (install Pillow for raster "
|
|
f"conversion). Convert it to PNG or JPEG and try again.",
|
|
)
|
|
|
|
|
|
# Full raster validation runs on untrusted images in a shared CPU executor:
|
|
# bound animated-image work by frame count AND total decoded area so a compact
|
|
# file cannot monopolize a worker with unbounded frames.
|
|
_VISION_MAX_VALIDATED_FRAME_COUNT = 100
|
|
_VISION_MAX_VALIDATED_AGGREGATE_PIXELS = 100_000_000
|
|
|
|
|
|
def _validate_raster_image_decodable(
|
|
image_path: Path,
|
|
max_frames: int = _VISION_MAX_VALIDATED_FRAME_COUNT,
|
|
max_pixels: int = _VISION_MAX_VALIDATED_AGGREGATE_PIXELS,
|
|
) -> Optional[str]:
|
|
"""Return an error unless Pillow can fully decode every frame.
|
|
|
|
Header sniffing and ``Image.open`` only inspect containers: a timed-out
|
|
download can look like a valid PNG with a truncated pixel stream. Without
|
|
Pillow the image passes unvalidated rather than rejecting everything.
|
|
"""
|
|
try:
|
|
from PIL import Image as _PILImage
|
|
from PIL import ImageSequence as _PILImageSequence
|
|
except ImportError:
|
|
return None
|
|
try:
|
|
with _PILImage.open(image_path) as image:
|
|
image.verify()
|
|
with _PILImage.open(image_path) as image:
|
|
validated_pixels = 0
|
|
for frame_number, frame in enumerate(
|
|
_PILImageSequence.Iterator(image), start=1
|
|
):
|
|
if frame_number > max_frames:
|
|
return (
|
|
"Image validation rejected animation: "
|
|
f"frame {frame_number} exceeds the maximum "
|
|
f"{max_frames} validated frames."
|
|
)
|
|
next_validated_pixels = validated_pixels + frame.width * frame.height
|
|
if next_validated_pixels > max_pixels:
|
|
return (
|
|
"Image validation rejected animation: aggregate decoded "
|
|
f"pixel count would reach {next_validated_pixels} at frame "
|
|
f"{frame_number}, exceeding the maximum "
|
|
f"{max_pixels}."
|
|
)
|
|
frame.load()
|
|
validated_pixels = next_validated_pixels
|
|
except Exception as exc:
|
|
return f"Image could not be fully decoded: {exc}"
|
|
return None
|
|
|
|
|
|
def _image_exceeds_dimension(image_path: Path, max_dimension: int) -> bool:
|
|
"""True if the longest side exceeds ``max_dimension`` px.
|
|
|
|
Anthropic enforces an 8000px per-side cap independently of the byte cap.
|
|
Returns False (no forced resize) without Pillow or on unreadable files —
|
|
a missing soft dependency must never break the embed path.
|
|
"""
|
|
try:
|
|
from PIL import Image as _PILImage
|
|
with _PILImage.open(image_path) as _img:
|
|
return max(_img.size) > max_dimension
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def _crop_image_region(
|
|
image_path: Path,
|
|
region: Any,
|
|
offset_out: Optional[dict] = None,
|
|
) -> tuple[Optional[Path], Optional[str], Optional[str]]:
|
|
"""Crop to ``region`` = [x1, y1, x2, y2] (original-image pixels).
|
|
|
|
Applied BEFORE downscaling so the crop gets the full resolution budget.
|
|
Coordinates clamp to the image bounds; a zero-area/inverted region is
|
|
rejected with an error naming the real dimensions. Returns
|
|
``(cropped_temp_path, mime, None)`` — caller owns cleanup — or
|
|
``(None, None, error)``. Ported from QwenLM/qwen-code zoom-image.ts (Apache-2.0).
|
|
"""
|
|
try:
|
|
from PIL import Image
|
|
except ImportError:
|
|
return None, None, (
|
|
"region cropping requires Pillow (`pip install Pillow`); "
|
|
"retry without the region parameter."
|
|
)
|
|
|
|
if (
|
|
not isinstance(region, (list, tuple))
|
|
or len(region) != 4
|
|
or not all(isinstance(v, (int, float)) and not isinstance(v, bool) for v in region)
|
|
):
|
|
return None, None, (
|
|
"Invalid region: expected [x1, y1, x2, y2] as four numbers "
|
|
"(pixel coordinates in the original image)."
|
|
)
|
|
|
|
try:
|
|
with Image.open(image_path) as img:
|
|
width, height = img.size
|
|
x1, y1, x2, y2 = (int(v) for v in region)
|
|
cx1 = max(0, min(x1, width))
|
|
cy1 = max(0, min(y1, height))
|
|
cx2 = max(0, min(x2, width))
|
|
cy2 = max(0, min(y2, height))
|
|
if cx2 <= cx1 or cy2 <= cy1:
|
|
return None, None, (
|
|
f"Invalid region [{x1}, {y1}, {x2}, {y2}]: crops to zero "
|
|
f"area after clamping to the image bounds. The image is "
|
|
f"{width}x{height} px — pick x1<x2 and y1<y2 inside "
|
|
f"[0, 0, {width}, {height}]."
|
|
)
|
|
cropped = img.crop((cx1, cy1, cx2, cy2))
|
|
if offset_out is not None:
|
|
offset_out.update(x=cx1, y=cy1, width=cx2 - cx1, height=cy2 - cy1)
|
|
out_path = image_path.with_name(
|
|
f"{image_path.stem}_region_{uuid.uuid4().hex[:8]}.png"
|
|
)
|
|
if cropped.mode not in ("RGB", "RGBA", "L", "LA", "P"):
|
|
cropped = cropped.convert("RGB")
|
|
cropped.save(out_path, format="PNG")
|
|
return out_path, "image/png", None
|
|
except Exception as exc:
|
|
return None, None, f"Failed to crop region: {exc}"
|