#!/usr/bin/env python3 """ Static checker + clean-room renderer + image comparator for the Agents' Last Exam Mitsuba Cornell-glass reconstruction task. Scale transforms are allowed; textures are forbidden. Default layout: input/reference.png input/scripts/render.py output/scene.xml no additional output files are allowed beyond scene.xml and evaluator-generated files Writes: output/render.exr output/render.png output/metrics.json """ from __future__ import annotations import argparse import json import os from pathlib import Path import re import shutil import struct import subprocess import sys import tempfile import zlib import numpy as np import xml.etree.ElementTree as ET WIDTH = 1024 HEIGHT = 1024 SCENE_SIZE_LIMIT = 150_000 MAX_SHAPES = 50 MAX_BSDFS = 20 MIN_EMITTERS = 1 MAX_EMITTERS = 4 GLOBAL_PASS = 0.0085 REGION_PASS = 0.014 ROBUST_CLIP_PERCENTILE = 99.5 RENDER_TIMEOUT = int(os.environ.get("EVAL_RENDER_TIMEOUT", "1800")) ALLOWED_SHAPES = {"sphere", "rectangle", "cube", "disk"} ALLOWED_BSDFS = {"diffuse", "dielectric", "conductor", "roughconductor", "plastic", "roughplastic", "twosided"} ALLOWED_EMITTERS = {"area", "spot", "point", "constant"} FORBIDDEN_XML_TYPES = { "bitmap", "checkerboard", "gridtexture", "mesh", "obj", "ply", "serialized", "envmap", "mesh_attribute", "shapegroup", "instance" } FORBIDDEN_TAGS = {"include", "mesh_attribute"} FORBIDDEN_WORDS = [ "base64", "data:image", ".png", ".jpg", ".jpeg", ".exr", ".hdr", ".pfm", ".obj", ".ply", ".serialized", "filename", "bitmap", "checkerboard", "gridtexture", "wireframe" ] B64_RE = re.compile(r"[A-Za-z0-9+/]{256,}={0,2}") LONG_NUMERIC_RE = re.compile(r"(?:[-+]?\d*\.?\d+(?:[eE][-+]?\d+)?[\s,]+){80,}") # Broad crops in 1024x1024 sRGB image space. They make it difficult to match # only the average color while missing the key geometry/material regions. REGIONS = { "left_stacked_blocks": (180, 310, 600, 900), "glass_cuboid_and_top_block": (430, 550, 900, 1000), "ceiling_light_and_upper_walls": (0, 0, 1024, 380), "floor_under_glass": (420, 820, 910, 1024), } def add_error(errors: list[str], msg: str) -> None: errors.append(msg) def png_read_rgb8(path: Path) -> np.ndarray: """Minimal PNG reader: 8-bit grayscale/RGB/RGBA, non-interlaced.""" data = path.read_bytes() if not data.startswith(b"\x89PNG\r\n\x1a\n"): raise ValueError(f"not a PNG: {path}") pos = 8 width = height = bit_depth = color_type = interlace = None idat = bytearray() while pos < len(data): if pos + 8 > len(data): raise ValueError("truncated PNG") length = struct.unpack(">I", data[pos:pos+4])[0] tag = data[pos+4:pos+8] chunk = data[pos+8:pos+8+length] pos += 12 + length if tag == b"IHDR": width, height, bit_depth, color_type, _comp, _filter, interlace = struct.unpack(">IIBBBBB", chunk) elif tag == b"IDAT": idat.extend(chunk) elif tag == b"IEND": break if bit_depth != 8 or interlace != 0: raise ValueError("only 8-bit non-interlaced PNGs are supported") channels = {0: 1, 2: 3, 4: 2, 6: 4}.get(color_type) if channels is None: raise ValueError(f"unsupported PNG color type: {color_type}") raw = zlib.decompress(bytes(idat)) stride = width * channels rows = [] prev = np.zeros(stride, dtype=np.uint8) i = 0 def paeth(a, b, c): p = a + b - c pa = abs(p - a); pb = abs(p - b); pc = abs(p - c) if pa <= pb and pa <= pc: return a if pb <= pc: return b return c for _y in range(height): filt = raw[i] i += 1 scan = np.frombuffer(raw[i:i+stride], dtype=np.uint8).copy() i += stride recon = np.empty_like(scan) bpp = channels for x in range(stride): left = int(recon[x-bpp]) if x >= bpp else 0 up = int(prev[x]) ul = int(prev[x-bpp]) if x >= bpp else 0 val = int(scan[x]) if filt == 0: out = val elif filt == 1: out = val + left elif filt == 2: out = val + up elif filt == 3: out = val + ((left + up) // 2) elif filt == 4: out = val + paeth(left, up, ul) else: raise ValueError(f"bad PNG filter {filt}") recon[x] = out & 255 rows.append(recon.reshape(width, channels)) prev = recon arr = np.stack(rows, axis=0) if color_type == 0: arr = np.repeat(arr[:, :, :1], 3, axis=2) elif color_type == 4: arr = np.repeat(arr[:, :, :1], 3, axis=2) elif color_type == 6: arr = arr[:, :, :3] return arr.astype(np.float32) / 255.0 def rel_mse(a: np.ndarray, b: np.ndarray) -> float: return float(np.mean(((a - b) ** 2) / (b ** 2 + 1e-2))) def robust_rel_mse(a: np.ndarray, b: np.ndarray) -> float: """relMSE with the highest-error firefly pixels clipped. This keeps the task geometric/material, not a random-noise lottery. The unclipped metric is still reported in metrics.json for debugging. """ e = ((a - b) ** 2) / (b ** 2 + 1e-2) if not np.isfinite(e).all(): return float("inf") cap = np.percentile(e, ROBUST_CLIP_PERCENTILE) return float(np.mean(np.minimum(e, cap))) def read_text(path: Path) -> str: return path.read_text(encoding="utf-8", errors="replace") def looks_encoded(text: str) -> bool: return bool(B64_RE.search(text) or LONG_NUMERIC_RE.search(text)) def child_value(parent: ET.Element, tag: str, name: str) -> str | None: for c in parent.findall(tag): if c.attrib.get("name") == name: return c.attrib.get("value") return None def check_camera(root: ET.Element, errors: list[str], info: dict) -> None: """Lightweight camera sanity check. The submitted camera is used by render.py. Exact pose/fov recovery is enforced by the image metric rather than by fragile XML-form checks, because equivalent Mitsuba cameras can be written with lookat, matrix transforms, omitted defaults, etc. """ sensors = root.findall("sensor") if len(sensors) != 1: add_error(errors, f"scene must contain exactly one top-level sensor, got {len(sensors)}") return sensor = sensors[0] stype = sensor.attrib.get("type") info["sensor_type"] = stype if stype != "perspective": add_error(errors, "sensor must be a perspective camera") fov_text = child_value(sensor, "float", "fov") if fov_text is not None: try: fov = float(fov_text) info["camera_fov"] = fov if not (10.0 <= fov <= 80.0): add_error(errors, "camera fov must be in a plausible perspective range") except ValueError: add_error(errors, "camera fov must be numeric if specified") # render.py replaces film and sampler settings, so their values are not checked. def static_check(output_dir: Path) -> tuple[bool, list[str], dict]: errors: list[str] = [] info: dict = {} scene_path = output_dir / "scene.xml" tex_dir = output_dir / "textures" if tex_dir.exists(): add_error(errors, "output/textures is not allowed for this task") if not scene_path.exists(): add_error(errors, "missing output/scene.xml") return False, errors, info # Only allow scene.xml and evaluator-generated files. No plugins/textures are needed or allowed. Transform scale tags are allowed. allowed_root = {"scene.xml", "render.exr", "render.png", "metrics.json"} for p in output_dir.rglob("*"): if p.is_dir(): continue rel = p.relative_to(output_dir) ok = rel.name in allowed_root and len(rel.parts) == 1 if not ok: add_error(errors, f"forbidden extra file in output: {rel}") scene_text = read_text(scene_path) if len(scene_text.encode("utf-8")) > SCENE_SIZE_LIMIT: add_error(errors, "scene.xml over 150 kB") if looks_encoded(scene_text): add_error(errors, "possible embedded/encoded pixel or numeric array in scene.xml") low_scene = scene_text.lower() for w in FORBIDDEN_WORDS: if w.lower() in low_scene: add_error(errors, f"forbidden reference/word in scene.xml: {w}") try: root = ET.fromstring(scene_text) except ET.ParseError as e: add_error(errors, f"XML parse error: {e}") return False, errors, info if root.tag != "scene": add_error(errors, "root element must be ") check_camera(root, errors, info) shape_nodes = root.findall(".//shape") bsdf_nodes = root.findall(".//bsdf") emitter_nodes = root.findall(".//emitter") texture_nodes = root.findall(".//texture") material_bsdf_nodes = [b for b in bsdf_nodes if b.attrib.get("type") != "twosided"] info.update({ "shape_count": len(shape_nodes), "bsdf_count": len(material_bsdf_nodes), "bsdf_wrapper_count": len(bsdf_nodes) - len(material_bsdf_nodes), "emitter_count": len(emitter_nodes), "texture_count": len(texture_nodes), }) if len(shape_nodes) > MAX_SHAPES: add_error(errors, f"too many shapes: {len(shape_nodes)} > {MAX_SHAPES}") if len(material_bsdf_nodes) > MAX_BSDFS: add_error(errors, f"too many BSDFs: {len(material_bsdf_nodes)} > {MAX_BSDFS}") if not (MIN_EMITTERS <= len(emitter_nodes) <= MAX_EMITTERS): add_error(errors, f"emitter count must be {MIN_EMITTERS}..{MAX_EMITTERS}, got {len(emitter_nodes)}") for node in root.iter(): if node.tag in FORBIDDEN_TAGS: add_error(errors, f"forbidden tag: <{node.tag}>") typ = node.attrib.get("type") if typ and typ.lower() in FORBIDDEN_XML_TYPES: add_error(errors, f"forbidden plugin/type: {typ}") if any(k.lower() == "filename" for k in node.attrib): add_error(errors, f"filename attribute forbidden on <{node.tag}>") for sh in shape_nodes: typ = sh.attrib.get("type") if typ not in ALLOWED_SHAPES: add_error(errors, f"shape type not allowed: {typ}") for b in bsdf_nodes: typ = b.attrib.get("type") if typ not in ALLOWED_BSDFS: add_error(errors, f"BSDF type not allowed: {typ}") if typ == "twosided": children = [c for c in list(b) if c.tag == "bsdf"] if len(children) != 1: add_error(errors, "twosided BSDF must wrap exactly one BSDF child") elif children[0].attrib.get("type") == "twosided": add_error(errors, "nested twosided BSDF is not allowed") for e in emitter_nodes: typ = e.attrib.get("type") if typ not in ALLOWED_EMITTERS: add_error(errors, f"emitter type not allowed: {typ}") if typ == "constant": for desc in e.iter(): if desc.tag == "texture" or "filename" in desc.attrib: add_error(errors, "constant environment may not reference texture or file") for tex in texture_nodes: typ = tex.attrib.get("type") add_error(errors, f"textures are not allowed in this task: {typ}") return len(errors) == 0, errors, info def run_clean_render(script_path: Path, output_dir: Path) -> tuple[bool, list[str], Path | None]: errors: list[str] = [] with tempfile.TemporaryDirectory(prefix="ale_mitsuba_clean_") as td: clean = Path(td) clean_out = clean / "output" clean_out.mkdir() shutil.copy2(output_dir / "scene.xml", clean_out / "scene.xml") # No additional files are copied into the clean room. cmd = [ sys.executable, str(script_path), "--scene", str(clean_out / "scene.xml"), "--out-exr", str(clean_out / "render.exr"), "--out-png", str(clean_out / "render.png"), ] try: proc = subprocess.run( cmd, cwd=str(clean), text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, timeout=RENDER_TIMEOUT, ) except subprocess.TimeoutExpired: add_error(errors, f"render timeout after {RENDER_TIMEOUT}s") return False, errors, None if proc.returncode != 0: add_error(errors, "render.py failed:\n" + proc.stdout[-4000:]) return False, errors, None png = clean_out / "render.png" exr = clean_out / "render.exr" if not png.exists() or not exr.exists(): add_error(errors, "render output missing render.png or render.exr") return False, errors, None shutil.copy2(exr, output_dir / "render.exr") shutil.copy2(png, output_dir / "render.png") return True, errors, output_dir / "render.png" def compare(render_png: Path, reference_png: Path) -> tuple[dict, list[str]]: errors: list[str] = [] r = png_read_rgb8(render_png) ref = png_read_rgb8(reference_png) metrics: dict = {} if r.shape != (HEIGHT, WIDTH, 3): add_error(errors, f"render size is {r.shape}, expected {(HEIGHT, WIDTH, 3)}") if ref.shape != (HEIGHT, WIDTH, 3): add_error(errors, f"reference size is {ref.shape}, expected {(HEIGHT, WIDTH, 3)}") if errors: return metrics, errors if not np.isfinite(r).all() or not np.isfinite(ref).all(): add_error(errors, "NaN or Inf in image data") return metrics, errors metrics["global_rel_mse_raw"] = rel_mse(r, ref) metrics["global_rel_mse"] = robust_rel_mse(r, ref) region_metrics = {} raw_region_metrics = {} for name, (x0, y0, x1, y1) in REGIONS.items(): rr = r[y0:y1, x0:x1, :] rf = ref[y0:y1, x0:x1, :] region_metrics[name] = robust_rel_mse(rr, rf) raw_region_metrics[name] = rel_mse(rr, rf) metrics["regions"] = region_metrics metrics["regions_raw"] = raw_region_metrics return metrics, errors def main() -> int: ap = argparse.ArgumentParser() ap.add_argument("--input-dir", default="input") ap.add_argument("--output-dir", default="output") ap.add_argument("--reference", default=None) ap.add_argument("--render-script", default=None) args = ap.parse_args() # Resolve paths before run_clean_render() changes the subprocess working # directory to the isolated clean room. This keeps the documented default # invocation functional while still copying only scene.xml into that room. input_dir = Path(args.input_dir).resolve() output_dir = Path(args.output_dir).resolve() reference = Path(args.reference).resolve() if args.reference else input_dir / "reference.png" render_script = Path(args.render_script).resolve() if args.render_script else input_dir / "scripts" / "render.py" metrics_path = output_dir / "metrics.json" output_dir.mkdir(parents=True, exist_ok=True) result = { "pass": False, "static_ok": False, "render_ok": False, "compare_ok": False, "errors": [], "static_info": {}, "global_rel_mse": None, "regions": {}, "thresholds": {"global": GLOBAL_PASS, "regions": REGION_PASS, "robust_clip_percentile": ROBUST_CLIP_PERCENTILE}, } static_ok, static_errors, static_info = static_check(output_dir) result["static_ok"] = static_ok result["static_info"] = static_info result["errors"].extend(static_errors) if static_ok: if not reference.exists(): result["errors"].append(f"missing reference image: {reference}") elif not render_script.exists(): result["errors"].append(f"missing render script: {render_script}") else: render_ok, render_errors, render_png = run_clean_render(render_script, output_dir) result["render_ok"] = render_ok result["errors"].extend(render_errors) if render_ok and render_png is not None: comp_metrics, comp_errors = compare(render_png, reference) result["errors"].extend(comp_errors) result.update(comp_metrics) if not comp_errors: global_ok = comp_metrics["global_rel_mse"] < GLOBAL_PASS regions_ok = all(v < REGION_PASS for v in comp_metrics["regions"].values()) result["compare_ok"] = global_ok and regions_ok if not global_ok: result["errors"].append(f"global relMSE too high: {comp_metrics['global_rel_mse']:.6f}") for k, v in comp_metrics["regions"].items(): if v >= REGION_PASS: result["errors"].append(f"region relMSE too high: {k}={v:.6f}") result["pass"] = bool(result["static_ok"] and result["render_ok"] and result["compare_ok"]) metrics_path.write_text(json.dumps(result, indent=2, sort_keys=True), encoding="utf-8") print(json.dumps(result, indent=2, sort_keys=True)) return 0 if result["pass"] else 2 if __name__ == "__main__": raise SystemExit(main())