opensportslib 0.3.3.dev3__tar.gz → 0.3.4.dev1__tar.gz

{opensportslib-0.3.3.dev3/opensportslib.egg-info → opensportslib-0.3.4.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: opensportslib
-Version: 0.3.3.dev3
+Version: 0.3.4.dev1
 Summary: OpenSportsLib is the professional library, designed for advanced video understanding in sports. It provides state-of-the-art tools for action recognition, spotting, retrieval, and captioning, making it ideal for researchers, analysts, and developers working with sports video data.
 Author: Jeet Vora
 Requires-Python: >=3.12
@@ -199,6 +199,12 @@ and [VQA guide](docs/tools/vqa.md).
 
 OpenSportsLib provides APIs and scripts for downloading and uploading OSL datasets with Hugging Face.
 
+For `tracking_tracab`, both transfer formats include the original DAT and all
+referenced metadata JSON/TF05/TF08 XML files. Uploads reject missing tracking
+assets; full downloads validate completeness. Annotation-only downloads defer
+assets until a sample/input download or missing-file hydration is requested.
+See [TRACAB transfer behavior](docs/tracking_tracab.md#dataset-loading-and-persistence).
+
 For SN-GAR classification and action-spotting configurations, setup, caching,
 and training commands, see the [SN-GAR examples](examples/sngar/README.md).
 
@@ -330,3 +336,13 @@ For synchronized native-FPS GAR action clips and optional exact frame reproducti
 
 For full-rate tracking clips with the same classification sample IDs, see the
 [tracking conversion script](docs/tools/spotting-tracking-to-classification.md).
+
+
+### Raw TRACAB tracking
+
+`tracking_tracab` is supported by classification and localization tracking loaders
+for local OSL JSON, Hugging Face JSON and indexed WebDataset sources. Set the
+input modality to `tracking_tracab`; reference the original metadata JSON with
+`metadata_path` and retain UTC period calibration and team codes in each input. `params.annotation_head` selects a head such as
+`action` or `headers`. See [the DAT contract and loading guide](docs/tracking_tracab.md) for the
+schema, configuration, caching and backward-compatible Parquet behavior.

{opensportslib-0.3.3.dev3 → opensportslib-0.3.4.dev1}/README.md

@@ -160,6 +160,12 @@ and [VQA guide](docs/tools/vqa.md).
 
 OpenSportsLib provides APIs and scripts for downloading and uploading OSL datasets with Hugging Face.
 
+For `tracking_tracab`, both transfer formats include the original DAT and all
+referenced metadata JSON/TF05/TF08 XML files. Uploads reject missing tracking
+assets; full downloads validate completeness. Annotation-only downloads defer
+assets until a sample/input download or missing-file hydration is requested.
+See [TRACAB transfer behavior](docs/tracking_tracab.md#dataset-loading-and-persistence).
+
 For SN-GAR classification and action-spotting configurations, setup, caching,
 and training commands, see the [SN-GAR examples](examples/sngar/README.md).
 
@@ -291,3 +297,13 @@ For synchronized native-FPS GAR action clips and optional exact frame reproducti
 
 For full-rate tracking clips with the same classification sample IDs, see the
 [tracking conversion script](docs/tools/spotting-tracking-to-classification.md).
+
+
+### Raw TRACAB tracking
+
+`tracking_tracab` is supported by classification and localization tracking loaders
+for local OSL JSON, Hugging Face JSON and indexed WebDataset sources. Set the
+input modality to `tracking_tracab`; reference the original metadata JSON with
+`metadata_path` and retain UTC period calibration and team codes in each input. `params.annotation_head` selects a head such as
+`action` or `headers`. See [the DAT contract and loading guide](docs/tracking_tracab.md) for the
+schema, configuration, caching and backward-compatible Parquet behavior.

{opensportslib-0.3.3.dev3 → opensportslib-0.3.4.dev1}/opensportslib/apis/base_task_model.py

@@ -48,7 +48,7 @@ def _manifest_media_references(payload: dict[str, Any]):
         for input_obj in inputs:
             if not isinstance(input_obj, dict):
                 continue
-            for key in ("path", "ball_path"):
+            for key in ("path", "ball_path", "metadata_path", "tf05_path", "tf08_path"):
                 if isinstance(input_obj.get(key), str):
                     yield input_obj[key], lambda value, input_obj=input_obj, key=key: input_obj.__setitem__(key, value)
             values = input_obj.get("paths")

{opensportslib-0.3.3.dev3 → opensportslib-0.3.4.dev1}/opensportslib/apis/classification.py

@@ -23,6 +23,7 @@ def _is_tracking_graph_modality(modality):
     return str(modality).lower() in {
         "tracking",
         "tracking_parquet",
+        "tracking_tracab",
         "tracking_h5",
         "player_centroids_h5",
         "player_joints_h5",
@@ -49,6 +50,13 @@ class ClassificationModel(BaseTaskModel):
         if override is not None:
             return expand(override)
 
+        from opensportslib.datasets.hf_json import hf_json_source, prepare_hf_json_split
+        if not self.is_remote and hf_json_source(self.config):
+            from opensportslib.core.config.accessors import set_split_source_path
+            prepared = prepare_hf_json_split(self.config, split)
+            set_split_source_path(self.config, split, str(prepared.source_path))
+            return str(prepared.annotation_path)
+
         path = get_split_annotation_path(self.config, split)
         if path:
             return expand(path)

{opensportslib-0.3.3.dev3 → opensportslib-0.3.4.dev1}/opensportslib/apis/localization.py

@@ -59,6 +59,15 @@ class LocalizationModel(BaseTaskModel):
                 set_split_source_path(self.config, "valid_data_frames", str(prepared.source_path))
             return str(prepared.annotation_path)
 
+        from opensportslib.datasets.hf_tracking import hf_tracking_source, prepare_hf_tracking_split, stage_tracking_files
+        if hf_tracking_source(self.config):
+            prepared = prepare_hf_tracking_split(self.config, split)
+            root = stage_tracking_files(prepared)
+            set_split_source_path(self.config, split, str(root))
+            if split == "valid":
+                set_split_source_path(self.config, "valid_data_frames", str(root))
+            return str(prepared.annotations_path)
+
         path = get_split_annotation_path(self.config, split)
         if path:
             return expand(path)

{opensportslib-0.3.3.dev3 → opensportslib-0.3.4.dev1}/opensportslib/core/config/loader.py

@@ -145,7 +145,11 @@ def _torch_cuda_available() -> bool:
 
 
 def _dali_available() -> bool:
-    return importlib.util.find_spec("nvidia.dali") is not None
+    try:
+        return importlib.util.find_spec("nvidia.dali") is not None
+    except ModuleNotFoundError:
+        # find_spec raises when the optional parent namespace is absent.
+        return False
 
 
 # DALI decodes video; any other input modality (tracking parquets, features)

{opensportslib-0.3.3.dev3 → opensportslib-0.3.4.dev1}/opensportslib/core/trainer/classification_trainer.py

@@ -74,6 +74,7 @@ def _is_tracking_graph_modality(modality):
     return str(modality).lower() in {
         "tracking",
         "tracking_parquet",
+        "tracking_tracab",
         "tracking_h5",
         "player_centroids_h5",
         "player_joints_h5",

{opensportslib-0.3.3.dev3 → opensportslib-0.3.4.dev1}/opensportslib/core/utils/load_annotations.py

@@ -12,6 +12,14 @@ from opensportslib.core.utils.config import load_json
 from collections import defaultdict
 
 
+def select_annotation_head(document, requested=None):
+    heads = document["labels"]
+    selected = requested or ("action" if "action" in heads else next(iter(heads)) if len(heads) == 1 else None)
+    if selected not in heads:
+        raise ValueError(f"Select annotation_head from {list(heads)}; got {requested!r}")
+    return selected
+
+
 LOCALIZATION_ANNOTATION_STATUSES = {"verified", "unlabeled", "excluded"}
 
 
@@ -164,7 +172,7 @@ def expand_localization_intervals(record):
 
 def load_annotations(
     annotations_path,
-    task_key="action",
+    task_key=None,
     exclude_labels=None,
     multiview=False,
     input_type="video",
@@ -175,6 +183,10 @@ def load_annotations(
     with open(annotations_path, "r") as f:
         data = json.load(f)
 
+    if not isinstance(data["data"], list):
+        raise TypeError("OSL data must be a list")
+    task_key = select_annotation_head(data, task_key)
+
     # this is used for data slicing experiments.
     # and doesn't affect the validation and test sets.
     # if you haven't added "data_scaling" to the config, this will be ignored.
@@ -517,7 +529,7 @@ def annotationstoe2eformat(
     return labels_e2e, task_name_list[0]
 
 
-def annotationstoe2eformat_tracking(label_files, video_dirs, extract_fps):
+def annotationstoe2eformat_tracking(label_files, video_dirs, extract_fps, input_type="tracking_parquet", annotation_head=None):
     """
     Adapt tracking-parquet action-spotting annotations to E2E format.
 
@@ -546,6 +558,7 @@ def annotationstoe2eformat_tracking(label_files, video_dirs, extract_fps):
         classes use to slice windows out of the cached native features.
     """
     from opensportslib.datasets.utils.tracking import load_or_build_game_cache
+    from opensportslib.tracking.tracab import parse_utc
 
     if not isinstance(label_files, list):
         label_files = [label_files]
@@ -562,20 +575,26 @@ def annotationstoe2eformat_tracking(label_files, video_dirs, extract_fps):
 
         annotations = load_json(label_path)
 
-        for task_name, task_data in annotations["labels"].items():
-            labels = task_data.get("labels", {})
-            task_name_list.append(task_name)
+        task_name = select_annotation_head(annotations, annotation_head)
+        labels = annotations["labels"][task_name].get("labels", {})
+        task_name_list.append(task_name)
         classes_by_label_dir.append(labels)
 
         videos = annotations["data"]
 
         for video in tqdm.tqdm(videos):
-            video_input = video["inputs"][0]
+            candidates = [i for i in video["inputs"] if i.get("type") == input_type]
+            if len(candidates) != 1:
+                raise ValueError(f"Expected one {input_type} input for {video.get('id')!r}; found {len(candidates)}")
+            video_input = dict(candidates[0])
+            for key in ("metadata_path", "tf05_path", "tf08_path"):
+                if video_input.get(key):
+                    video_input[key] = os.path.abspath(os.path.join(video_dir, video_input[key]))
             video_path = video_input["path"]
             full_video_path = os.path.join(video_dir, video_path)
             assert os.path.isfile(full_video_path), full_video_path
 
-            _, times, _ = load_or_build_game_cache(full_video_path)
+            _, times, _ = load_or_build_game_cache(full_video_path, descriptor=video_input)
             num_frames = len(times)
 
             native_fps = video_input.get("fps")
@@ -591,7 +610,18 @@ def annotationstoe2eformat_tracking(label_files, video_dirs, extract_fps):
 
             events = []
             for ann in video.get("events", []):
-                position_ms = float(ann["position_ms"])
+                if ann.get("head", task_name) != task_name:
+                    continue
+                if ann["label"] not in labels:
+                    raise ValueError(f"Event label {ann['label']!r} absent from head {task_name!r}")
+                if ann.get("timestamp_utc"):
+                    if not video_input.get("UTC_time_start"):
+                        raise ValueError("UTC events require input.UTC_time_start")
+                    position_ms = (parse_utc(ann["timestamp_utc"]) - parse_utc(video_input["UTC_time_start"])).total_seconds() * 1000
+                else:
+                    position_ms = float(ann["position_ms"])
+                if position_ms < 0 or (video_input.get("duration_ms") and position_ms >= video_input["duration_ms"]):
+                    raise ValueError(f"Event timestamp {position_ms} outside tracking timeline")
                 native_idx = int(np.searchsorted(times, position_ms))
                 native_idx = min(max(native_idx, 0), num_frames - 1)
                 adj_frame = native_idx // stride
@@ -614,6 +644,7 @@ def annotationstoe2eformat_tracking(label_files, video_dirs, extract_fps):
             events.sort(key=lambda x: x["frame"])
 
             labels_e2e.append({
+                "input": dict(video_input),
                 "events": events,
                 "fps": sample_fps,
                 "num_frames": num_frames_after,
@@ -834,7 +865,7 @@ def check_config(cfg, split="train"):
         if get_loader_backend(cfg) == "dali":
             cfg.TRAIN.execution.repartitions = get_repartition_gpu(cfg.SYSTEM.gpu.count)
         primary_modality = get_data_modality(cfg)
-        assert primary_modality in ["rgb", "video", "tracking_parquet"]
+        assert primary_modality in ["rgb", "video", "tracking_parquet", "tracking_tracab"]
         backbone_type = get_component_name_by_kind(cfg, "encoder")
         assert backbone_type in [
             # From torchvision
@@ -900,8 +931,13 @@ def check_config(cfg, split="train"):
             and data_path.endswith(".json")
             and "labels" in load_json(data_path).keys()
         ):
-            for task_name, task_data in load_json(data_path)["labels"].items():
-                classes = task_data.get("labels", {})
+            document = load_json(data_path)
+            if primary_modality in {"tracking_parquet", "tracking_tracab"}:
+                head = select_annotation_head(document, get_data_params(cfg).get("annotation_head"))
+                classes = document["labels"][head].get("labels", {})
+            else:
+                for task_data in document["labels"].values():
+                    classes = task_data.get("labels", {})
         else:
             assert isinstance(cfg.DATA.common.classes, (list, ListConfig))
             classes = cfg.DATA.common.classes

{opensportslib-0.3.3.dev3 → opensportslib-0.3.4.dev1}/opensportslib/datasets/classification_dataset.py

@@ -64,7 +64,7 @@ def build(config, annotations_path, processor=None, split="train"):
     """
     modality = get_data_modality(config).lower()
 
-    if modality in ("tracking", "tracking_parquet"):
+    if modality in ("tracking", "tracking_parquet", "tracking_tracab"):
         from opensportslib.datasets.hf_tracking import hf_tracking_source
 
         if hf_tracking_source(config):
@@ -133,7 +133,8 @@ class ClassificationDataset(Dataset):
             annotation_input_type = "frames_npy"
 
         self.samples, self.label_map = load_annotations(
-            annotations_path, 
+            annotations_path,
+            task_key=data_params.get("annotation_head"),
             exclude_labels=self.exclude_labels, 
             multiview=is_multiview,
             input_type=annotation_input_type,
@@ -540,7 +541,7 @@ class TrackingDataset(ClassificationDataset):
                 continue
             
             clip_path = clip_paths[0]
-            df = self._load_tracking_clip(clip_path)
+            df = self._load_tracking_clip(clip_path, item["inputs"][0])
 
             num_frames = len(df)
             all_features = np.zeros(
@@ -568,7 +569,7 @@ class TrackingDataset(ClassificationDataset):
             # regardless of transforms.
             edge_indices = []
             for t in range(num_frames):
-                edge_index = self.build_edge_index(
+                edge_index = self._build_edge_index(
                     all_features[t],
                     all_positions[t],
                     self.edge_type,
@@ -587,7 +588,7 @@ class TrackingDataset(ClassificationDataset):
         
         print(f"Loaded {len(self.processed_samples)} {self.split} samples")
     
-    def _load_tracking_clip(self, path):
+    def _load_tracking_clip(self, path, descriptor=None):
         """read a single parquet tracking clip.
         
         Args:
@@ -600,6 +601,10 @@ class TrackingDataset(ClassificationDataset):
         import pandas as pd
 
         full_path = os.path.join(self.video_path, path)
+        if descriptor and descriptor.get("type") == "tracking_tracab":
+            from opensportslib.tracking import TracabTrackingReader
+            with TracabTrackingReader(full_path, descriptor, metadata_base_dir=self.video_path) as reader:
+                return pd.DataFrame(reader.iter_uniform_rows())
         return pd.read_parquet(full_path)
     
     def __getitem__(self, idx):
@@ -687,7 +692,7 @@ class TrackingDataset(ClassificationDataset):
             raise ValueError(f"No tracking paths found for item {idx}")
         
         clip_path = clip_paths[0]
-        df = self._load_tracking_clip(clip_path)
+        df = self._load_tracking_clip(clip_path, item["inputs"][0])
         
         num_frames = len(df)
         all_features = np.zeros(
@@ -765,8 +770,8 @@ class HFTrackingDataset(TrackingDataset):
         self._hf_reader = HFTarTrackingReader(prepared)
         super().__init__(config, str(prepared.annotations_path), split)
 
-    def _load_tracking_clip(self, path):
-        return self._hf_reader.read_parquet(path)
+    def _load_tracking_clip(self, path, descriptor=None):
+        return self._hf_reader.read_tracking(path, descriptor)
 
 
 class H5TrackingDataset(ClassificationDataset):

{opensportslib-0.3.3.dev3 → opensportslib-0.3.4.dev1}/opensportslib/datasets/hf_json.py

@@ -130,10 +130,13 @@ def prepare_hf_json_split(config, split: str) -> PreparedHFJsonSplit:
                         f"Expected one {media_type!r} input for sample {sample.get('game_id', sample.get('id'))!r}"
                     )
                 item = dict(selected[0])
-                filename = _safe_path(str(item.get("path") or ""))
-                if not filename.startswith(f"{split}/"):
-                    raise ValueError(f"Media path {filename!r} is outside split {split!r}")
-                filenames.add(filename)
+                for key in ("path", "metadata_path", "tf05_path", "tf08_path"):
+                    if key != "path" and not item.get(key):
+                        continue
+                    filename = _safe_path(str(item.get(key) or ""))
+                    if key == "path" and not filename.startswith(f"{split}/"):
+                        raise ValueError(f"Media path {filename!r} is outside split {split!r}")
+                    filenames.add(filename)
                 sample["inputs"] = [item]
 
             def stage_media(filename: str) -> None:

{opensportslib-0.3.3.dev3 → opensportslib-0.3.4.dev1}/opensportslib/datasets/hf_tracking.py

@@ -13,7 +13,7 @@ from collections import OrderedDict
 from dataclasses import dataclass
 from pathlib import Path
 
-from opensportslib.core.config.accessors import get_input_cfg
+from opensportslib.core.config.accessors import get_input_cfg, get_data_modality
 
 
 logger = logging.getLogger(__name__)
@@ -126,7 +126,14 @@ def prepare_hf_tracking_split(config, split: str) -> PreparedTrackingSplit:
 
     _set_resolved_revision(config, resolved)
     safe_repo = re.sub(r"[^A-Za-z0-9_.-]+", "--", repo_id)
+    media_type = str(source.get("input_type") or get_data_modality(config))
+    if media_type == "tracking":
+        media_type = "tracking_parquet"
+    if media_type not in {"tracking_parquet", "tracking_tracab"}:
+        raise ValueError(f"Unsupported tracking input type: {media_type}")
     directory = cache_dir / safe_repo / resolved / split
+    if media_type != "tracking_parquet":
+        directory = directory / media_type
     directory.mkdir(parents=True, exist_ok=True)
     annotations_path = directory / "annotations.json"
     index_path = directory / "index.json"
@@ -181,7 +188,7 @@ def prepare_hf_tracking_split(config, split: str) -> PreparedTrackingSplit:
 
             manifest = {}
             for row in rows("shard_manifest.parquet"):
-                if row.get("file_role") not in (None, "", "primary") or row.get("input_type") != "tracking_parquet":
+                if row.get("file_role") not in (None, "", "primary", "metadata", "tf05", "tf08") or row.get("input_type") != media_type:
                     continue
                 key = (str(row["sample_id"]), str(row["relative_path"]))
                 if row.get("status") == "ok" and row.get("wds_member"):
@@ -190,17 +197,20 @@ def prepare_hf_tracking_split(config, split: str) -> PreparedTrackingSplit:
             members = {}
             for sample_id, (shard_name, sample) in locations.items():
                 for input_item in sample.get("inputs", []):
-                    if input_item.get("type") != "tracking_parquet":
+                    if input_item.get("type") != media_type:
                         continue
-                    path = str(input_item.get("path") or "")
-                    location = manifest.get((sample_id, path))
-                    if not path or location is None or location[0] != shard_name:
-                        raise ValueError(f"Missing tracking TAR member for {sample_id!r}: {path!r}")
-                    if path in members and members[path] != location:
-                        raise ValueError(f"Tracking path maps to multiple TAR members: {path!r}")
-                    members[path] = location
+                    for key in ("path", "metadata_path", "tf05_path", "tf08_path"):
+                        if key != "path" and not input_item.get(key):
+                            continue
+                        path = str(input_item.get(key) or "")
+                        location = manifest.get((sample_id, path))
+                        if not path or location is None or location[0] != shard_name:
+                            raise ValueError(f"Missing tracking TAR member for {sample_id!r} ({key}): {path!r}")
+                        if path in members and members[path] != location:
+                            raise ValueError(f"Tracking path maps to multiple TAR members: {path!r}")
+                        members[path] = location
             if not members:
-                raise ValueError(f"{split} contains no tracking_parquet inputs")
+                raise ValueError(f"{split} contains no {media_type} inputs")
 
             _write_json_atomic(annotations_path, {**header, "data": samples})
             _write_json_atomic(index_path, {"members": members})
@@ -277,10 +287,33 @@ class HFTarTrackingReader:
         return tar
 
     def read_parquet(self, path: str):
+        return self.read_tracking(path)
+
+    def read_tracking(self, path: str, descriptor=None):
         import pandas as pd
 
         if path not in self.members:
             raise FileNotFoundError(f"Tracking clip is absent from Hub manifest: {path}")
+        metadata = None
+        xml_metadata = {}
+        for key in ("metadata_path", "tf05_path", "tf08_path"):
+            if not descriptor or not descriptor.get(key):
+                continue
+            metadata_path = descriptor[key]
+            if metadata_path not in self.members:
+                raise FileNotFoundError(f"Tracking metadata is absent from Hub manifest: {metadata_path}")
+            metadata_shard, metadata_member = self.members[metadata_path]
+            try:
+                from opensportslib.tracking.native import read_document
+                with self._tar(metadata_shard).extractfile(metadata_member) as stream:
+                    document = read_document(stream, format="json" if key == "metadata_path" else "xml")
+                    document.source = metadata_path
+                if key == "metadata_path":
+                    metadata = document
+                else:
+                    xml_metadata[key] = document
+            except (KeyError, ValueError, TypeError) as exc:
+                raise ValueError(f"TRACAB metadata {metadata_path}: {exc}") from exc
         shard_name, member_name = self.members[path]
         tar = self._tar(shard_name)
         try:
@@ -289,9 +322,42 @@ class HFTarTrackingReader:
             if stream is None:
                 raise KeyError(member_name)
             with stream:
+                if descriptor and descriptor.get("type") == "tracking_tracab":
+                    from opensportslib.tracking import TracabTrackingReader
+                    with TracabTrackingReader(stream, descriptor, metadata=metadata, xml_metadata=xml_metadata) as reader:
+                        return pd.DataFrame(reader.iter_uniform_rows())
                 contents = stream.read()
         except KeyError as exc:
             raise FileNotFoundError(
                 f"Tracking TAR member {member_name!r} is missing from {shard_name!r}"
             ) from exc
         return pd.read_parquet(io.BytesIO(contents))
+
+
+def stage_tracking_files(prepared):
+    """Materialize raw members once for localization's mmap feature cache."""
+    import shutil
+    from opensportslib.datasets.hf_json import _safe_path
+
+    root = prepared.annotations_path.parent / "media"
+    root.mkdir(exist_ok=True)
+    with (root / ".stage.lock").open("w") as lock:
+        fcntl.flock(lock, fcntl.LOCK_EX)
+        reader = HFTarTrackingReader(prepared)
+        try:
+            for path, (shard, member) in reader.members.items():
+                destination = root / _safe_path(path)
+                if destination.is_file():
+                    continue
+                destination.parent.mkdir(parents=True, exist_ok=True)
+                temporary = destination.with_name(destination.name + f".tmp-{os.getpid()}")
+                try:
+                    with reader._tar(shard).extractfile(member) as source, temporary.open("wb") as target:
+                        shutil.copyfileobj(source, target)
+                    os.replace(temporary, destination)
+                finally:
+                    temporary.unlink(missing_ok=True)
+        finally:
+            for tar in reader._open_tars.values():
+                tar.close()
+    return root

{opensportslib-0.3.3.dev3 → opensportslib-0.3.4.dev1}/opensportslib/datasets/localization_dataset.py

@@ -470,6 +470,8 @@ class LocalizationDataset(Dataset):
                 dilate_len=self.data_cfg.dilate_len,
                 augmentations=_build_tracking_spotting_augmentations(self.config),
                 normalize=normalize,
+                input_type=get_runtime_modality(self.config),
+                annotation_head=get_data_params(self.config).get("annotation_head"),
             )
         elif dataset_type == "TrackingActionSpotVideo":
             encoder_params = get_component_params_by_kind(self.config, "encoder")
@@ -485,6 +487,8 @@ class LocalizationDataset(Dataset):
                 r=encoder_params.get("radius"),
                 overlap_len=getattr(cfg, "overlap_len", self.data_cfg.clip_len // 2),
                 normalize=normalize,
+                input_type=get_runtime_modality(self.config),
+                annotation_head=get_data_params(self.config).get("annotation_head"),
             )
         else:
             dataset = None
@@ -1276,17 +1280,24 @@ class _TrackingGraphWindowMixin:
         """
         return batch["frame"].to(device)
 
-    def _get_game_cache(self, video_path):
+    def _get_game_cache(self, video_path, descriptor=None):
         """lazily load (and memoize for the lifetime of this dataset
         instance/worker) a game's cached whole-match features."""
-        if video_path not in self._game_cache:
-            self._game_cache[video_path] = self._load_or_build_game_cache(video_path)
-        return self._game_cache[video_path]
+        companion_signatures = []
+        if descriptor and descriptor.get("type") == "tracking_tracab":
+            for field in ("metadata_path", "tf05_path", "tf08_path"):
+                if descriptor.get(field):
+                    stat = os.stat(descriptor[field])
+                    companion_signatures.append((field, stat.st_size, stat.st_mtime_ns))
+        key = (video_path, json.dumps(descriptor, sort_keys=True), tuple(companion_signatures))
+        if key not in self._game_cache:
+            self._game_cache[key] = self._load_or_build_game_cache(video_path, descriptor=descriptor)
+        return self._game_cache[key]
 
     def _slice_window(self, video_meta, base_idx):
         """Fetch the game cache and slice a decimated feature window +
         per-frame position lists out of it (see tracking.slice_window)."""
-        features_cache, _, positions_cache = self._get_game_cache(video_meta["video"])
+        features_cache, _, positions_cache = self._get_game_cache(video_meta["video"], video_meta.get("input"))
         return self._slice_window_fn(
             features_cache, positions_cache, base_idx, self._clip_len, video_meta["stride"]
         )
@@ -1364,12 +1375,14 @@ class TrackingActionSpotDataset(Dataset, _TrackingGraphWindowMixin):
         fg_upsample=-1,
         augmentations=None,
         normalize=True,
+        input_type="tracking_parquet",
+        annotation_head=None,
     ):
         from opensportslib.core.utils.load_annotations import annotationstoe2eformat_tracking
 
         self._src_file = label_file
         self._labels, self.task_name = annotationstoe2eformat_tracking(
-            label_file, video_dir, extract_fps
+            label_file, video_dir, extract_fps, input_type, annotation_head
         )
         self._class_dict = _normalize_class_dict(classes)
         self._video_idxs = {x["video"]: i for i, x in enumerate(self._labels)}
@@ -1528,12 +1541,14 @@ class TrackingActionSpotVideoDataset(Dataset, _TrackingGraphWindowMixin, Dataset
         pad_len=5,
         skip_partial_end=True,
         normalize=True,
+        input_type="tracking_parquet",
+        annotation_head=None,
     ):
         from opensportslib.core.utils.load_annotations import annotationstoe2eformat_tracking
 
         self._src_file = label_file
         self._labels, self.task_name = annotationstoe2eformat_tracking(
-            label_file, video_dir, extract_fps
+            label_file, video_dir, extract_fps, input_type, annotation_head
         )
         self._class_dict = _normalize_class_dict(classes)
         self._video_idxs = {x["path"]: i for i, x in enumerate(self._labels)}