133 lines
5.1 KiB
Python
133 lines
5.1 KiB
Python
import math
|
|
import base64
|
|
import logging
|
|
import io
|
|
import datetime
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Tuple
|
|
from PIL import Image, ImageOps
|
|
import config
|
|
from tools.utils import format_relative_time, ToolError, get_file_info_list, sort_file_list, parse_indices
|
|
|
|
|
|
logger = logging.getLogger("MooseCP")
|
|
|
|
def calculate_patch_dimensions(orig_w: int, orig_h: int, target_tokens: int = None):
|
|
"""
|
|
Calculates optimal pixel dimensions to hit a token budget of at least target_tokens.
|
|
Validates budget based on actual rounded pixel dimensions to avoid float precision errors.
|
|
"""
|
|
if target_tokens is None:
|
|
target_tokens = config.PREVIEW_TOKEN_BUDGET
|
|
|
|
ar = orig_w / orig_h
|
|
|
|
# Case A: Width is the anchor (non-padded)
|
|
w_anchor = 1
|
|
while True:
|
|
w_px = w_anchor * config.PATCH_SIZE
|
|
h_px = round(w_px / ar)
|
|
# Calculate actual tokens based on resulting pixel dimensions
|
|
tokens = math.ceil(w_px / config.PATCH_SIZE) * math.ceil(h_px / config.PATCH_SIZE)
|
|
if tokens >= target_tokens:
|
|
res_a = {"w": w_px, "h": h_px, "tokens": tokens}
|
|
break
|
|
w_anchor += 1
|
|
|
|
# Case B: Height is the anchor (non-padded)
|
|
h_anchor = 1
|
|
while True:
|
|
h_px = h_anchor * config.PATCH_SIZE
|
|
w_px = round(h_px * ar)
|
|
# Calculate actual tokens based on resulting pixel dimensions
|
|
tokens = math.ceil(w_px / config.PATCH_SIZE) * math.ceil(h_px / config.PATCH_SIZE)
|
|
if tokens >= target_tokens:
|
|
res_b = {"w": w_px, "h": h_px, "tokens": tokens}
|
|
break
|
|
h_anchor += 1
|
|
|
|
# Pick the one with the smaller token count
|
|
best = res_a if res_a["tokens"] <= res_b["tokens"] else res_b
|
|
return max(1, best["w"]), max(1, best["h"])
|
|
|
|
async def handle(args: Dict[str, Any]):
|
|
"""
|
|
Generates small thumbnails of images with detailed info.
|
|
Can take a single file path, or a directory with indices/ranges from contact_sheet.
|
|
Thumbnails are optimized for ~70 token usage.
|
|
"""
|
|
path_str = args.get("path")
|
|
indices_str = args.get("indices")
|
|
sort_by = args.get("sort_by", "mtime")
|
|
|
|
if not path_str:
|
|
raise ToolError("Missing path argument")
|
|
|
|
path = Path(path_str)
|
|
# Store as (path, index_label) where index_label is None if direct path
|
|
target_files: List[Tuple[Path, Any]] = []
|
|
|
|
if path.is_file():
|
|
target_files.append((path, None))
|
|
elif path.is_dir():
|
|
if not indices_str:
|
|
raise ToolError("Indices are required when providing a directory path. Example: '1, 5-10'")
|
|
|
|
exts = ('.png', '.jpg', '.jpeg', '.webp', '.bmp')
|
|
file_info_list = get_file_info_list(path, extensions=exts)
|
|
file_info_list = sort_file_list(file_info_list, sort_by)
|
|
|
|
indices = parse_indices(indices_str)
|
|
for idx in indices:
|
|
# contact_sheet uses 1-based indexing
|
|
if 1 <= idx <= len(file_info_list):
|
|
target_files.append((file_info_list[idx-1]["path"], idx))
|
|
else:
|
|
logger.warning(f"Index {idx} out of range for directory {path_str}")
|
|
else:
|
|
raise ToolError(f"Path {path_str} is neither a file nor a directory.")
|
|
|
|
if not target_files:
|
|
return {"text": "No valid images found to preview."}
|
|
|
|
results = []
|
|
for file_path, index_label in target_files:
|
|
try:
|
|
stats = file_path.stat()
|
|
with Image.open(file_path) as img:
|
|
# Apply EXIF orientation to fix sideways images
|
|
img = ImageOps.exif_transpose(img)
|
|
orig_w, orig_h = img.size
|
|
target_w, target_h = calculate_patch_dimensions(orig_w, orig_h)
|
|
|
|
# Use thumbnail() to preserve aspect ratio and fit within the target bounding box.
|
|
thumb = img.convert("RGB")
|
|
thumb.thumbnail((target_w, target_h), Image.Resampling.LANCZOS)
|
|
|
|
buf = io.BytesIO()
|
|
thumb.save(buf, format='JPEG', quality=config.IMAGE_QUALITY)
|
|
img_data = base64.b64encode(buf.getvalue()).decode("utf-8")
|
|
|
|
# Gather info
|
|
mtime_rel = format_relative_time(stats.st_mtime)
|
|
mtime_abs = datetime.datetime.fromtimestamp(stats.st_mtime).strftime('%Y-%m-%d %H:%M')
|
|
size_kb = stats.st_size / 1024
|
|
|
|
# Prepend index if this image was selected via index/range
|
|
idx_prefix = f"[Index: {index_label}] " if index_label is not None else ""
|
|
|
|
info_text = (
|
|
f"{idx_prefix}File: {file_path.name}\n"
|
|
f"Resolution: {orig_w}x{orig_h}\n"
|
|
f"Size: {size_kb:.1f}KB\n"
|
|
f"Modified: {mtime_rel} ({mtime_abs})"
|
|
)
|
|
|
|
results.append({"type": "text", "text": info_text})
|
|
results.append({"type": "image", "data": img_data, "mimeType": "image/jpeg"})
|
|
|
|
except Exception as e:
|
|
results.append({"type": "text", "text": f"Error processing {file_path.name}: {e}"})
|
|
|
|
return results
|