Spaces:
Running on Zero
Running on Zero
| """Prepare missing labels before data workers start; never overwrite input labels.""" | |
| from contextlib import contextmanager | |
| import fcntl | |
| import hashlib | |
| import io | |
| import json | |
| import os | |
| from pathlib import Path | |
| import tempfile | |
| import pandas as pd | |
| import torch | |
| from .labels import generate_labels | |
| from .state import GroupingState, row_identity | |
| from .validation import validate_generated_labels, validate_existing_labels | |
| # Included only in the cache filename when labels are generated because | |
| # group_path is missing. Bundled groups.pt files are loaded by path. | |
| # Change this string if the generated label numbers change. | |
| ALGORITHM_VERSION = "grouping" | |
| DEFAULT_CACHE = "outputs/grouping_cache" | |
| def _optional_path(base_path, value): | |
| if value is None or pd.isna(value) or not str(value).strip(): | |
| return None | |
| path = Path(str(value)) | |
| return path if path.is_absolute() else Path(base_path) / path | |
| def _cache_lock(path): | |
| with path.open("a") as handle: | |
| fcntl.flock(handle, fcntl.LOCK_EX) | |
| try: | |
| yield | |
| finally: | |
| fcntl.flock(handle, fcntl.LOCK_UN) | |
| def _atomic_save(payload, path): | |
| fd, temporary = tempfile.mkstemp( | |
| prefix=path.name + ".", suffix=".tmp", dir=path.parent | |
| ) | |
| try: | |
| with os.fdopen(fd, "wb") as handle: | |
| torch.save(payload, handle) | |
| handle.flush() | |
| os.fsync(handle.fileno()) | |
| os.replace(temporary, path) | |
| finally: | |
| if os.path.exists(temporary): | |
| os.unlink(temporary) | |
| def ensure_group_path( | |
| row, | |
| base_path, | |
| *, | |
| num_players, | |
| cache_dir=DEFAULT_CACHE, | |
| group_path_column="group_path", | |
| ): | |
| """Return an existing label path, or deterministically generate a missing one.""" | |
| existing = _optional_path(base_path, row.get(group_path_column)) | |
| if existing is not None and existing.exists(): | |
| if not existing.is_file(): | |
| raise ValueError(f"Grouping path is not a file: {existing}") | |
| validate_existing_labels( | |
| torch.load(existing, map_location="cpu", weights_only=True) | |
| ) | |
| return str(existing) | |
| state_path = _optional_path(base_path, row.get("grouping_state_path")) | |
| if state_path is None or not state_path.is_file(): | |
| raise FileNotFoundError( | |
| f"Missing grouping labels ({existing}) and grouping_state_path ({state_path}). " | |
| "Supply the lossless tick state including 48 frames of history; " | |
| "camera_pose_path files do not contain enough information." | |
| ) | |
| raw = state_path.read_bytes() | |
| state = GroupingState.load(io.BytesIO(raw)) | |
| identity = row_identity(row, num_players) | |
| actual = ( | |
| state.match_id, | |
| state.map_name, | |
| state.round_no, | |
| state.player_slots.tolist(), | |
| state.clip_start_frame, | |
| state.clip_num_frames, | |
| ) | |
| if actual != identity: | |
| raise ValueError( | |
| f"Grouping state does not match CSV round/player order/clip: {state_path}" | |
| ) | |
| key = hashlib.sha256( | |
| json.dumps( | |
| { | |
| "algorithm": ALGORITHM_VERSION, | |
| "state_sha256": hashlib.sha256(raw).hexdigest(), | |
| "row": identity, | |
| }, | |
| sort_keys=True, | |
| ).encode() | |
| ).hexdigest() | |
| cache = Path(cache_dir).resolve() | |
| cache.mkdir(parents=True, exist_ok=True) | |
| output = cache / (key + ".pt") | |
| with _cache_lock(cache / (key + ".lock")): | |
| if output.exists(): | |
| payload = torch.load(output, map_location="cpu", weights_only=True) | |
| validate_generated_labels(payload, num_players) | |
| else: | |
| payload = { | |
| k: torch.from_numpy(v) for k, v in generate_labels(state).items() | |
| } | |
| validate_generated_labels(payload, num_players) | |
| _atomic_save(payload, output) | |
| return str(output) | |
| def prepare_dataset_groups(dataset, sample_indices, cache_dir, model_profile): | |
| """Prepare only selected samples (including every constituent CSV row).""" | |
| if sample_indices is None: | |
| sample_indices = range(dataset.n_rows) | |
| sample_indices = tuple(sorted(set(sample_indices))) | |
| dataset.grouping_sample_indices = sample_indices | |
| rows = set() | |
| for index in sample_indices: | |
| if not 0 <= index < dataset.n_rows: | |
| raise ValueError(f"Grouping sample index out of range: {index}") | |
| start = dataset.sample_start_rows[index] | |
| rows.update(range(start, start + dataset.consecutive_rows)) | |
| paths = [None] * len(dataset.metadata) | |
| dataset.group_paths = paths | |
| for index in sorted(rows): | |
| row = dataset.metadata.iloc[index] | |
| existing = _optional_path(dataset.base_path, row.get(dataset.group_path_column)) | |
| if existing is None or not existing.exists(): | |
| if ( | |
| dataset.hier_group_format != "membership" | |
| or dataset.hier_group_mode != "csv" | |
| or dataset.group_label_player_ids | |
| != list(range(1, dataset.num_players + 1)) | |
| ): | |
| raise ValueError( | |
| "Automatic grouping requires CSV membership in compact player order" | |
| ) | |
| expected = (5, 6, 1.0, 3) | |
| if tuple(model_profile) != expected: | |
| raise ValueError( | |
| f"Automatic grouping requires local/global hubs, direct fraction, block = {expected}" | |
| ) | |
| if dataset.num_groups != (dataset.num_players - 1).bit_length(): | |
| raise ValueError( | |
| "Automatic grouping requires ceil(log2(num_players)) group slots" | |
| ) | |
| try: | |
| paths[index] = ensure_group_path( | |
| row, | |
| dataset.base_path, | |
| num_players=dataset.num_players, | |
| cache_dir=cache_dir, | |
| group_path_column=dataset.group_path_column, | |
| ) | |
| rgb_frames = min( | |
| int(row.get(f"video{i}_frame", dataset.target_rgb_frames)) | |
| for i in range(1, dataset.num_players + 1) | |
| ) | |
| # Reuse the loader's canonical checks, before __getitem__'s retry loop. | |
| dataset._load_group_labels(index, (rgb_frames - 1) // 4 + 1) | |
| except Exception as error: | |
| raise ValueError( | |
| f"Grouping preparation failed at CSV row {index} " | |
| f"(case {row.get('case_id', index)}): {error}" | |
| ) from error | |
| dataset.prepared_group_rows = frozenset(rows) | |