import argparse
import math
import re
import xml.etree.ElementTree as ET
from pathlib import Path
from perceptron import collect_annotations, resolve_asset_idx
from PIL import Image, ImageDraw, ImageOps
def parse_boxes(text):
annotations = collect_annotations(text, strict=True)
tag_names = ("point_box", "point", "polygon", "clip", "collection", "track")
coordinate = r"\(\s*[0-9]+\s*,\s*[0-9]+\s*\)"
box_body = rf"\s*{coordinate}\s+{coordinate}\s*"
for segment in annotations.parsed:
if segment["kind"] == "text":
remaining = segment["text"]
tag_start = re.search(r"</?(?:point_box|point|polygon|clip|collection|track)\b", remaining, re.I)
partial = re.search(r"</?([a-z_]+)$", remaining.rstrip(), re.I)
if tag_start or (partial and any(tag.startswith(partial[1].lower()) for tag in tag_names)):
raise ValueError("Unparsed annotation markup remains in the answer")
continue
span = segment["span"]
try:
root = ET.fromstring(text[span["start"]:span["end"]])
except ET.ParseError as error:
raise ValueError("Annotation markup must be valid XML") from error
for node in root.iter():
if node.tag not in {"point_box", "collection"}:
raise ValueError("Expected static image boxes only")
if set(node.attrib) - {"mention", "asset_idx"}:
raise ValueError("Only mention and asset_idx attributes are supported")
selector = node.get("asset_idx")
if selector is not None and re.fullmatch(r"[0-9]+", selector) is None:
raise ValueError("asset_idx must be a non-negative integer when present")
if node.tag == "point_box":
if len(node) or re.fullmatch(box_body, node.text or "") is None:
raise ValueError("A box must contain exactly two coordinate pairs")
elif (node.text or "").strip() or any((child.tail or "").strip() for child in node):
raise ValueError("Collections must contain only annotation elements")
if annotations.tracks or annotations.points or annotations.polygons or annotations.clips:
raise ValueError("Expected static image boxes only")
for box in annotations.boxes:
if box.t is not None:
raise ValueError("Timestamped boxes need a video frame")
x1, y1 = box.top_left.x, box.top_left.y
x2, y2 = box.bottom_right.x, box.bottom_right.y
if not all(math.isfinite(value) for value in (x1, y1, x2, y2)):
raise ValueError("Coordinates must be finite")
if not (0 <= x1 < x2 <= 1000 and 0 <= y1 < y2 <= 1000):
raise ValueError("Box corners must be ordered and within 0–1000")
return annotations.boxes
def render(image_path, response_path, asset_idx, last_asset_idx, output_path):
if last_asset_idx < 0:
raise ValueError("last_asset_idx must be non-negative")
if not 0 <= asset_idx <= last_asset_idx:
raise ValueError("asset_idx must be between 0 and last_asset_idx")
boxes = parse_boxes(Path(response_path).read_text(encoding="utf-8"))
selected = [
box for box in boxes
if resolve_asset_idx(box, n_assets=last_asset_idx + 1) == asset_idx
]
with Image.open(image_path) as source:
canvas = ImageOps.exif_transpose(source).convert("RGB")
width, height = canvas.size
draw = ImageDraw.Draw(canvas)
for box in selected:
x1, y1 = box.top_left.x, box.top_left.y
x2, y2 = box.bottom_right.x, box.bottom_right.y
pixels = (
min(width - 1, round(x1 * width / 1000)),
min(height - 1, round(y1 * height / 1000)),
min(width - 1, round(x2 * width / 1000)),
min(height - 1, round(y2 * height / 1000)),
)
draw.rectangle(pixels, outline="red", width=3)
draw.text((pixels[0], max(0, pixels[1] - 12)), box.mention or "object", fill="red")
canvas.save(output_path)
print(f"Drew {len(selected)} boxes for asset {asset_idx} in {output_path}")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("image", type=Path)
parser.add_argument("response", type=Path)
parser.add_argument("--asset-idx", type=int, required=True)
parser.add_argument("--last-asset-idx", type=int, required=True,
help="Index of the last asset available when this answer was produced")
parser.add_argument("--output", type=Path, default=Path("annotated.png"))
args = parser.parse_args()
render(args.image, args.response, args.asset_idx, args.last_asset_idx, args.output)
