FastFill v4 base (step 6552): LoRA adapter, geometry heads, tokenizer, training code, predict.py

e6f80f3 verified 6 days ago

5.29 kB

"""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)

	"""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)