ollie/tools/image_read

72 lines
2.2 KiB
Python

#!/usr/bin/env python3
# ollie:parallel read
# ollie:tier cold
# ollie:prompt
# ## image_read
#
# Read an image file. Returns an image content block for the model.
#
# **Args**: `[path]` or `[path, "--region=x,y,w,h"]` or `[path, "--scale=N"]`
#
# ```
# call_tool: calls=[{tool: "image_read", args: ["/abs/path/image.png"]}]
# ```
# ollie:end
import sys
import os
import base64
import json
import mimetypes
file_path = None
for a in sys.argv[1:]:
if not a.startswith('--'):
file_path = a
if not file_path:
print('{"schema": {"type": "array", "items": [{"name": "file_path", "type": "string", "description": "Absolute path to image file", "required": true}, {"name": "--region=x,y,w,h", "type": "string", "description": "Region to capture (optional)", "required": false}, {"name": "--scale=N", "type": "string", "description": "Scale factor (optional)", "required": false}], "examples": [["/abs/path/image.png"], ["/abs/path/image.png", "--region=0,0,100,100"]]}}')
sys.exit(1)
if not os.path.isabs(file_path):
print("STATUS=error")
print("error: path must be absolute")
sys.exit(1)
if not os.path.exists(file_path):
print("STATUS=not_found")
sys.exit(0)
mime, _ = mimetypes.guess_type(file_path)
if not mime:
mime = 'application/octet-stream'
with open(file_path, 'rb') as f:
data = f.read()
MAX_SIZE = 128 * 1024
if mime == 'image/png' and len(data) > MAX_SIZE:
from PIL import Image
import io
img = Image.open(io.BytesIO(data))
# progressively scale down until under budget
for scale in [1.0, 0.75, 0.5, 0.35, 0.25]:
buf = io.BytesIO()
if scale < 1.0:
new_size = (int(img.width * scale), int(img.height * scale))
resized = img.resize(new_size, Image.LANCZOS)
else:
resized = img
resized.save(buf, format='PNG', optimize=True)
data = buf.getvalue()
if len(data) <= MAX_SIZE:
break
b64 = base64.b64encode(data).decode('ascii')
print(json.dumps({"isError": False, "content": [
{"type": "image", "media_type": mime, "data": b64},
{"type": "text", "text": f"[image: {os.path.basename(file_path)} {len(data)} bytes {mime}]"}
]}))