wehub-resource-sync
346 lines
12 KiB
Python
346 lines
12 KiB
Python
"""Pixel Query: render text queries as images.
|
|
|
|
This module converts text queries to pixel images so VLMs receive the query
|
|
as a visual input instead of text tokens. This tests whether the model can
|
|
"read" the question from an image and still perform retrieval + answering.
|
|
|
|
Usage:
|
|
renderer = PixelQueryRenderer(output_dir="pixel_queries", font_size=16, img_width=600)
|
|
img_path = renderer.render(example_id, query_text)
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Default font paths (tried in order)
|
|
_FONT_CANDIDATES = [
|
|
"/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
|
|
"/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf",
|
|
"/usr/share/fonts/truetype/noto/NotoSans-Regular.ttf",
|
|
]
|
|
|
|
|
|
def _find_font(font_path: str | None = None) -> str:
|
|
"""Find a usable TrueType font on this system."""
|
|
if font_path and os.path.exists(font_path):
|
|
return font_path
|
|
for candidate in _FONT_CANDIDATES:
|
|
if os.path.exists(candidate):
|
|
return candidate
|
|
raise FileNotFoundError(
|
|
"No suitable TTF font found. Install dejavu or liberation fonts, "
|
|
"or pass font_path explicitly."
|
|
)
|
|
|
|
|
|
def _wrap_text_by_pixel_width(
|
|
text: str, font: ImageFont.FreeTypeFont, max_width: int
|
|
) -> list[str]:
|
|
"""Word-wrap text so each line fits within *max_width* pixels."""
|
|
words = text.split()
|
|
lines: list[str] = []
|
|
current_line = ""
|
|
|
|
for word in words:
|
|
test_line = f"{current_line} {word}".strip()
|
|
bbox = font.getbbox(test_line)
|
|
line_width = bbox[2] - bbox[0]
|
|
if line_width <= max_width:
|
|
current_line = test_line
|
|
else:
|
|
if current_line:
|
|
lines.append(current_line)
|
|
current_line = word
|
|
if current_line:
|
|
lines.append(current_line)
|
|
|
|
return lines
|
|
|
|
|
|
class PixelQueryRenderer:
|
|
"""Renders text queries as small, clear PNG images.
|
|
|
|
Images are cached on disk so each query is only rendered once.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
output_dir: str = "pixel_queries",
|
|
font_path: str | None = None,
|
|
font_size: int = 16,
|
|
img_width: int = 600,
|
|
padding_x: int = 16,
|
|
padding_y: int = 12,
|
|
line_spacing: int = 4,
|
|
):
|
|
self.output_dir = output_dir
|
|
self.font_path = _find_font(font_path)
|
|
self.font_size = font_size
|
|
self.img_width = img_width
|
|
self.padding_x = padding_x
|
|
self.padding_y = padding_y
|
|
self.line_spacing = line_spacing
|
|
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
self.font = ImageFont.truetype(self.font_path, self.font_size)
|
|
logger.info(
|
|
f"PixelQueryRenderer: dir={output_dir}, font={os.path.basename(self.font_path)}, "
|
|
f"size={font_size}, width={img_width}"
|
|
)
|
|
|
|
def _render_image(self, text: str) -> Image.Image:
|
|
"""Render *text* to a PIL Image (white bg, black text)."""
|
|
max_text_width = self.img_width - 2 * self.padding_x
|
|
lines = _wrap_text_by_pixel_width(text, self.font, max_text_width)
|
|
|
|
# Measure line height using a reference string
|
|
line_height = self.font.getbbox("Ay")[3] - self.font.getbbox("Ay")[1]
|
|
total_text_height = (
|
|
len(lines) * line_height + (len(lines) - 1) * self.line_spacing
|
|
)
|
|
img_height = total_text_height + 2 * self.padding_y
|
|
|
|
img = Image.new("RGB", (self.img_width, img_height), color=(255, 255, 255))
|
|
draw = ImageDraw.Draw(img)
|
|
|
|
y = self.padding_y
|
|
for line in lines:
|
|
draw.text((self.padding_x, y), line, font=self.font, fill=(0, 0, 0))
|
|
y += line_height + self.line_spacing
|
|
|
|
return img
|
|
|
|
def render(self, example_id: str, query_text: str) -> str:
|
|
"""Render a query and return the path to the saved PNG.
|
|
|
|
If the image already exists on disk it is *not* re-rendered.
|
|
"""
|
|
out_path = os.path.join(self.output_dir, f"{example_id}_query.png")
|
|
if os.path.exists(out_path):
|
|
return out_path
|
|
|
|
img = self._render_image(query_text)
|
|
img.save(out_path)
|
|
logger.debug(f"Rendered pixel query: {out_path} ({img.size[0]}x{img.size[1]})")
|
|
return out_path
|
|
|
|
def render_all(self, examples: list[dict]) -> dict[str, str]:
|
|
"""Batch-render pixel queries for a list of examples.
|
|
|
|
Args:
|
|
examples: List of dicts with at least ``id`` and ``problem`` keys.
|
|
|
|
Returns:
|
|
Dict mapping example_id → pixel query image path.
|
|
"""
|
|
id_to_path: dict[str, str] = {}
|
|
rendered, cached = 0, 0
|
|
for ex in examples:
|
|
eid = ex["id"]
|
|
path = os.path.join(self.output_dir, f"{eid}_query.png")
|
|
if os.path.exists(path):
|
|
cached += 1
|
|
else:
|
|
img = self._render_image(ex["problem"])
|
|
img.save(path)
|
|
rendered += 1
|
|
id_to_path[eid] = path
|
|
|
|
logger.info(
|
|
f"PixelQueryRenderer: {rendered} rendered, {cached} cached, "
|
|
f"{rendered + cached} total in {self.output_dir}"
|
|
)
|
|
return id_to_path
|
|
|
|
|
|
class QueryImageTextRenderer:
|
|
"""Renders query text + query image together into a single card image.
|
|
|
|
Reuses PixelQueryRenderer for text rendering. Layout: image on top, text below
|
|
(similar to VQA task cards). Text is centered and uses a larger font.
|
|
Images are cached on disk.
|
|
|
|
Usage:
|
|
renderer = QueryImageTextRenderer(output_dir="query_cards", tiles_dir="tiles/evqa")
|
|
path = renderer.render(example_id, query_text, query_image_path)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
output_dir: str = "query_cards",
|
|
tiles_dir: str = "tiles/evqa",
|
|
font_path: str | None = None,
|
|
font_size: int = 22,
|
|
card_width: int = 600,
|
|
padding_x: int = 24,
|
|
padding_y: int = 20,
|
|
line_spacing: int = 6,
|
|
image_padding: int = 16,
|
|
text_section_padding: int = 24,
|
|
border_radius: int = 12,
|
|
):
|
|
self.output_dir = output_dir
|
|
self.tiles_dir = tiles_dir
|
|
self.card_width = card_width
|
|
self.image_padding = image_padding
|
|
self.text_section_padding = text_section_padding
|
|
self.border_radius = border_radius
|
|
|
|
font_path_resolved = _find_font(font_path)
|
|
self.font_path = font_path_resolved
|
|
self.font_size = font_size
|
|
self.font = ImageFont.truetype(font_path_resolved, font_size)
|
|
self.padding_x = padding_x
|
|
self.padding_y = padding_y
|
|
self.line_spacing = line_spacing
|
|
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
logger.info(
|
|
f"QueryImageTextRenderer: dir={output_dir}, card_width={card_width}, "
|
|
f"font_size={font_size}"
|
|
)
|
|
|
|
def _render_query_text_centered(
|
|
self, text: str, width: int | None = None
|
|
) -> Image.Image:
|
|
"""Render query text centered, reusing wrap logic from PixelQueryRenderer."""
|
|
width = width or self.card_width
|
|
max_text_width = width - 2 * self.padding_x
|
|
|
|
# Calculate dynamic font size based on width
|
|
# Ratio: width / 30 seems reasonable (600px -> 20px, 1500px -> 50px)
|
|
dynamic_font_size = max(22, int(width / 30))
|
|
|
|
# Use dynamic font if size is different from default
|
|
if dynamic_font_size != self.font_size:
|
|
try:
|
|
font = ImageFont.truetype(self.font_path, dynamic_font_size)
|
|
except Exception:
|
|
font = self.font # Fallback
|
|
else:
|
|
font = self.font
|
|
|
|
lines = _wrap_text_by_pixel_width(text, font, max_text_width)
|
|
|
|
line_height = font.getbbox("Ay")[3] - font.getbbox("Ay")[1]
|
|
# Scale spacing proportionally
|
|
spacing = max(4, int(self.line_spacing * (dynamic_font_size / self.font_size)))
|
|
|
|
total_text_height = len(lines) * line_height + (len(lines) - 1) * spacing
|
|
# Scale padding proportionally
|
|
padding_y = max(
|
|
self.padding_y, int(self.padding_y * (dynamic_font_size / self.font_size))
|
|
)
|
|
text_height = total_text_height + 2 * padding_y
|
|
|
|
img = Image.new("RGB", (width, text_height), color=(255, 255, 255))
|
|
draw = ImageDraw.Draw(img)
|
|
|
|
y = padding_y
|
|
for line in lines:
|
|
bbox = font.getbbox(line)
|
|
line_w = bbox[2] - bbox[0]
|
|
x = (width - line_w) // 2
|
|
draw.text((x, y), line, font=font, fill=(50, 50, 50))
|
|
y += line_height + spacing
|
|
|
|
return img
|
|
|
|
def render(
|
|
self,
|
|
example_id: str,
|
|
query_text: str,
|
|
query_image_path: str | None,
|
|
force: bool = False,
|
|
) -> str:
|
|
"""Render query image + text into one card, save to disk.
|
|
|
|
Layout: image on top, text below (same for iNaturalist and Landmarks).
|
|
|
|
Args:
|
|
example_id: Example identifier for filename.
|
|
query_text: The question text to render below the image.
|
|
query_image_path: Path to the query image (iNaturalist or Landmark photo).
|
|
If None or file missing, only the text is rendered.
|
|
force: If True, re-render even if output exists (e.g. when images were added later).
|
|
|
|
Returns:
|
|
Path to the saved PNG.
|
|
"""
|
|
out_path = os.path.join(self.output_dir, f"{example_id}_query_card.png")
|
|
if os.path.exists(out_path) and not force:
|
|
return out_path
|
|
|
|
# Load query image
|
|
if query_image_path and os.path.exists(query_image_path):
|
|
try:
|
|
query_img = Image.open(query_image_path).convert("RGB")
|
|
# Resize if too large to ensure font size is readable and image isn't massive
|
|
max_dim = 1536 # Standard reasonable max dimension
|
|
if max(query_img.size) > max_dim:
|
|
ratio = max_dim / max(query_img.size)
|
|
new_size = (
|
|
int(query_img.width * ratio),
|
|
int(query_img.height * ratio),
|
|
)
|
|
query_img = query_img.resize(new_size, Image.Resampling.LANCZOS)
|
|
except Exception as e:
|
|
logger.warning(f"Failed to load query image {query_image_path}: {e}")
|
|
query_img = None
|
|
else:
|
|
query_img = None
|
|
|
|
# Card width adapts to image: expand if image is wider than card_width
|
|
if query_img is not None:
|
|
effective_width = max(
|
|
self.card_width, query_img.width + 2 * self.image_padding
|
|
)
|
|
else:
|
|
effective_width = self.card_width
|
|
|
|
# Render query text at effective width
|
|
text_img = self._render_query_text_centered(query_text, effective_width)
|
|
|
|
# Compose: image on top, text below with padding for balanced look
|
|
if query_img is not None:
|
|
img_section_height = query_img.height + 2 * self.image_padding
|
|
else:
|
|
img_section_height = 0
|
|
|
|
text_section_height = text_img.height + 2 * self.text_section_padding
|
|
total_height = img_section_height + text_section_height
|
|
|
|
card = Image.new("RGB", (effective_width, total_height), color=(255, 255, 255))
|
|
ImageDraw.Draw(card)
|
|
|
|
y_offset = 0
|
|
if query_img is not None:
|
|
x_center = (effective_width - query_img.width) // 2
|
|
card.paste(query_img, (x_center, self.image_padding))
|
|
y_offset = img_section_height
|
|
|
|
# Center text block in its section
|
|
text_y = y_offset + self.text_section_padding
|
|
card.paste(text_img, (0, text_y))
|
|
|
|
# Optional: rounded corners (simplified - draw white rounded rect overlay)
|
|
if self.border_radius > 0:
|
|
# Create mask for rounded corners
|
|
mask = Image.new("L", card.size, 255)
|
|
m_draw = ImageDraw.Draw(mask)
|
|
m_draw.rounded_rectangle(
|
|
(0, 0, card.width - 1, card.height - 1),
|
|
radius=self.border_radius,
|
|
fill=255,
|
|
outline=0,
|
|
)
|
|
# For simple output we keep the card as-is; full rounded crop would need alpha
|
|
pass
|
|
|
|
card.save(out_path)
|
|
logger.debug(f"Rendered query card: {out_path} ({card.size[0]}x{card.size[1]})")
|
|
return out_path
|