186 lines
6.6 KiB
Python
186 lines
6.6 KiB
Python
import os
|
|
import base64
|
|
import mimetypes
|
|
import logging
|
|
import asyncio
|
|
import io
|
|
import datetime
|
|
from typing import Any, Callable, Dict, List, Union
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
logger = logging.getLogger("MattCP")
|
|
|
|
class Tool:
|
|
def __init__(self, name: str, description: str, schema: Dict[str, Any], handler: Callable):
|
|
self.name = name
|
|
self.description = description
|
|
self.schema = schema
|
|
self.handler = handler
|
|
|
|
def to_mcp_dict(self):
|
|
return {
|
|
"name": self.name,
|
|
"description": self.description,
|
|
"inputSchema": self.schema
|
|
}
|
|
|
|
async def read_image_handler(args: Dict[str, Any]):
|
|
path = args.get("path")
|
|
if not path:
|
|
return {"text": "Error: Missing path argument"}
|
|
|
|
mime_type, _ = mimetypes.guess_type(path)
|
|
if not mime_type or not mime_type.startswith("image/"):
|
|
return {"text": f"Error: File {path} is not a supported image format."}
|
|
|
|
try:
|
|
with open(path, "rb") as f:
|
|
data = base64.b64encode(f.read()).decode("utf-8")
|
|
return {"type": "image", "data": data, "mimeType": mime_type}
|
|
except Exception as e:
|
|
return {"text": f"Error reading image: {str(e)}"}
|
|
|
|
async def read_png_metadata_handler(args: Dict[str, Any]):
|
|
path = args.get("path")
|
|
if not path:
|
|
return {"text": "Error: Missing path argument"}
|
|
|
|
try:
|
|
with Image.open(path) as img:
|
|
if img.format != 'PNG':
|
|
return {"text": "Error: File is not a PNG image."}
|
|
metadata = img.info
|
|
if not metadata:
|
|
return {"text": "No metadata found in this PNG."}
|
|
output = "\n".join([f"{k}: {v}" for k, v in metadata.items()])
|
|
return {"text": f"PNG Metadata:\n{output}"}
|
|
except Exception as e:
|
|
return {"text": f"Error reading PNG metadata: {str(e)}"}
|
|
|
|
async def list_thumbnails_handler(args: Dict[str, Any]):
|
|
dir_path = args.get("path")
|
|
page = int(args.get("page", 1))
|
|
page_size = int(args.get("page_size", 64))
|
|
|
|
if not dir_path:
|
|
return {"text": "Error: Missing path argument"}
|
|
if not os.path.isdir(dir_path):
|
|
return {"text": f"Error: {dir_path} is not a directory."}
|
|
|
|
exts = ('.png', '.jpg', '.jpeg', '.webp', '.bmp')
|
|
all_files = sorted([f for f in os.listdir(dir_path) if f.lower().endswith(exts)])
|
|
|
|
if not all_files:
|
|
return {"text": "No supported images found in the directory."}
|
|
|
|
total_files = len(all_files)
|
|
start_idx = (page - 1) * page_size
|
|
end_idx = start_idx + page_size
|
|
files = all_files[start_idx:end_idx]
|
|
|
|
if not files:
|
|
return {"text": f"No images found on page {page}."}
|
|
|
|
thumb_size = 160
|
|
padding = 15
|
|
label_height = 30
|
|
|
|
num_files = len(files)
|
|
cols = int(num_files**0.5) if num_files > 0 else 1
|
|
if cols == 0: cols = 1
|
|
rows = (num_files + cols - 1) // cols
|
|
|
|
canvas_w = cols * (thumb_size + padding) + padding
|
|
canvas_h = rows * (thumb_size + label_height + padding) + padding
|
|
|
|
canvas = Image.new('RGB', (canvas_w, canvas_h), (30, 30, 30))
|
|
draw = ImageDraw.Draw(canvas)
|
|
|
|
try:
|
|
font = ImageFont.load_default()
|
|
except Exception:
|
|
font = None
|
|
|
|
file_details = []
|
|
|
|
for i, filename in enumerate(files):
|
|
full_path = os.path.join(dir_path, filename)
|
|
row = i // cols
|
|
col = i % cols
|
|
|
|
x = padding + col * (thumb_size + padding)
|
|
y = padding + row * (thumb_size + label_height + padding)
|
|
|
|
# Gather metadata for the text list
|
|
try:
|
|
stats = os.stat(full_path)
|
|
with Image.open(full_path) as img:
|
|
width, height = img.size
|
|
# Process thumbnail
|
|
img.thumbnail((thumb_size, thumb_size))
|
|
off_x = (thumb_size - img.width) // 2
|
|
off_y = (thumb_size - img.height) // 2
|
|
canvas.paste(img, (x + off_x, y + off_y))
|
|
|
|
mod_time = datetime.datetime.fromtimestamp(stats.st_mtime).strftime('%Y-%m-%d %H:%M')
|
|
size_kb = stats.st_size / 1024
|
|
file_details.append(f"{i+1}. {filename} | {width}x{height} | {size_kb:.1f}KB | {mod_time}")
|
|
|
|
# Label the thumbnail with just the index number
|
|
draw.text((x, y + thumb_size + 2), str(i + 1), fill=(220, 220, 220), font=font)
|
|
except Exception as e:
|
|
logger.error(f"Failed to process {filename}: {e}")
|
|
draw.text((x, y + thumb_size // 2), "Error", fill=(255, 0, 0), font=font)
|
|
file_details.append(f"{i+1}. {filename} | ERROR")
|
|
|
|
buf = io.BytesIO()
|
|
canvas.save(buf, format='PNG')
|
|
img_data = base64.b64encode(buf.getvalue()).decode("utf-8")
|
|
|
|
total_pages = (total_files + page_size - 1) // page_size
|
|
header_text = f"Directory: {dir_path}\nPage {page} of {total_pages} ({total_files} total images). Showing {start_idx+1}-{min(end_idx, total_files)}."
|
|
|
|
details_text = "\n".join(file_details)
|
|
|
|
return [
|
|
{"type": "text", "text": f"{header_text}\n\nFile List:\n{details_text}"},
|
|
{"type": "image", "data": img_data, "mimeType": "image/png"}
|
|
]
|
|
|
|
TOOL_REGISTRY = [
|
|
Tool(
|
|
name="read_image",
|
|
description="Reads an image from the disk and returns it as an image content object.",
|
|
schema={
|
|
"type": "object",
|
|
"properties": {"path": {"type": "string", "description": "Path to the image file"}},
|
|
"required": ["path"],
|
|
},
|
|
handler=read_image_handler
|
|
),
|
|
Tool(
|
|
name="read_png_metadata",
|
|
description="Reads the metadata (tEXt, zTXt, iTXt) from a PNG image to fetch info such as AI prompts.",
|
|
schema={
|
|
"type": "object",
|
|
"properties": {"path": {"type": "string", "description": "Path to the PNG file"}},
|
|
"required": ["path"],
|
|
},
|
|
handler=read_png_metadata_handler
|
|
),
|
|
Tool(
|
|
name="list_thumbnails",
|
|
description="Generates a thumbnail grid of images in a directory. Returns a numbered grid and a detailed file list.",
|
|
schema={
|
|
"type": "object",
|
|
"properties": {
|
|
"path": {"type": "string", "description": "Path to the directory containing images"},
|
|
"page": {"type": "integer", "description": "The page number to retrieve (1-indexed)", "default": 1},
|
|
"page_size": {"type": "integer", "description": "Number of images per page", "default": 64},
|
|
},
|
|
"required": ["path"],
|
|
},
|
|
handler=list_thumbnails_handler
|
|
),
|
|
]
|