| """Detached assignment only within explicitly certified exchangeable groups.""" |
| from __future__ import annotations |
|
|
| import json |
| import math |
|
|
| import torch |
| from scipy.optimize import linear_sum_assignment |
|
|
|
|
| def _rename_references(row, mapping, scalar_fields, list_fields=()): |
| """Copy a schema row, renaming only its explicit object-reference fields. |
| |
| IDs share string representations with ordinary descriptions and attributes; |
| those values, dictionary keys, group labels and asset-local surface IDs are |
| independent namespaces. Nested metadata must never be interpreted as refs. |
| ``object_ids`` is the existing top-level constraint reference-list extension. |
| """ |
| renamed = {} |
| for key, value in row.items(): |
| if key in scalar_fields and isinstance(value, str): |
| renamed[key] = mapping.get(value, value) |
| elif key in list_fields and isinstance(value, list): |
| renamed[key] = [mapping.get(ref, ref) if isinstance(ref, str) else ref for ref in value] |
| else: |
| renamed[key] = value |
| return renamed |
|
|
|
|
| def certify_group(objects, constraints, indices): |
| """Check conditioning invariance under generators of the group's permutations.""" |
| fingerprints = [json.dumps({k: v for k, v in objects[i].items() if k not in {"id", "exchangeable_group"}}, |
| sort_keys=True) for i in indices] |
| if len(set(fingerprints)) != 1: |
| raise ValueError("exchangeable group has distinct descriptions/requirements/roles") |
| canonical = lambda rows: sorted(json.dumps(r, sort_keys=True) for r in rows) |
| original = canonical(constraints) |
| object_graph = canonical(objects) |
| for i in indices[1:]: |
| mapping = {objects[indices[0]]["id"]: objects[i]["id"], objects[i]["id"]: objects[indices[0]]["id"]} |
| renamed_constraints = [_rename_references(c, mapping, ("object_id", "target_id", "parent_id"), |
| ("target_ids", "object_ids")) for c in constraints] |
| renamed_objects = [_rename_references(o, mapping, ("id", "support_parent")) for o in objects] |
| if canonical(renamed_constraints) != original or canonical(renamed_objects) != object_graph: |
| raise ValueError("exchange would change a constraint or support role; use fixed identities") |
|
|
|
|
| def group_labels(objects, constraints, position_valid): |
| """C5 exchangeable labels per request: identical non-id fields, complete position labels, certified swap. |
| |
| A signature set (identical request fields except ``id``) whose whole swap fails certification |
| falls back to its members that no constraint reference field (the ones ``certify_group`` renames) |
| and no ``support_parent`` names; that remainder always certifies and is grouped under the set's |
| label when it has at least two members. |
| """ |
| candidates = {} |
| for i, obj in enumerate(objects): |
| if all(position_valid[i]): |
| signature = json.dumps({k: v for k, v in obj.items() if k != "id"}, sort_keys=True) |
| candidates.setdefault(signature, []).append(i) |
| referenced = {o.get("support_parent") for o in objects} |
| for c in constraints: |
| referenced.update(c.get(k) for k in ("object_id", "target_id", "parent_id")) |
| referenced.update(ref for k in ("target_ids", "object_ids") for ref in c.get(k) or ()) |
| labels = [None] * len(objects) |
| for number, indices in enumerate(candidates.values()): |
| for members in (indices, [i for i in indices if objects[i]["id"] not in referenced]): |
| if len(members) < 2: |
| break |
| try: |
| certify_group(objects, constraints, members) |
| except ValueError: |
| continue |
| for i in members: |
| labels[i] = f"anonymous_{number}" |
| break |
| return labels |
|
|
|
|
| def validate_matching_options(enabled, alpha_position, alpha_size): |
| if not isinstance(enabled, bool): |
| raise ValueError("matching enabled/hungarian flag must be boolean") |
| weights = (alpha_position, alpha_size) |
| if (any(type(value) not in {int, float} or not math.isfinite(value) or value < 0 for value in weights) |
| or alpha_position + alpha_size <= 0): |
| raise ValueError("matching weights must be native finite nonnegative numbers and not both zero") |
|
|
|
|
| @torch.no_grad() |
| def match_batch(predictions, batch, enabled=True, alpha_position=1., alpha_size=1., *, loss_config=None): |
| """Groups come from collate's ``batch["exchangeable_group"]`` (per sample, per |
| request slot), never from the rendered condition objects. A group needs |
| complete position labels on every member; otherwise it keeps fixed identity. |
| Position cost (times ``alpha_position``): the decoded L1 for regression |
| predictions; for grid_residual predictions (``position_cell_logits``) each |
| (prediction slot, target) pair costs exactly the loss's ``position`` terms, |
| ``losses._grid_position`` under ``loss_config`` (default ``LossConfig()``): |
| ``position_cell * CE + position_residual * GT-cell residual + z / 3`` on the |
| prediction slot's learned (not ``fixed_position_mask``) coordinates, summed in |
| float64 (on the CPU for Metal, which has none). |
| Size enters the cost only when every member also has a complete size label; |
| a pair whose target object is box-symmetric (its own |
| ``batch["size_axis_swap_allowed"]``) costs the minimum log-size error over the |
| target's (sx, sy) and (sy, sx) orders, every other pair the written order. |
| Members marked in ``batch["given_pose_mask"]`` (masked layout training) keep |
| fixed identity: the rest of their group matches among itself. |
| """ |
| validate_matching_options(enabled, alpha_position, alpha_size) |
| if enabled and "exchangeable_group" not in batch: |
| raise ValueError("exchangeable matching requires batch exchangeable_group from collate") |
| grid = "position_cell_logits" in predictions |
| if grid: |
| from fastfill.v2.losses import GRID_KEYS, LossConfig, _grid_position |
| cfg = LossConfig() if loss_config is None else loss_config |
| valid = batch["slot_mask"] |
| bsz, slots = valid.shape |
| assignments = [] |
| for b in range(bsz): |
| indices = list(range(slots)) |
| objects = batch["objects"][b] |
| if len(objects) != int(valid[b].sum()) or not valid[b, :len(objects)].all(): |
| raise ValueError("slots must be contiguous and bound to every request") |
| groups = {} |
| if enabled: |
| labels = batch["exchangeable_group"][b] |
| if len(labels) != len(objects) or any(label is not None and not isinstance(label, str) for label in labels): |
| raise ValueError("exchangeable_group must hold one string-or-null label per request slot") |
| for i, label in enumerate(labels): |
| if label: |
| groups.setdefault(label, []).append(i) |
| if any(len(group) > 1 for group in groups.values()) and "conditions" not in batch: |
| raise ValueError("exchangeable matching requires original conditions and constraint graph") |
| conditions = batch.get("conditions", [{} for _ in range(bsz)]) |
| for group in groups.values(): |
| certify_group(objects, conditions[b].get("constraints", []), group) |
| if "given_pose_mask" in batch: |
| group = [i for i in group if not bool(batch["given_pose_mask"][b, i])] |
| if len(group) < 2 or not batch["validity"]["position"][b, group].all(): |
| continue |
| pos, gt_pos = predictions["position_normalized"][b, group], batch["targets"]["position_normalized"][b, group] |
| if not torch.isfinite(pos).all() or not torch.isfinite(gt_pos).all(): |
| raise ValueError("nonfinite position geometry in matching group") |
| if grid: |
| n = len(group) |
| rows = torch.tensor(group, device=valid.device).repeat_interleave(n) |
| cols = torch.tensor(group, device=valid.device).repeat(n) |
| learn = ~batch.get("fixed_position_mask", torch.zeros_like(batch["validity"]["position"]))[b, rows] |
| parts = _grid_position({k: predictions[k][b, rows] for k in ("position_normalized", *GRID_KEYS)}, |
| batch["targets"]["position_normalized"][b, cols], torch.ones_like(rows, dtype=torch.bool), learn, cfg) |
| |
| host = torch.device("cpu") if valid.device.type == "mps" else valid.device |
| part, learn = (lambda name: parts[name].to(host).double()), learn.to(host) |
| pair = torch.zeros(n * n, dtype=torch.float64, device=host) |
| pair[learn[:, 0] & learn[:, 1]] = (cfg.position_cell * part("position_cell") + |
| cfg.position_residual * part("position_residual")) |
| pair[learn[:, 2]] += part("position_z") |
| cost = alpha_position * pair.view(n, n) |
| else: |
| cost = alpha_position * (pos[:, None] - gt_pos[None]).abs().sum(-1) |
| if batch["validity"]["size"][b, group].all(): |
| size, gt_size = predictions["size"][b, group], batch["targets"]["size"][b, group] |
| if not torch.isfinite(size).all() or not torch.isfinite(gt_size).all() or (size <= 0).any() or (gt_size <= 0).any(): |
| raise ValueError("nonfinite or nonpositive size geometry in matching group") |
| size_cost = (size.log()[:, None] - gt_size.log()[None]).abs().sum(-1) |
| if "size_axis_swap_allowed" in batch: |
| swapped = (size.log()[:, None] - gt_size.log()[None, :, [1, 0, 2]]).abs().sum(-1) |
| size_cost = torch.where(batch["size_axis_swap_allowed"][b, group][None], torch.minimum(size_cost, swapped), size_cost) |
| cost = cost + (alpha_size * size_cost).to(cost.device) |
| if not torch.isfinite(cost).all(): |
| raise ValueError("matching cost is not finite") |
| |
| rows, cols = linear_sum_assignment(cost.cpu().double().numpy()) |
| chosen = {group[int(row)]: group[int(col)] for row, col in zip(rows, cols)} |
| indices = [chosen.get(k, old) for k, old in enumerate(indices)] |
| assignments.append(indices) |
| return torch.tensor(assignments, dtype=torch.long, device=valid.device) |
|
|
|
|
| def permute_relations(assignment, parents=None, relation_matrix=None): |
| """Negative parent indices represent external supports and remain unchanged.""" |
| if (assignment.ndim != 1 or assignment.dtype not in {torch.int32, torch.int64} or |
| not torch.equal(assignment.sort().values, torch.arange(len(assignment), device=assignment.device, dtype=assignment.dtype))): |
| raise ValueError("relation assignment must be a one-to-one permutation") |
| if parents is not None and parents.shape != assignment.shape: |
| raise ValueError("parent targets must have one value per assigned object") |
| if relation_matrix is not None and relation_matrix.shape != (len(assignment), len(assignment)): |
| raise ValueError("relation matrix shape must match assigned objects") |
| inverse = assignment.argsort() |
| new_parents = None |
| if parents is not None: |
| selected = parents[assignment] |
| external = selected < 0 |
| if ((selected >= len(assignment)) & ~external).any(): |
| raise ValueError("parent index out of range") |
| new_parents = torch.where(external, selected, inverse[selected.clamp_min(0)]) |
| matrix = None if relation_matrix is None else relation_matrix[assignment][:, assignment] |
| return new_parents, matrix |
|
|
