File size: 5,289 Bytes
e6f80f3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | """Leakage-safe train/dev/test assignment.
Rooms are linked when they share
- a `group` (same house / scan across exports),
- a meta.group_aliases entry (e.g. several scans of one ARKit visit),
- a furniture content key (3D-FRONT reuses whole furniture sets across houses; exports repackage rooms),
- a layout key (same boundary and object poses up to a 90-degree rotation, whatever the assets).
Each connected component goes to one split, chosen by a stable hash; a component holding an eval-only room
(meta.eval_only) goes to test as a whole.
"""
import hashlib
import math
from fastfill.scene import num, rot90
# Procedurally generated sources share one asset library across all plans: equal furniture sets there are not
# repackaged rooms, and linking them fuses thousands of rooms into one component.
GROUP_ONLY = {"MansionWorld"}
def content_key(room, step=0.05):
"""Rotation/translation-invariant fingerprint: multiset of sorted object dims (5 cm grid)."""
if len(room["objects"]) < 3 or room.get("source") in GROUP_ONLY:
return None
dims = sorted(tuple(sorted(round(v / step) for v in o["size"])) for o in room["objects"])
return hashlib.sha1(repr(dims).encode()).hexdigest()
def layout_key(room):
"""Asset-independent design fingerprint: boundary + poses (cm, integer degrees), min over 4 rotations."""
if not room.get("boundary"):
return None
keys = []
for k in range(4):
r = rot90(room, k)
poses = sorted((num(o["pos"][0]), num(o["pos"][1]), round(math.degrees(o["yaw"])) % 360) for o in r["objects"])
keys.append(hashlib.sha1(repr((r["boundary"], poses)).encode()).hexdigest())
return min(keys)
def assign_splits(rooms, dev=0.05, test=0.05):
"""rooms: list of IR dicts -> list of 'train' | 'dev' | 'test' (same order)."""
parent = {}
def find(x):
while parent.setdefault(x, x) != x:
parent[x] = parent[parent[x]]
x = parent[x]
return x
def union(a, b):
ra, rb = find(a), find(b)
if ra != rb:
parent[max(ra, rb)] = min(ra, rb)
for r in rooms:
g = "g:" + r["group"]
find(g)
for key in ("c:%s" % content_key(r), "l:%s" % layout_key(r)):
if not key.endswith("None"):
union(g, key)
for al in r.get("meta", {}).get("group_aliases") or []:
union(g, "a:" + al)
held_out = {find("g:" + r["group"]) for r in rooms if r.get("meta", {}).get("eval_only")}
# a component is hashed by its smallest GROUP name, not by its union-find root (which may be a furniture or layout
# key): the split of a house then does not depend on object content, and a prep change does not move houses
name = {}
for r in rooms:
c = find("g:" + r["group"])
name[c] = min(name.get(c, "g:" + r["group"]), "g:" + r["group"])
out = []
for r in rooms:
c = find("g:" + r["group"])
h = int(hashlib.sha1(name[c].encode()).hexdigest(), 16) % 10000 / 10000
out.append("test" if c in held_out or h < test else "dev" if h < test + dev else "train")
return out
if __name__ == "__main__":
def room(group, sizes, poses=None, boundary=((0, 0), (5, 0), (5, 4), (0, 4)), source="S", meta=None):
poses = poses or [(1 + i, 1, 0) for i in range(len(sizes))]
return {"group": group, "source": source, "boundary": [list(p) for p in boundary], "meta": meta or {},
"objects": [{"size": s, "pos": [x, y, 0], "yaw": math.radians(d)} for s, (x, y, d) in zip(sizes, poses)]}
a = room("house1", [[2, 1.6, .5], [.5, .5, .6], [.5, .5, .6]])
b = room("house2", [[1.6, 2, .5], [.5, .5, .6], [.5, .6, .5]]) # same furniture, other house, axes permuted
c = room("house1", [[1, 1, 1]]) # same house, different room
v1 = room("vid1", [[1, 2, 3]], meta={"group_aliases": ["visit:9"]})
v2 = room("vid2", [[3, 2, 1.5]], meta={"group_aliases": ["visit:9"]}) # two videos of one visit
t1 = room("h_t1", [[1, 1, 1], [2, 2, 2]], [(1, 1, 0), (3, 2, 90)])
t2 = room("h_t2", [[.9, .9, 1.2], [2, 2, 1]], [(3, 1, 90), (2, 3, 180)], # same design stored rotated, new assets
boundary=((0, 0), (4, 0), (4, 5), (0, 5)))
e1 = room("eval1", [[1, 1, 1], [1, 2, 1], [2, 2, 2]], meta={"eval_only": True})
e2 = room("twin", [[1, 1, 1], [2, 1, 1], [2, 2, 2]]) # content twin of an eval-only room
m1 = room("mw1", [[1, 1, 1]] * 3, source="MansionWorld")
m2 = room("mw2", [[1, 1, 1]] * 3, [(2, 2, 0), (3, 2, 0), (4, 3, 0)], source="MansionWorld")
rooms = [a, b, c, v1, v2, t1, t2, e1, e2, m1, m2] + [
room(f"h{i}", [[i / 10 + .1, 1, 1]] * 3, [(1 + i / 1000, 1, 0), (2, 2, 0), (3, 3, 0)]) for i in range(2000)]
s = assign_splits(rooms)
assert s[0] == s[1] == s[2] and s[3] == s[4] and s[5] == s[6], s[:7]
assert s[7] == s[8] == "test", s[7:9]
assert content_key(m1) is None and layout_key(m1) != layout_key(m2) and layout_key(t1) == layout_key(t2)
frac = {k: s.count(k) / len(s) for k in ("train", "dev", "test")}
assert 0.85 < frac["train"] < 0.95 and 0.02 < frac["test"] < 0.08, frac
assert assign_splits(rooms) == s
print("split.py self-check ok", frac)
|