"""Stable group slots and the compatible 24-tensor label schema.""" from __future__ import annotations import numpy as np from .partition import cohesion, flexible_partition from .relevance import relevance_periods, threshold_graph from .budget import resolve from .state import eligible_players LATENT_STRIDE = 4 def latent_frame_count(num_frames: int, stride: int = LATENT_STRIDE) -> int: return (num_frames - 1) // stride + 1 def track_labels(seq: list[list[list[int]]], n_slots: int) -> list[list[int]]: """Match slots by intersection, Jaccard, then stable group/slot indices.""" out: list[list[int]] = [] prev: list[tuple[int, set[int]]] = [] # 上一周期的 (编号, 成员集合) for groups in seq: sets = [set(g) for g in groups] cand = [] for gi, s in enumerate(sets): for ci, ps in prev: inter = len(s & ps) if inter: # 交集大的先配；gi/ci 兜底定序，保证可复现 cand.append((-inter, -inter / len(s | ps), gi, ci)) cand.sort() lab = [-1] * len(groups) taken: set[int] = set() for _, _, gi, ci in cand: if lab[gi] < 0 and ci not in taken: lab[gi] = ci taken.add(ci) free = [c for c in range(n_slots) if c not in taken] for gi in range(len(groups)): if lab[gi] < 0: lab[gi] = free.pop(0) if free else gi % n_slots taken.add(lab[gi]) out.append(lab) prev = [(lab[gi], sets[gi]) for gi in range(len(groups))] return out def flex_label_arrays( groupings, periods, keep: np.ndarray, *, n_players: int, max_groups: int, start_frame: int, num_frames: int, threshold: float, budget, stride: int = LATENT_STRIDE, ) -> dict[str, np.ndarray]: """把逐周期的 `FlexGrouping` 摊成定长 tensor 字典。 列序 = CSV 的 `video_path1..P`（world 按 `row.slots` 加载时，world 下标 == 列下标）。 组号经 `track_labels` 跨周期跟踪，槽位落在 `0 .. G-1`（G = max_groups）。 """ P, G, K = n_players, max_groups, len(groupings) if K == 0: raise ValueError("没有更新周期，写不出标签") labs = track_labels([g.groups for g in groupings], n_slots=max(G, 1)) p_start = np.zeros(K, np.int32) p_frames = np.zeros(K, np.int32) p_mem = np.zeros((K, G, P), bool) p_go = np.zeros((K, P), bool) p_anchor = np.zeros((K, P), bool) p_score = np.full((K, G), -1.0, np.float32) p_size = np.zeros((K, G), np.int8) p_coh = np.full((K, G), -1.0, np.float32) p_pair = np.zeros(K, np.int32) p_cost = np.full(K, -1.0, np.float32) eligible = np.zeros(P, bool) eligible[np.asarray(keep, int)] = True for k, (gr, per, sid) in enumerate(zip(groupings, periods, labs)): p_start[k] = int(per.start_frame - start_frame) p_frames[k] = int(per.num_frames) W = threshold_graph(per.smooth, threshold) for gi, g in enumerate(gr.groups): s = int(sid[gi]) % G p_size[k, s] = len(g) p_score[k, s] = gr.scores[gi] if gi < len(gr.scores) else 0.0 p_coh[k, s] = cohesion(W, g) if len(g) > 1 else 0.0 for j in g: p_mem[k, s, int(keep[j])] = True for j in gr.global_only: p_go[k, int(keep[j])] = True for j in gr.shared: p_anchor[k, int(keep[j])] = True p_pair[k] = gr.pair_used # gr.groups 的下标是 keep 内部下标；未进 keep 的 player 在实际 # membership/global_only 里两者都是 False，不能误算成 global-only。 p_cost[k] = budget.topology_cost(gr.groups, n_players=len(keep)) has = p_mem.any(axis=1) p_lab = np.where(has, p_mem.argmax(axis=1).astype(np.int8), np.int8(-1)).astype( np.int8 ) raw = np.minimum( np.arange(latent_frame_count(num_frames, stride)) * stride, max(num_frames - 1, 0), ) idx = np.clip(np.searchsorted(p_start, raw, side="right") - 1, 0, K - 1) n_local = has.sum(axis=1).astype(np.float32) n_go = p_go.sum(axis=1).astype(np.float32) n_anchor = p_anchor.sum(axis=1).astype(np.float32) denom = np.float32(max(P, 1)) membership_latent = p_mem[idx] global_only_latent = p_go[idx] clip_cost = budget.topology_cost_timeline(membership_latent, global_only_latent) if clip_cost > budget.cost_limit: raise RuntimeError( f"整段 dynamic membership 通信超预算：{clip_cost:,.0f} > " f"{budget.cost_limit:,.0f}" ) return { "membership": membership_latent, "global_only": global_only_latent, "eligible": eligible, "labels": p_lab[idx], "latent_period": idx.astype(np.int16), "period_membership": p_mem, "period_global_only": p_go, "period_anchor": p_anchor, "period_labels": p_lab, "period_start": p_start, "period_frames": p_frames, "period_scores": p_score, "period_sizes": p_size, "period_cohesion": p_coh, "period_pair_used": p_pair, "period_total_cost": p_cost, "clip_total_cost": np.array(clip_cost, dtype=np.float64), "cost_limit": np.array(budget.cost_limit, dtype=np.float64), "coverage": n_local / denom, "orphan_ratio": n_go / denom, "anchor_ratio": n_anchor / denom, "pair_budget": np.array(budget.pair_budget, dtype=np.int32), "max_groups": np.array(G, dtype=np.int32), "max_size": np.array(budget.max_size, dtype=np.int32), } def generate_labels(state): """Generate pose-only labels without model state or random-number draws.""" state.validate() if state.clip_num_frames != 165: raise ValueError( "Automatic grouping currently requires 165 RGB frames per CSV row" ) budget = resolve(state.num_players) span = state.frame_window(state.clip_start_frame, state.clip_num_frames) keep = np.flatnonzero(eligible_players(state, span)) if len(keep) < 2: raise ValueError("Fewer than two eligible players in grouping clip") periods = relevance_periods(state, span, state.clip_start_frame, keep) groups = [flexible_partition(p.smooth, budget=budget) for p in periods] return flex_label_arrays( groups, periods, keep, n_players=state.num_players, max_groups=budget.groups, start_frame=state.clip_start_frame, num_frames=state.clip_num_frames, threshold=0.3, budget=budget, )