summaryrefslogtreecommitdiff
path: root/src/ollama_web/tools/helper/image.py
diff options
context:
space:
mode:
Diffstat (limited to 'src/ollama_web/tools/helper/image.py')
-rw-r--r--src/ollama_web/tools/helper/image.py65
1 files changed, 65 insertions, 0 deletions
diff --git a/src/ollama_web/tools/helper/image.py b/src/ollama_web/tools/helper/image.py
new file mode 100644
index 0000000..e89a96f
--- /dev/null
+++ b/src/ollama_web/tools/helper/image.py
@@ -0,0 +1,65 @@
+"""Image preprocessing helpers for ollama vision models."""
+
+from __future__ import annotations
+
+import io
+from pathlib import Path
+
+from PIL import Image
+
+from ...config import settings
+
+Image.MAX_IMAGE_PIXELS = settings.max_image_pixels
+
+
+_IMAGE_FORMATS: dict[str, str] = {
+ "png": "PNG",
+ "jpg": "JPEG",
+ "jpeg": "JPEG",
+ "gif": "GIF",
+ "bmp": "BMP",
+ "webp": "WEBP",
+ "tif": "TIFF",
+ "tiff": "TIFF",
+}
+
+
+def _format_from_name(name: str) -> str:
+ ext = Path(name).suffix.lstrip(".").lower()
+ return _IMAGE_FORMATS.get(ext, "JPEG")
+
+
+def resize_image(
+ data: bytes,
+ name: str = "image.jpg",
+ max_dimension: int = 1024,
+ quality: int = 85,
+) -> bytes:
+ """Resize and re-encode an image to reduce payload size for ollama."""
+ if not data:
+ return data
+
+ try:
+ with Image.open(io.BytesIO(data)) as img:
+ if img.width * img.height > settings.max_image_pixels:
+ return data
+ processed = img.convert("RGB")
+
+ width, height = processed.size
+ if width > max_dimension or height > max_dimension:
+ ratio = min(max_dimension / width, max_dimension / height)
+ new_size = (int(width * ratio), int(height * ratio))
+ processed = processed.resize(new_size, Image.Resampling.LANCZOS)
+
+ fmt = _format_from_name(name)
+ output = io.BytesIO()
+ save_kwargs: dict[str, object] = {}
+ if fmt == "JPEG":
+ save_kwargs["quality"] = quality
+ save_kwargs["optimize"] = True
+ elif fmt == "WEBP":
+ save_kwargs["quality"] = quality
+ processed.save(output, format=fmt, **save_kwargs)
+ return output.getvalue()
+ except Exception: # noqa: BLE001
+ return data