| """IR rooms -> filtered, deduplicated, leakage-safe SFT files (spec: FastFill I/O design review, 2026-09-24). |
| |
| python -m fastfill.build --ir /Volumes/harddisk/fastfill_ir --out data/fastfill_v1 [--sources ...] [--constraint_frac 0.3] |
| |
| Per room: prep (reject or clean, see `prep`; structure and unlabelled floor boxes become `fixed` context, never |
| placed) -> cross-source dedup by group -> split by group/content |
| -> train: rot90 augmentation -> canonical order/ids -> optional constraints (S2) -> field dropout -> messages |
| dev/test: canonical -> messages (no augmentation), IR written to {dev,test}_rooms.jsonl for fastfill.evaluate, |
| test_constrained_rooms.jsonl adds 1-4 GT-true constraints per test room (rooms without any are skipped). |
| Eval-only sources (meta.eval_only) go to test with their whole split component (split.assign_splits). |
| stats.json counts every rejection reason, quality flag, dropped object and inferred vs evidence support per source. |
| Rejections are only what makes a training target impossible (no verified front, missing boxes, bad numbers, no |
| boundary, no placeable object); everything else (reference objects through walls, same-category overlaps, tiny rooms, |
| unexplained dropped floor area, single-object rooms, snapped or tilted boxes) is a FLAG written next to the sample. |
| """ |
| import argparse |
| import collections |
| import hashlib |
| import json |
| import math |
| import os |
| import random |
| import re |
| import tempfile |
| from contextlib import ExitStack |
| from pathlib import Path |
|
|
| import numpy as np |
| import shapely |
| from shapely import STRtree |
| from shapely.geometry import LineString, Point, Polygon |
| from shapely.ops import unary_union |
|
|
| from fastfill.anchors import (DIVIDERS, FINISH, FLAT, FLOOR_Z, FURNITURE_SUNK_Z, SOFT_COVERS, SOFT_HANGINGS, WALL_ART, |
| annotate, closed, fixed_geometry, hanging, is_structure) |
| from fastfill.scene import (GENERIC, canonical, clean_boundary, footprint, head, messages, norm_cat, num, rot90, |
| shown_size, slug) |
| from fastfill.split import assign_splits |
| from fastfill.validate import check, holds, support_ok |
|
|
| |
| FLUSH = re.compile(r"^(doors?|doorframes?|windows?|curtains?|drapes?|blinds?|frames?|openings?|baseboards?|rails?|" |
| r"railings?|radiators?|beams?|shades?|screening|designs?|molds?|moldings?|panels?|paneling|combination|" |
| r"walls?|floors?|ceilings?|skirting|trim|tiles?|boards?)$") |
|
|
|
|
| WALL_BAND_M = 0.10 |
| THICK_WALL_M = 0.30 |
| MIN_AREA_M2 = 0.5 |
| INTRUSION_M2 = 0.05 |
| |
| SOFT_HEADS = {"plant", "plants", "flower", "flowers", "pillow", "pillows", "cushion", "towel", "book", "books", |
| "ornament", "decoration", "vase", "toy", "toys", "clothes", "bag", "lamp", "light"} |
| COMPOSITE = re.compile(r"\b(combination|l shaped|composite|kids)\b") |
|
|
|
|
| def _words(category): |
| return re.sub(r"[^a-z]+", " ", (category or "").lower()).split() |
|
|
|
|
| def _flush(o, inner, polygon): |
| """A dropped floor box that leaves no hole: a flat covering, a soft cover or curtain, or a structural |
| opening/finish ('bed frame' is not). In a polygon room (the boundary is the walls) such a box must also keep to |
| the walls, under INTRUSION_M2 inside `inner` (the room shrunk by WALL_BAND_M): an open door leaf, a shower |
| enclosure or a 'wall' box standing in the room is an obstacle whatever its label. Source-flagged architecture |
| (stairs/pillars are shown as fixed boxes when they qualify) needs only the geometry there. In a hull room (a scan's |
| floor hull, walls stand inside it) the name decides, flagged boxes may also pass by geometry.""" |
| w = _words(o["category"]) |
| if head(o) in FLAT or head(o) in SOFT_COVERS or SOFT_HANGINGS & set(w) or \ |
| ((head(o) in FINISH or head(o).rstrip("s") in FINISH) and not DIVIDERS & set(w)): |
| return True |
| stays = Polygon(footprint(o)).intersection(inner).area < INTRUSION_M2 |
| if o.get("structure"): |
| return stays if polygon else \ |
| (any(FLUSH.match(x) for x in w) and not {"partition", "screen", "clothes"} & set(w)) or stays |
| if not (polygon is False or stays): |
| return False |
| while len(w) > 1 and w[-1] in ("floor", "wall", "ceiling"): |
| w.pop() |
| return bool(w) and is_structure(o["category"]) and not {"partition", "screen", "clothes"} & set(w) \ |
| and bool(FLUSH.match(w[-1])) |
|
|
|
|
| FRONT = {"chair", "armchair", "seat", "sofa", "couch", "loveseat", "bed", "desk", "tv", "television", "monitor", "toilet", |
| "piano"} |
| PRIORITY = ["on", "between", "faces", "against_wall", "near"] |
|
|
|
|
| def _drop_with_dependants(objs, bad): |
| bad = set(bad) |
| while True: |
| more = {o["id"] for o in objs if o.get("parent") in bad} - bad |
| if not more: |
| return [o for o in objs if o["id"] not in bad] |
| bad |= more |
|
|
|
|
| FLAT_Z = 0.15 |
|
|
|
|
| def _finite(o): |
| return all(isinstance(o.get(k), (list, tuple)) and len(o[k]) == 3 for k in ("size", "pos")) and \ |
| all(_finite_number(v) for v in list(o["size"]) + list(o["pos"]) + [o.get("yaw")]) |
|
|
|
|
| def _finite_number(value): |
| if not isinstance(value, (int, float)) or isinstance(value, bool): |
| return False |
| try: |
| return math.isfinite(value) |
| except OverflowError: |
| return False |
|
|
|
|
| def _room_geometry_error(room): |
| boundary = room.get("boundary") |
| if not isinstance(boundary, (list, tuple)) or len(boundary) < 3 or \ |
| any(not isinstance(p, (list, tuple)) or len(p) != 2 for p in boundary): |
| return "boundary_shape" |
| if not all(_finite_number(v) for p in boundary for v in p): |
| return "bad_numbers" |
| if len({tuple(p) for p in boundary}) < 3: |
| return "boundary_shape" |
| height = room.get("height") |
| if height is not None and (not _finite_number(height) or height <= 0): |
| return "bad_numbers" |
| return None |
|
|
|
|
| def _keep_fixed(fixed, objs, shell, cover=0.5, slab=0.5): |
| """Fixed boxes worth showing next to the reference layout: touching the room, not a slab; a box whose top is |
| below FLAT_Z is a covering or trim, not an obstacle; a box sunk below the floor is cut at the floor. A NON-structure |
| box (unlabelled, wall-anchored, wall-hung, no-front) mostly covered by a placed object is that object's second |
| annotation and is dropped; a structure box (column, stair, door, window...) is real whatever stands on it, so it |
| is always kept (a reference layout through it shows up as a GT fixed collision).""" |
| kept = [(Polygon(footprint(o)), o["pos"][2], o["pos"][2] + o["size"][2]) for o in objs] |
| near = shell.buffer(WALL_BAND_M) |
| wall = shell.buffer(THICK_WALL_M) |
| out = [] |
| for f in fixed: |
| z0, z1 = f["pos"][2], f["pos"][2] + f["size"][2] |
| if z1 < FLAT_Z: |
| continue |
| if z0 < -FLOOR_Z: |
| f, z0 = {**f, "z_src": z0, "pos": [f["pos"][0], f["pos"][1], 0.0], "size": [f["size"][0], f["size"][1], z1]}, 0.0 |
| fp = Polygon(footprint(f)) |
| if fp.area <= 0 or fp.area >= slab * shell.area or \ |
| not fp.intersects(wall if f.get("fixed_kind") == "structure" else near): |
| continue |
| over = [p for p, lo, hi in kept if min(z1, hi) - max(z0, lo) > 0.05] |
| if f.get("fixed_kind") != "structure" and over and unary_union(over).intersection(fp).area >= cover * fp.area: |
| continue |
| out.append(f) |
| return out |
|
|
|
|
| def prep(room, a): |
| """IR room -> (clean room, None) or (None, reject reason). The clean room carries `fixed`: boxes the model is |
| told about with their pose and never places (each tagged fixed_kind for stats): |
| structure doors, windows, columns, stairs... (anchors.fixed_geometry) |
| generic unlabelled boxes ('unknown', 'otherprop', 'object') standing on the floor |
| wall wall-anchored items whose box stands on the floor (a wall cabinet run from z=0) |
| floating wall-hung floor furniture in designer rooms (anchor "fixed" in anchors.infer_anchors) |
| no_front placeable boxes whose front the source does not give (InteriorGS custom cabinetry): the answer needs |
| every object's yaw, so they are shown as built-in instead of rejecting the room (flag no_front_fixed) |
| all of them filtered by _keep_fixed. Tilted boxes are kept as upright boxes with their yaw; floor boxes sunk |
| up to 30 cm and floor furniture floating up to 50 cm with nothing below (scans, or anything that cannot hang) |
| are snapped to the floor; wall / ceiling / inside-cabinet items are dropped without rejecting the room.""" |
| meta = room.get("meta", {}) |
| if meta.get("front_known") is not True: |
| return None, "front_unknown" |
| if meta.get("n_incomplete"): |
| return None, "incomplete_objects" |
| if len({o["id"] for o in room["objects"]}) != len(room["objects"]): |
| return None, "duplicate_source_ids" |
| if room.get("boundary_type") not in a.boundary_types or not room.get("boundary"): |
| return None, "boundary" |
| if reason := _room_geometry_error(room): |
| return None, reason |
| if not all(_finite(o) for o in room["objects"]): |
| return None, "bad_numbers" |
| fixed = [{**o, "fixed_kind": "structure"} for o in fixed_geometry(room)] |
| allobj = annotate(room) |
| objs = closed(allobj, a.source_anchors.get(room["source"], a.anchors)) |
| generic = {o["id"] for o in objs if GENERIC.match(norm_cat(o["category"])) or not norm_cat(o["category"])} |
| sunk = {o["id"] for o in objs if not o.get("parent") and o["pos"][2] < -FURNITURE_SUNK_Z} |
| objs = _drop_with_dependants(objs, generic | sunk) |
| byid = {o["id"]: o for o in objs} |
| objs = _drop_with_dependants(objs, [ |
| o["id"] for o in objs if o.get("parent") and ( |
| (p := byid.get(o["parent"])) is None or not support_ok(o, p))]) |
| nofront = {o["id"] for o in objs if o.get("front_known") is False} |
| objs = _drop_with_dependants(objs, nofront) |
| if any(min(o["size"]) < 0 for o in objs): |
| return None, "bad_numbers" |
| kept = {o["id"] for o in objs} |
| children = {o["id"] for o in allobj if o.get("parent")} |
| |
| for o in allobj: |
| nc = norm_cat(o["category"]) |
| if o["id"] in kept or o["id"] in children or o.get("parent") or o["size"][2] < FLAT_Z: |
| continue |
| if not (-FURNITURE_SUNK_Z <= o["pos"][2] <= FLOOR_Z or o["anchor"] == "fixed" or o["id"] in nofront): |
| continue |
| if hanging(o["category"]) or WALL_ART.search(nc) or head(o) in SOFT_COVERS: |
| continue |
| kind = "generic" if o["id"] in generic else "no_front" if o["id"] in nofront else \ |
| "floating" if o["anchor"] == "fixed" else "wall" |
| fixed.append({**{k: v for k, v in o.items() if k not in ("anchor", "parent", "tilted", "anchor_inferred")}, |
| "fixed_kind": kind}) |
| shell = Polygon(room["boundary"]).buffer(0) |
| fixed = _keep_fixed(fixed, objs, shell) |
| shown = kept | {f["id"] for f in fixed} |
| |
| |
| inner = shell.buffer(-WALL_BAND_M) |
| |
| covered = unary_union([Polygon(footprint(o)) for o in objs if not o.get("parent") and o["size"][2] >= FLAT_Z] |
| + [Polygon(footprint(f)) for f in fixed if f["size"][2] >= FLAT_Z]) |
| polygon = room.get("boundary_type") == "polygon" |
| low = lambda o: o["pos"][2] <= FLOOR_Z and o["pos"][2] + o["size"][2] >= FLAT_Z |
| hidden = sum(Polygon(footprint(o)).intersection(shell).difference(covered).area for o in room["objects"] |
| if o["id"] not in shown and o["id"] not in children and not o.get("parent") and low(o) |
| and not hanging(o["category"]) |
| and not WALL_ART.search(norm_cat(o["category"])) and not _flush(o, inner, polygon)) |
| if len(objs) < a.min_objects: |
| return None, "object_count" |
| |
| |
| rounded = lambda o: {**o, "size": shown_size(o), "yaw": math.radians(round(math.degrees(o["yaw"])) % 360)} |
| objs = [{**rounded(o), **({"z_src": o["pos"][2]} if not o.get("parent") and abs(o["pos"][2]) > 1e-9 else {})} |
| for o in objs] |
| fixed = [{**rounded(o), "pos": [num(o["pos"][0]), num(o["pos"][1]), |
| 0 if -FLOOR_Z <= o["pos"][2] <= FLOOR_Z else num(o["pos"][2])]} for o in fixed] |
| byid = {o["id"]: o for o in objs} |
| done = set() |
|
|
| def z_of(o): |
| if o["id"] not in done: |
| p = byid.get(o.get("parent")) |
| o["pos"] = [num(o["pos"][0]), num(o["pos"][1]), num(z_of(p) + p["size"][2]) if p else 0] |
| done.add(o["id"]) |
| return o["pos"][2] |
| for o in objs: |
| z_of(o) |
| b = clean_boundary(room["boundary"]) |
| if b is None or (a.max_vertices and len(b) > a.max_vertices) or Polygon(b).area < MIN_AREA_M2: |
| return None, "boundary_shape" |
| out = {**room, "boundary": b, "objects": objs, "fixed": fixed} |
| |
| |
| flags = {} |
| area = Polygon(b).area |
| geometry = check(out, oob_tol=a.oob_tol) |
| oob = geometry["oob"] |
| if oob: |
| flags["oob_objects"] = len(oob) |
| if hidden > 0: |
| flags["hidden_m2"] = round(hidden, 2) |
| if hidden >= a.hidden_max: |
| flags["hidden_obstacle"] = True |
| if area < 1.0: |
| flags["small_area_m2"] = round(area, 2) |
| if any(not o.get("parent") and head(o) not in FLAT and o["size"][2] >= 0.5 |
| and o["size"][0] * o["size"][1] >= 0.8 * area for o in objs): |
| flags["room_filling_object"] = True |
| if _bad_overlap(objs): |
| flags["overlapping_furniture"] = True |
| through = geometry["fixed_blocking"] |
| if through: |
| flags["fixed_collision"] = len(through) |
| if len(objs) == 1: |
| flags["single_object"] = True |
| snapped = sum(abs(o.get("z_src") or 0) > FLOOR_Z for o in objs) |
| if snapped: |
| flags["z_snapped"] = snapped |
| if any(o.get("tilted") for o in objs): |
| flags["tilted"] = sum(bool(o.get("tilted")) for o in objs) |
| if nofront: |
| flags["no_front_fixed"] = len(nofront) |
| hit = [k for k in a.reject_flagged if k in flags] |
| if hit: |
| return None, "flagged:" + hit[0] |
| out["meta"] = {**out.get("meta", {}), "flags": flags} |
| h = out.get("height") |
| if h and any(o["pos"][2] + o["size"][2] > h + 0.05 for o in objs): |
| out = {**out, "height": None, "meta": {**out["meta"], "height_dropped": True}} |
| elif h: |
| out = {**out, "height": num(h)} |
| return out, None |
|
|
|
|
| def _bad_overlap(objs): |
| """Copies of the same furniture that pass through each other (3D box overlap over 30% of the smaller footprint): |
| a wardrobe or bookcase entered twice in the source. Different categories are not judged here: at box level an |
| appliance built into a cabinet run, a loft bed over its desk or an L-sofa's box over its table look the same as a |
| real interpenetration, so those stay collisions reported by evaluate against the GT rate.""" |
| fl = [o for o in objs if not o.get("parent") and head(o) not in FLAT | SOFT_HEADS |
| and not COMPOSITE.search(norm_cat(o["category"]))] |
| fps = [Polygon(footprint(o)) for o in fl] |
| for i, a in enumerate(fl): |
| for j in range(i + 1, len(fl)): |
| b = fl[j] |
| if norm_cat(a["category"]) != norm_cat(b["category"]): |
| continue |
| if min(a["pos"][2] + a["size"][2], b["pos"][2] + b["size"][2]) - max(a["pos"][2], b["pos"][2]) <= 0.05: |
| continue |
| small, big = sorted((fps[i].area, fps[j].area)) |
| if fps[i].intersection(fps[j]).area > 0.3 * small and small >= big / 3: |
| return True |
| return False |
|
|
|
|
| def exact_key(room): |
| """Same boundary and the same objects and fixed boxes (category, size, pose) up to a 90-degree turn.""" |
| keys = [] |
| for k in range(4): |
| r = rot90(room, k) |
| pose = lambda o, c: (c, tuple(o["size"]), num(o["pos"][0]), num(o["pos"][1]), num(o["pos"][2]), |
| round(math.degrees(o["yaw"])) % 360) |
| keys.append(repr((r["boundary"], sorted(pose(o, norm_cat(o["category"])) for o in r["objects"]), |
| sorted(pose(o, slug(o["category"])) for o in r.get("fixed") or [])))) |
| return hashlib.sha1(min(keys).encode()).hexdigest() |
|
|
|
|
| def dedup(rooms): |
| """Same group in several sources (e.g. a ScanNet scan in Scan2CAD and InternScenes): keep the source whose |
| rooms of that group keep more objects; ties go to a polygon boundary, then source name.""" |
| by = collections.defaultdict(lambda: collections.defaultdict(list)) |
| for r in rooms: |
| by[r["group"]][r["source"]].append(r) |
| keep, dropped = [], collections.Counter() |
| for g, srcs in by.items(): |
| best = max(srcs, key=lambda s: (sum(len(r["objects"]) for r in srcs[s]), |
| any(r["boundary_type"] == "polygon" for r in srcs[s]), s)) |
| keep += srcs[best] |
| for s, rs in srcs.items(): |
| if s != best: |
| dropped[s] += len(rs) |
| return keep, dropped |
|
|
|
|
| def rng_for(uid, salt): |
| return random.Random(int(hashlib.sha1(f"{salt}:{uid}".encode()).hexdigest(), 16)) |
|
|
|
|
| def extract_constraints(room, rng, max_k=4): |
| """1..max_k relations that hold on the reference layout, one per type first (round robin, rarest type first). |
| Only a few relations are given, never the full relation graph: the model must still lay out everything else.""" |
| b = room["boundary"] |
| objs = room["objects"] |
| fps = {o["id"]: Polygon(footprint(o)) for o in objs} |
| cand = collections.defaultdict(list) |
| objs = [o for o in objs if head(o) not in FLAT and o["size"][2] >= 0.05] |
| idx = {o["id"]: o for o in objs} |
| for o in objs: |
| if o.get("parent"): |
| cand["on"].append(["on", o["id"], o["parent"]]) |
| if room["boundary_type"] == "polygon" and not o.get("parent") and holds(["against_wall", o["id"]], objs, b, idx): |
| cand["against_wall"].append(["against_wall", o["id"]]) |
| for x in objs: |
| for y in objs: |
| if x is y or x.get("parent") or y.get("parent"): |
| continue |
| if head(x) in FRONT and not fps[x["id"]].intersects(fps[y["id"]]) and \ |
| holds(["faces", x["id"], y["id"]], objs, b, idx): |
| cand["faces"].append(["faces", x["id"], y["id"]]) |
| if x["id"] < y["id"]: |
| gap = fps[x["id"]].distance(fps[y["id"]]) |
| d = next((d for d in (0.1, 0.2, 0.3, 0.5) if d >= gap), None) |
| small = min(fps[x["id"]].area, fps[y["id"]].area) |
| if small > 0 and fps[x["id"]].intersection(fps[y["id"]]).area > 0.3 * small: |
| continue |
| if d is not None and holds(["near", x["id"], y["id"], d], objs, b, idx): |
| cand["near"].append(["near", x["id"], y["id"], d]) |
| floor = [o for o in objs if not o.get("parent")] |
| cat = [norm_cat(o["category"]) for o in floor] |
| |
| |
| |
| |
| geo = np.array([fps[m["id"]] for m in floor], dtype=object) |
| tree = STRtree(list(geo)) |
| for i, p in enumerate(floor): |
| for j in range(i + 1, len(floor)): |
| if cat[i] != cat[j]: |
| continue |
| q = floor[j] |
| seg = LineString([p["pos"][:2], q["pos"][:2]]) |
| ks = np.sort(tree.query(seg)) |
| if not len(ks): |
| continue |
| g = geo[ks] |
| ok = ~shapely.intersects(g, Point(p["pos"][:2])) & ~shapely.intersects(g, Point(q["pos"][:2])) \ |
| & shapely.intersects(g, seg) & ~shapely.touches(g, seg) & (ks != i) & (ks != j) |
| for k in ks[ok]: |
| cand["between"].append(["between", floor[k]["id"], p["id"], q["id"]]) |
| types = [t for t in PRIORITY if cand[t]] |
| for t in types: |
| rng.shuffle(cand[t]) |
| k, out = rng.randint(1, max_k), [] |
| while len(out) < k and any(cand[t] for t in types): |
| for t in types: |
| if cand[t] and len(out) < k: |
| out.append(cand[t].pop()) |
| return out |
|
|
|
|
| def count_constraints(st, split, room, cons): |
| """stats: constraints per type; 'on' split by evidence (support given by the source vs box-contact inference).""" |
| objs = {o["id"]: o for o in room["objects"]} |
| for c in cons or []: |
| t = c[0] |
| if t == "on": |
| t += "(inferred)" if objs[c[1]].get("anchor_inferred") else "(source)" |
| st[f"{split}_constraints:{t}"] += 1 |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--ir", required=True, help="dir of <source>.jsonl from fastfill.adapters") |
| ap.add_argument("--out", required=True, help="new dataset directory (must be absent or empty)") |
| ap.add_argument("--sources", nargs="*", default=None, help="default: every file in --ir") |
| ap.add_argument("--min_objects", type=int, default=1, help="rooms with fewer placeable objects have no target") |
| ap.add_argument("--max_vertices", type=int, default=0, help="reject rooms with more boundary vertices (0: no limit)") |
| ap.add_argument("--reject_flagged", nargs="*", default=[], |
| help="quality flags that reject a room instead of being recorded (oob_objects hidden_obstacle " |
| "small_area_m2 room_filling_object overlapping_furniture single_object)") |
| ap.add_argument("--keep_duplicates", action="store_true", help="keep exact duplicate layouts inside a split") |
| ap.add_argument("--boundary_types", nargs="*", default=["polygon", "hull"]) |
| ap.add_argument("--anchors", nargs="*", default=["floor", "object"], help="object anchors kept (see anchors.py)") |
| ap.add_argument("--oob_tol", type=float, default=0.10, help="flag oob_objects: reference objects outside boundary+tol") |
| ap.add_argument("--hidden_max", type=float, default=0.3, help="flag hidden_obstacle: m2 of unexplained dropped floor boxes") |
| ap.add_argument("--cap", nargs="*", default=[], help="SOURCE=N: keep N train rooms of SOURCE (hash-sampled); " |
| "default none: balance sources at sampling time, not here") |
| ap.add_argument("--source_anchors", nargs="*", default=[], help="SOURCE=a,b: anchors kept for that source") |
| ap.add_argument("--dev", type=float, default=0.05) |
| ap.add_argument("--test", type=float, default=0.05) |
| ap.add_argument("--constraint_frac", type=float, default=0.0, help="share of train rooms given as S2 (v1.1: 0.3)") |
| ap.add_argument("--no_rot90", action="store_true") |
| ap.add_argument("--field_dropout", type=float, default=0.2, help="train: drop room_type / height each with p") |
| ap.add_argument("--desc_dropout", type=float, default=0.5, help="train: drop all desc of a room with p") |
| a = ap.parse_args() |
| a.source_anchors = {k: v.split(",") for k, v in (x.split("=") for x in a.source_anchors)} |
| caps = {k: int(v) for k, v in (x.split("=") for x in a.cap)} |
|
|
| files = _input_files(a, ap) |
| _validate_destination(a.out) |
| |
| here = Path(__file__).resolve().parent |
| code_files = _code_files(here) |
| code_sha = {str(p.relative_to(here)): _sha256(p) for p in code_files} |
| ir_sha = {p.name: _sha256(p) for p in files} |
| rooms, stats = [], collections.defaultdict(collections.Counter) |
| for f in files: |
| for line in _json_lines(f): |
| r0 = line |
| st = stats[r0["source"]] |
| r, why = prep(r0, a) |
| st["scanned"] += 1 |
| if why: |
| st["rejected:" + why] += 1 |
| continue |
| rooms.append({**r, "meta": {**r["meta"], "n_src_objects": len(r0["objects"])}}) |
| print(f"prepared {f.name}: {len(rooms)} retained rooms so far", flush=True) |
| if not rooms: |
| ap.error("no usable rooms; output was not changed") |
| rooms, dropped = dedup(rooms) |
| for s, n in dropped.items(): |
| stats[s]["dedup_dropped"] += n |
|
|
| splits = assign_splits(rooms, a.dev, a.test) |
| |
| |
| first = {} |
| for i in sorted(range(len(rooms)), key=lambda i: hashlib.sha1(rooms[i]["uid"].encode()).hexdigest()): |
| key = (splits[i], exact_key(rooms[i])) |
| if key in first and not a.keep_duplicates: |
| stats[rooms[i]["source"]]["duplicate_layout"] += 1 |
| splits[i] = "duplicate" |
| else: |
| first[key] = i |
| for src, n in caps.items(): |
| idx = sorted((i for i, r in enumerate(rooms) if r["source"] == src and splits[i] == "train"), |
| key=lambda i: hashlib.sha1(rooms[i]["uid"].encode()).hexdigest()) |
| for i in idx[n:]: |
| splits[i] = "capped" |
| if not any(s in ("train", "dev", "test") for s in splits): |
| ap.error("no rooms remain after filtering; output was not changed") |
| destination = Path(a.out).absolute() |
| destination.parent.mkdir(parents=True, exist_ok=True) |
| with tempfile.TemporaryDirectory(prefix=f".{destination.name}-", dir=destination.parent) as staging: |
| _write_dataset(rooms, splits, stats, a, Path(staging), code_sha, ir_sha) |
| if any(_sha256(p) != ir_sha[p.name] for p in files): |
| raise RuntimeError("input IR changed during build; output was not published") |
| dependencies = {"__init__.py", "build.py", "anchors.py", "scene.py", "split.py", "validate.py"} |
| current_code = {str(p.relative_to(here)): _sha256(p) for p in _code_files(here)} |
| if any(current_code.get(name) != code_sha.get(name) for name in dependencies): |
| raise RuntimeError("implementation changed during build; output was not published") |
| |
| |
| manifest_path = Path(staging) / "MANIFEST.json" |
| with manifest_path.open(encoding="utf-8") as f: |
| manifest = json.load(f) |
| with manifest_path.open("w", encoding="utf-8") as f: |
| json.dump({**manifest, "code_sha256": current_code, "build_start_code_sha256": code_sha}, |
| f, indent=1, allow_nan=False) |
| _validate_destination(destination) |
| os.rename(staging, destination) |
| for k, v in sorted(stats.items()): |
| print(f"{k:24s} {dict(v)}") |
|
|
|
|
| def _write_dataset(rooms, splits, stats, a, directory, code_sha, ir_sha): |
| with ExitStack() as stack: |
| open_output = lambda name: stack.enter_context(open(directory / f"{name}.jsonl", "w", encoding="utf-8")) |
| outs = {s: open_output(s) for s in ("train", "dev", "test")} |
| ir_outs = {s: open_output(s) for s in |
| ("dev_rooms", "test_rooms", "dev_constrained_rooms", "test_constrained_rooms")} |
| for r, s in zip(rooms, splits): |
| st = stats[r["source"]] |
| if s in ("capped", "duplicate"): |
| st["capped"] += s == "capped" |
| continue |
| |
| st["height_dropped"] += bool(r.get("meta", {}).get("height_dropped")) |
| st["objects_dropped"] += r["meta"]["n_src_objects"] - len(r["objects"]) - len(r["fixed"]) |
| st["fixed_items"] += len(r["fixed"]) |
| for f in r["fixed"]: |
| st["fixed:" + f["fixed_kind"]] += 1 |
| st["rooms_with_fixed"] += bool(r["fixed"]) |
| st["objects_max"] = max(st["objects_max"], len(r["objects"])) |
| for k in r["meta"]["flags"]: |
| st["flag:" + k] += 1 |
| st["supports_inferred"] += sum(bool(o.get("anchor_inferred")) for o in r["objects"] if o.get("parent")) |
| st["supports_evidence"] += sum(not o.get("anchor_inferred") for o in r["objects"] if o.get("parent")) |
| if s == "train": |
| |
| |
| aug, s2, drop = (rng_for(r["uid"], k) for k in ("aug", "s2", "dropout")) |
| c = canonical(r if a.no_rot90 else rot90(r, aug.randrange(4))) |
| cons = extract_constraints(c, s2) if s2.random() < a.constraint_frac else None |
| st["train_S2" if cons else "train_S1"] += 1 |
| count_constraints(st, "train", c, cons) |
| if drop.random() < a.field_dropout: |
| c = {**c, "room_type": None} |
| if drop.random() < a.field_dropout: |
| c = {**c, "height": None} |
| m = messages(c, cons, with_desc=drop.random() >= a.desc_dropout) |
| else: |
| c = canonical(r) |
| st[s] += 1 |
| m = messages(c) |
| ir_outs[f"{s}_rooms"].write(json.dumps(c, ensure_ascii=False, allow_nan=False) + "\n") |
| cons = extract_constraints(c, rng_for(r["uid"], "eval")) |
| if cons: |
| count_constraints(st, s, c, cons) |
| ir_outs[f"{s}_constrained_rooms"].write(json.dumps({**c, "constraints": cons}, ensure_ascii=False, allow_nan=False) + "\n") |
| outs[s].write(json.dumps({"uid": r["uid"], "source": r["source"], "flags": r["meta"]["flags"], "messages": m}, |
| ensure_ascii=False, allow_nan=False) + "\n") |
| with open(directory / "stats.json", "w", encoding="utf-8") as f: |
| json.dump({k: dict(v) for k, v in sorted(stats.items())}, f, indent=1, allow_nan=False) |
| manifest = {"args": vars(a), |
| "files": {p.name: {"sha256": _sha256(p), "lines": _line_count(p)} |
| for p in sorted(directory.iterdir()) if p.suffix == ".jsonl" or p.name == "stats.json"}, |
| "code_sha256": code_sha, "ir_sha256": ir_sha} |
| with open(directory / "MANIFEST.json", "w", encoding="utf-8") as f: |
| json.dump(manifest, f, indent=1, allow_nan=False) |
|
|
|
|
| def _sha256(path): |
| digest = hashlib.sha256() |
| with open(path, "rb") as f: |
| for chunk in iter(lambda: f.read(1024 * 1024), b""): |
| digest.update(chunk) |
| return digest.hexdigest() |
|
|
|
|
| def _line_count(path): |
| with open(path, "rb") as f: |
| return sum(1 for _ in f) |
|
|
|
|
| def _code_files(here): |
| return sorted(p for p in here.rglob("*.py") |
| if not any(part.startswith(("review", ".", "__pycache__")) or part == "tests" |
| for part in p.relative_to(here).parts[:-1])) |
|
|
|
|
| def _json_lines(path): |
| with open(path, encoding="utf-8") as f: |
| for line_number, line in enumerate(f, 1): |
| try: |
| yield json.loads(line) |
| except json.JSONDecodeError as exc: |
| raise ValueError(f"{path}:{line_number}: invalid JSON: {exc.msg}") from exc |
|
|
|
|
| def _input_files(a, ap): |
| if not Path(a.ir).is_dir(): |
| ap.error(f"IR directory does not exist: {a.ir}") |
| files = sorted(Path(a.ir).glob("*.jsonl")) |
| if a.sources is not None: |
| missing = set(a.sources) - {p.stem for p in files} |
| if missing: |
| ap.error("missing IR sources: " + ", ".join(sorted(missing))) |
| files = [p for p in files if p.stem in a.sources] |
| if not files or any(p.stat().st_size == 0 for p in files): |
| ap.error("selected IR files must exist and be nonempty") |
| if any(not 0 <= x <= 1 for x in (a.dev, a.test, a.constraint_frac, a.field_dropout, a.desc_dropout)) or \ |
| a.dev + a.test > 1: |
| ap.error("split/dropout/constraint fractions must be in [0, 1], with dev + test <= 1") |
| return files |
|
|
|
|
| def _validate_destination(path): |
| path = Path(path) |
| if path.is_symlink() or (path.exists() and (not path.is_dir() or any(path.iterdir()))): |
| raise FileExistsError(f"refusing to replace existing dataset: {path}; choose a new version directory") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
