xAdmin
41e08611c9
- Add image intake service with format validation (JPEG, PNG, WebP) and EXIF/GPS stripping - Enforce strict single-image rule across Web and Signal attachment channels - Implement token-optimized vision downscaling and JPEG compression - Add IMAGE_CONTEXT pipeline stage with OmniRoute vision routing and resilient failover - Seed and manage versioned idea-image-interpreter prompt in catalog - Update Web UI with responsive image picker, preview chip, and Visual Context tab - Add comprehensive automated test suite in test_image_intake.py - Update README and Labyricorn devlog
250 lines
8.8 KiB
Python
250 lines
8.8 KiB
Python
"""
|
|
ThinkStorm Image Intake & Sanitization Service
|
|
Handles safe image validation, EXIF/GPS metadata stripping, token-optimized vision resizing,
|
|
durable artifact persistence, and canonical metadata generation.
|
|
"""
|
|
|
|
import io
|
|
import os
|
|
import hashlib
|
|
from pathlib import Path
|
|
from typing import Optional, Dict, Any, Tuple
|
|
from PIL import Image, ImageOps
|
|
|
|
from ..config import config, BASE_DIR
|
|
|
|
ARTIFACTS_DIR = BASE_DIR / "data" / "artifacts"
|
|
ARTIFACTS_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Map PIL format strings to canonical MIME types and file extensions
|
|
FORMAT_TO_MIME = {
|
|
"JPEG": "image/jpeg",
|
|
"PNG": "image/png",
|
|
"WEBP": "image/webp"
|
|
}
|
|
|
|
MIME_TO_EXT = {
|
|
"image/jpeg": ".jpg",
|
|
"image/png": ".png",
|
|
"image/webp": ".webp"
|
|
}
|
|
|
|
class ImageValidationError(Exception):
|
|
"""Raised when an uploaded or received file fails image validation."""
|
|
pass
|
|
|
|
class ImageTooLargeError(ImageValidationError):
|
|
"""Raised when an image exceeds configured size limit."""
|
|
pass
|
|
|
|
class ImageFormatError(ImageValidationError):
|
|
"""Raised when image format is unsupported or invalid."""
|
|
pass
|
|
|
|
def sanitize_filename(filename: Optional[str], default_ext: str = ".jpg") -> str:
|
|
"""Strips path traversal elements and unsafe characters from original filename."""
|
|
if not filename:
|
|
return f"submission_image{default_ext}"
|
|
base = os.path.basename(filename).strip()
|
|
# Remove null bytes and non-printable characters
|
|
clean = "".join(c for c in base if c.isalnum() or c in "._- ")
|
|
if not clean or clean.startswith("."):
|
|
return f"submission_image{default_ext}"
|
|
return clean[:128]
|
|
|
|
def validate_and_sanitize_image(
|
|
raw_bytes: bytes,
|
|
original_filename: str = "",
|
|
source: str = "web",
|
|
idea_id: Optional[str] = None
|
|
) -> Tuple[bytes, Dict[str, Any]]:
|
|
"""
|
|
Validates, strips metadata, and generates canonical sanitized image bytes and metadata.
|
|
|
|
Returns:
|
|
Tuple[bytes, Dict[str, Any]]: (sanitized_bytes, metadata_dict)
|
|
|
|
Raises:
|
|
ImageTooLargeError: If payload exceeds max size.
|
|
ImageFormatError: If data is corrupted or unsupported format.
|
|
"""
|
|
if not raw_bytes:
|
|
raise ImageFormatError("Image payload cannot be empty.")
|
|
|
|
max_bytes = config.max_image_upload_bytes
|
|
if len(raw_bytes) > max_bytes:
|
|
raise ImageTooLargeError(
|
|
f"Image size ({len(raw_bytes)} bytes) exceeds maximum limit of {max_bytes} bytes ({max_bytes // (1024 * 1024)}MB)."
|
|
)
|
|
|
|
# 1. Open and verify image structure with Pillow
|
|
try:
|
|
in_stream = io.BytesIO(raw_bytes)
|
|
img = Image.open(in_stream)
|
|
img_format = img.format
|
|
except Exception as e:
|
|
raise ImageFormatError(f"Could not decode image header: {str(e)}")
|
|
|
|
if not img_format or img_format.upper() not in FORMAT_TO_MIME:
|
|
allowed_str = ", ".join(config.allowed_image_mime_types)
|
|
raise ImageFormatError(
|
|
f"Unsupported image format: '{img_format}'. Supported formats are: {allowed_str} (JPEG, PNG, WebP)."
|
|
)
|
|
|
|
canonical_mime = FORMAT_TO_MIME[img_format.upper()]
|
|
|
|
# 2. Decode pixel data to ensure file is not truncated or corrupted
|
|
try:
|
|
# Respect EXIF orientation before stripping EXIF metadata
|
|
img = ImageOps.exif_transpose(img)
|
|
img.load()
|
|
except Exception as e:
|
|
raise ImageFormatError(f"Malformed or corrupt image pixel stream: {str(e)}")
|
|
|
|
# 3. Create sanitized copy (stripping EXIF, GPS, camera metadata, comments)
|
|
out_format = img_format.upper()
|
|
out_stream = io.BytesIO()
|
|
|
|
try:
|
|
if out_format == "JPEG":
|
|
# Convert palette/RGBA modes to RGB for JPEG
|
|
if img.mode in ("RGBA", "LA", "P"):
|
|
rgb_img = Image.new("RGB", img.size, (255, 255, 255))
|
|
if img.mode == "P":
|
|
img = img.convert("RGBA")
|
|
rgb_img.paste(img, mask=img.split()[-1] if img.mode == "RGBA" else None)
|
|
img = rgb_img
|
|
elif img.mode != "RGB":
|
|
img = img.convert("RGB")
|
|
img.save(out_stream, format="JPEG", quality=92, optimize=True)
|
|
|
|
elif out_format == "PNG":
|
|
# Preserve RGB / RGBA transparency
|
|
if img.mode not in ("RGB", "RGBA", "L", "LA"):
|
|
img = img.convert("RGBA" if "transparency" in img.info else "RGB")
|
|
img.save(out_stream, format="PNG", optimize=True)
|
|
|
|
elif out_format == "WEBP":
|
|
if img.mode not in ("RGB", "RGBA"):
|
|
img = img.convert("RGBA" if "transparency" in img.info else "RGB")
|
|
img.save(out_stream, format="WEBP", quality=90, method=6)
|
|
else:
|
|
raise ImageFormatError(f"Unsupported save format: {out_format}")
|
|
|
|
except Exception as e:
|
|
raise ImageFormatError(f"Failed to encode sanitized image: {str(e)}")
|
|
|
|
sanitized_bytes = out_stream.getvalue()
|
|
width, height = img.size
|
|
sha256_hash = hashlib.sha256(sanitized_bytes).hexdigest()
|
|
ext = MIME_TO_EXT.get(canonical_mime, ".jpg")
|
|
safe_filename = sanitize_filename(original_filename, default_ext=ext)
|
|
|
|
artifact_id = f"IMG-{idea_id}" if idea_id else f"IMG-{sha256_hash[:12].upper()}"
|
|
|
|
metadata: Dict[str, Any] = {
|
|
"present": True,
|
|
"artifact_id": artifact_id,
|
|
"original_filename": safe_filename,
|
|
"mime_type": canonical_mime,
|
|
"size_bytes": len(sanitized_bytes),
|
|
"width": width,
|
|
"height": height,
|
|
"sha256": sha256_hash,
|
|
"source": source,
|
|
"vision_analysis_run_id": None
|
|
}
|
|
|
|
return sanitized_bytes, metadata
|
|
|
|
def save_image_artifact(idea_id: str, sanitized_bytes: bytes, metadata: Dict[str, Any]) -> str:
|
|
"""
|
|
Saves the canonical sanitized image file to the idea's artifact directory.
|
|
Returns relative or absolute path to saved artifact.
|
|
"""
|
|
idea_dir = ARTIFACTS_DIR / idea_id
|
|
idea_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
artifact_id = metadata.get("artifact_id") or f"IMG-{idea_id}"
|
|
mime = metadata.get("mime_type", "image/jpeg")
|
|
ext = MIME_TO_EXT.get(mime, ".jpg")
|
|
|
|
file_path = idea_dir / f"{artifact_id}{ext}"
|
|
file_path.write_bytes(sanitized_bytes)
|
|
|
|
# Also keep a predictable reference file for easy retrieval
|
|
canonical_link = idea_dir / f"reference_image{ext}"
|
|
if canonical_link != file_path:
|
|
canonical_link.write_bytes(sanitized_bytes)
|
|
|
|
return str(file_path)
|
|
|
|
def get_image_artifact_path(idea_id: str, metadata: Optional[Dict[str, Any]] = None) -> Optional[Path]:
|
|
"""Retrieves file path to an idea's saved reference image."""
|
|
idea_dir = ARTIFACTS_DIR / idea_id
|
|
if not idea_dir.exists():
|
|
return None
|
|
|
|
if metadata and metadata.get("artifact_id"):
|
|
artifact_id = metadata["artifact_id"]
|
|
mime = metadata.get("mime_type", "image/jpeg")
|
|
ext = MIME_TO_EXT.get(mime, ".jpg")
|
|
path = idea_dir / f"{artifact_id}{ext}"
|
|
if path.exists():
|
|
return path
|
|
|
|
# Fallback to search any image in idea artifact directory
|
|
for candidate_ext in [".jpg", ".jpeg", ".png", ".webp"]:
|
|
ref = idea_dir / f"reference_image{candidate_ext}"
|
|
if ref.exists():
|
|
return ref
|
|
img_match = list(idea_dir.glob(f"IMG-*{candidate_ext}"))
|
|
if img_match:
|
|
return img_match[0]
|
|
|
|
return None
|
|
|
|
def create_token_optimized_vision_payload(
|
|
sanitized_bytes: bytes,
|
|
mime_type: str = "image/jpeg"
|
|
) -> Tuple[bytes, str]:
|
|
"""
|
|
Optimizes and downscales image dimensions to bound vision LLM token consumption.
|
|
Bounds maximum dimension to `config.vision_max_dimension` (default 1536px)
|
|
and compresses to efficient JPEG to minimize token expenditure across vision grid tiles.
|
|
|
|
Returns:
|
|
Tuple[bytes, str]: (optimized_bytes, optimized_mime_type)
|
|
"""
|
|
try:
|
|
img = Image.open(io.BytesIO(sanitized_bytes))
|
|
width, height = img.size
|
|
max_dim = config.vision_max_dimension
|
|
|
|
# Downscale if larger than max_dim while preserving aspect ratio
|
|
if width > max_dim or height > max_dim:
|
|
img.thumbnail((max_dim, max_dim), Image.Resampling.LANCZOS)
|
|
|
|
# Convert to RGB if needed for JPEG compression
|
|
if img.mode in ("RGBA", "LA", "P"):
|
|
rgb_img = Image.new("RGB", img.size, (255, 255, 255))
|
|
if img.mode == "P":
|
|
img = img.convert("RGBA")
|
|
rgb_img.paste(img, mask=img.split()[-1] if img.mode == "RGBA" else None)
|
|
img = rgb_img
|
|
elif img.mode != "RGB":
|
|
img = img.convert("RGB")
|
|
|
|
out_buf = io.BytesIO()
|
|
img.save(
|
|
out_buf,
|
|
format="JPEG",
|
|
quality=config.vision_jpeg_quality,
|
|
optimize=True
|
|
)
|
|
return out_buf.getvalue(), "image/jpeg"
|
|
|
|
except Exception:
|
|
# Fallback to original sanitized bytes if downscaling fails
|
|
return sanitized_bytes, mime_type
|