Data Extension API

Extension point. Dataset dispatch is task-owned. Preserve the output contract expected by the matching trainer when adding modality support.

Routes and classes

Module	Route/classes	Role
datasets.builder	build_dataset	Routes classification, localization, and vqa tasks.
datasets.classification_dataset	build, ClassificationDataset, VideoDataset, TrackingDataset, HFTrackingDataset, H5TrackingDataset	Classification modalities.
datasets.localization_dataset	LocalizationDataset and concrete spotting/feature/SoccerNet classes	Localization dataset selection and sampling.
datasets.vqa_dataset	VQADataset	VQA manifest rows and answer references.
datasets.utils.tracking / utils.h5_tracking	tracking transforms, graph helpers, and HDF5 reader/manifest helpers	Tracking and HDF5 data support.
datasets.hf_json, datasets.hf_tracking	prepared-split types and prepare_*_split	Optional Hub-backed split staging.

Dataset classes are internal implementations; the dispatch path is the extension point. The following sections document all non-private local dataset classes before the optional Hugging Face integration.

build_dataset(config, annotation_file=None, processor=None, split='train')

Return a dataset instance based on model type

Source code in opensportslib/datasets/builder.py

def build_dataset(config, annotation_file=None, processor=None, split="train"):
    """Return a dataset instance based on model type"""
    task = config.TASK.lower()

    if "classification" in task:
        from opensportslib.datasets import classification_dataset
        return classification_dataset.build(config, annotation_file, processor, split)

    elif "localization" in task:
        from opensportslib.datasets.localization_dataset import LocalizationDataset
        return LocalizationDataset(config, annotation_file, processor, split=split)

    elif task == "vqa":
        from opensportslib.datasets.vqa_dataset import VQADataset
        return VQADataset(config, annotation_file=annotation_file, split=split)

    else:
        raise ValueError(f"No dataset found for task: {task}")

Classification datasets

classification_dataset.build(...) chooses the configured classification modality. VideoDataset handles video/frames-style inputs; TrackingDataset handles tracking data; HFTrackingDataset and H5TrackingDataset are concrete tracking specializations. These classes are consumed by Trainer_Classification, so a new dataset must preserve the batch fields that its selected trainer reads.

classification dataset implmentations for video and tracking modalities.

provides three concrete dataset classes:

• VideoDataset (MVFoul, SN-GAR video)
• TrackingDataset (SN-GAR tracking / parquet)

both inherit from ClassificationDataset, which handles annotation loading, label mapping, and class-weight computation.

ClassificationDataset

Bases: Dataset

shared base for all classification datasets.

loads annotations, builds a label map, and exposes helpers for computing sample-level and class-level weights (useful for balanced sampling and weighted loss respectively).

Parameters:

Name	Type	Description	Default
config		the loaded YAML configuration.	required
annotations_path		path to the annotation JSON file.	required
split		one of "train", "valid", "test".	'train'

Source code in opensportslib/datasets/classification_dataset.py

class ClassificationDataset(Dataset):
    """shared base for all classification datasets.

    loads annotations, builds a label map, and exposes helpers for
    computing sample-level and class-level weights (useful for
    balanced sampling and weighted loss respectively).

    Args:
        config: the loaded YAML configuration.
        annotations_path: path to the annotation JSON file.
        split: one of "train", "valid", "test".
    """

    def __init__(self, config, annotations_path, processor, split="train"):
        self.config = config
        self.split = split
        self.exclude_labels = ["Unknown", "Dont know"]
        self.video_path = get_split_source_path(config, split)
        self.processor = None

        # view_type is optional; only MVFoul uses it as of now
        data_params = get_data_params(config)
        is_multiview = data_params.get("view_type") == "multi"

        allow_missing_labels = split in ["test", "infer"]

        # these lines of code are used for data scaling experiments.
        # if you want to check how the model performance changes with different 
        # number of games in the training set, you can use this code.
        # to use this, you need to add the following to the config:
        # DATA:
        #   data_slicing:
        #     enabled: true
        #     training_matches: <number of games to include in the training set>
        # we will refer to this as "data slicing" in the rest of the code.
        max_games = None
        slicing_cfg = data_params.get("data_slicing", {})
        if slicing_cfg and slicing_cfg.get("enabled", False) and split == "train":
            max_games = slicing_cfg.get("training_matches")

        annotation_input_type = str(get_data_modality(config)).lower()
        if annotation_input_type == "tracking":
            annotation_input_type = "tracking_parquet"
        elif annotation_input_type == "tracking_h5":
            annotation_input_type = get_data_params(config).get(
                "h5_input_type", "player_centroids_h5"
            )
        elif annotation_input_type in {"frames", "frames_npy"}:
            annotation_input_type = "frames_npy"

        self.samples, self.label_map = load_annotations(
            annotations_path, 
            exclude_labels=self.exclude_labels, 
            multiview=is_multiview,
            input_type=annotation_input_type,
            allow_missing_labels=allow_missing_labels,
            max_games=max_games
        )

        # this is used for quick testing of the model.
        # we can only use a small subset of data (ideally 100 samples) to 
        # test the overall integration quickly.
        # to use this, you need to add the following to the config:
        # DATA:
        #   max_samples: <number of samples to include in the training set>
        max_samples = data_params.get("max_samples")
        if max_samples:
            self.samples = self.samples[:max_samples]

        # invert to id -> name and propagate into the config so
        # downstream components (metrics, logging) can look it up.
        annotation_label_map = {v: k for k, v in self.label_map.items()}
        runtime = get_data_runtime(self.config)
        inference_source = str(runtime.get("inference_class_source") or "").strip().lower()
        configured_classes = list(get_data_classes(self.config))
        configured_label_map = {idx: name for idx, name in enumerate(configured_classes)}

        if split in {"test", "infer"} and inference_source in {"model", "local"} and configured_classes:
            self.label_map = configured_label_map
            set_data_num_classes(self.config, len(configured_classes))

            annotation_classes = [annotation_label_map[idx] for idx in sorted(annotation_label_map.keys())]
            if annotation_classes != configured_classes:
                logger.warning(
                    "Inference annotation labels differ from resolved %s classes. "
                    "Using resolved runtime classes for predictions.",
                    inference_source,
                )
                for sample in self.samples:
                    sample.pop("label", None)
            self.has_labels = len(self.samples) > 0 and "label" in self.samples[0]
        else:
            self.label_map = annotation_label_map
            set_data_classes(self.config, list(self.label_map.values()))
            self.has_labels = len(self.samples) > 0 and "label" in self.samples[0]

        print(len(get_data_classes(self.config)), "classes:", get_data_classes(self.config))
        print("Label Map : ", self.label_map)

    # -- Sampling / loss weights ------------------------------------------

    def _normalized_label_tensor(self):
        """Return labels as validated torch.long class ids."""
        normalized_labels = []

        for idx, item in enumerate(self.samples):
            if "label" not in item or item["label"] is None:
                raise ValueError(
                    f"Missing label for sample index {idx}; cannot compute class weights."
                )

            label = item["label"]

            if isinstance(label, torch.Tensor):
                if label.numel() != 1:
                    raise ValueError(
                        f"Expected scalar label at sample index {idx}, got shape {tuple(label.shape)}."
                    )
                label = label.item()

            try:
                label_float = float(label)
            except (TypeError, ValueError) as exc:
                raise ValueError(
                    f"Label at sample index {idx} is not numeric: {label!r}."
                ) from exc

            if not label_float.is_integer():
                raise ValueError(
                    f"Label at sample index {idx} must be an integer class id, got {label!r}."
                )

            normalized_labels.append(int(label_float))

        if not normalized_labels:
            raise ValueError("Cannot compute class weights from an empty dataset.")

        return torch.tensor(normalized_labels, dtype=torch.long)

    def get_sample_weights(self):
        """per-sample inverse-frequency weights for WeightedRandomSampler.

        Returns:
            torch.Tensor of length len(self) with one weight per sample.
        """
        labels = self._normalized_label_tensor()

        class_counts = torch.bincount(labels)
        class_weights = 1.0 / class_counts.float()
        sample_weights = torch.tensor(
            [class_weights[label] for label in labels], 
            dtype=torch.float
        )

        return sample_weights


    def get_class_weights(self, num_classes=None, normalize=True, sqrt=False):
        """per-class inverse-frequency weights for WeightedRandomSampler.

        Args:
            num_classes: if None, inferred from the label tensor.
            normalize: if True, weights are scaled so they sum to num_classes.
            sqrt: if True, use inverse square-root frequency instead of raw counts.

        Returns:
            torch.Tensor of shape (num_classes,).
        """
        labels = self._normalized_label_tensor()

        if num_classes is None:
            num_classes = int(labels.max().item() + 1)

        counts = torch.bincount(labels, minlength=num_classes).float()
        counts[counts == 0] = 1.0  # avoid division by zero for unseen classes

        weights = 1.0 / torch.sqrt(counts) if sqrt else 1.0 / counts

        if normalize:
            weights = weights / weights.sum() * num_classes

        return weights

    def __len__(self):
        return len(self.samples)

    def num_classes(self):
        return len(self.label_map)

get_class_weights(num_classes=None, normalize=True, sqrt=False)

per-class inverse-frequency weights for WeightedRandomSampler.

Parameters:

Name	Type	Description	Default
num_classes		if None, inferred from the label tensor.	None
normalize		if True, weights are scaled so they sum to num_classes.	True
sqrt		if True, use inverse square-root frequency instead of raw counts.	False

Returns:

Type	Description
	torch.Tensor of shape (num_classes,).

Source code in opensportslib/datasets/classification_dataset.py

def get_class_weights(self, num_classes=None, normalize=True, sqrt=False):
    """per-class inverse-frequency weights for WeightedRandomSampler.

    Args:
        num_classes: if None, inferred from the label tensor.
        normalize: if True, weights are scaled so they sum to num_classes.
        sqrt: if True, use inverse square-root frequency instead of raw counts.

    Returns:
        torch.Tensor of shape (num_classes,).
    """
    labels = self._normalized_label_tensor()

    if num_classes is None:
        num_classes = int(labels.max().item() + 1)

    counts = torch.bincount(labels, minlength=num_classes).float()
    counts[counts == 0] = 1.0  # avoid division by zero for unseen classes

    weights = 1.0 / torch.sqrt(counts) if sqrt else 1.0 / counts

    if normalize:
        weights = weights / weights.sum() * num_classes

    return weights

get_sample_weights()

per-sample inverse-frequency weights for WeightedRandomSampler.

Returns:

Type	Description
	torch.Tensor of length len(self) with one weight per sample.

Source code in opensportslib/datasets/classification_dataset.py

def get_sample_weights(self):
    """per-sample inverse-frequency weights for WeightedRandomSampler.

    Returns:
        torch.Tensor of length len(self) with one weight per sample.
    """
    labels = self._normalized_label_tensor()

    class_counts = torch.bincount(labels)
    class_weights = 1.0 / class_counts.float()
    sample_weights = torch.tensor(
        [class_weights[label] for label in labels], 
        dtype=torch.float
    )

    return sample_weights

VideoDataset

Bases: ClassificationDataset

frame-sampled video clips for classification.

for MVFoul: supports single-view and multi-view modes. In multi-view training, two views are randomly sampled per clip; at test time all available views are returned and stacked along a view dimension.

Parameters:

Name	Type	Description	Default
config		the loaded YAML configuration.	required
annotations_path		path to the annotation JSON file.	required
processor		HuggingFace image processor (used only for HuggingFace models).	required
split		one of "train", "valid", "test".	'train'

Source code in opensportslib/datasets/classification_dataset.py

class VideoDataset(ClassificationDataset):
    """frame-sampled video clips for classification.

    for MVFoul: supports single-view and multi-view modes. In multi-view 
    training, two views are randomly sampled per clip; at test time all available
    views are returned and stacked along a view dimension.

    Args:
        config: the loaded YAML configuration.
        annotations_path: path to the annotation JSON file.
        processor: HuggingFace image processor (used only for HuggingFace models).
        split: one of "train", "valid", "test".
    """

    def __init__(self, config, annotations_path, processor, split="train"):
        super().__init__(config, annotations_path, processor, split=split)

        self.processor = processor
        self.view_type = get_data_params(config).get("view_type", "single")
        sampling_cfg = get_data_sampling(config)
        self.num_frames = sampling_cfg.get("num_frames")
        self.input_fps = sampling_cfg.get("input_fps")
        self.transform = build_transform(config, mode=self.split)

    def _select_views(self, video_paths):
        """ choose which camera views to laod for this sample.

        Args:
            video_paths: list of available view paths for the clip.

        Returns:
            a (possibly subsampled) list of paths.

        """
        if self.view_type == "single":
            return [video_paths[0]]

        if self.split.lower() == "train" and self.view_type == "multi":
            return random.sample(video_paths, min(2, len(video_paths)))

        return video_paths

    def _load_and_sample_clip(self, path):
        """read a video file, temporally sub-sample, and apply transforms.

        Args:
            path: realtive path (under video_path) to the video file.

        Returns:
            numpy.ndarray of shape (T, H, W, C).
        """
        if path.endswith(".npy"):
            frames = np.load(os.path.join(self.video_path, path)).astype(np.float32) / 255.0
            if self.transform is not None:
                frames = self.transform(frames)
            mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
            std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
            frames = (frames - mean) / std
            return frames

        v = read_video(os.path.join(self.video_path, path))

        v = process_frames(
            v,
            self.num_frames,
            self.input_fps,
            get_data_sampling(self.config).get("target_fps"),
            start_frame=get_data_sampling(self.config).get("start_frame"),
            end_frame=get_data_sampling(self.config).get("end_frame"),
        )

        if isinstance(v, list):
            v = np.stack(v)  # (T, H, W, C)

        if self.transform is not None:
            v = self.transform(v)

        return v  # (T, H, W, C)

    def __getitem__(self, idx):
        item = self.samples[idx]
        label = item.get("label", None)
        if label is not None:
            label = torch.tensor(label, dtype=torch.long)
        video_paths = item["video_paths"]
        sample_id = item["id"]

        # --- Choose which clips to load ---
        if not video_paths:
            raise ValueError(f"No video paths found for item {idx}")

        selected_paths = self._select_views(video_paths)

        # --- Load and process frames for selected clips ---
        if get_component_provider_by_kind(self.config, "encoder") == "huggingface":
            path = selected_paths[0]
            v = self._load_and_sample_clip(path)
            # convert clip -> list of frames
            v = list(v)   # each element is (H, W, C)

            #print(type(v), v)
            v = self.processor(v, return_tensors="pt")#, do_rescale=False) 
            pixel_values = v["pixel_values"].float()
            pixel_values = pixel_values.squeeze(0)
            out = {"pixel_values": pixel_values, "id": sample_id}
            if label is not None:
                out["labels"] = label
            return out

        else:
            view_tensors = []
            for path in selected_paths:
                v = self._load_and_sample_clip(path)

                if path.endswith(".npy"):
                    # frames_npy
                    v = torch.from_numpy(v)  # (T, H, W, C)
                else:
                    # existing raw video path: apply torchvision model transforms
                    v = torch.from_numpy(v).permute(0, 3, 1, 2)  # (T, C, H, W)
                    encoder_name = get_component_params_by_kind(self.config, "encoder").get(
                        "pretrained_model",
                        get_component_name_by_kind(self.config, "encoder"),
                    )
                    v = get_transforms_model(encoder_name)(v)  # (C, T, H, W)

                view_tensors.append(v)

            if selected_paths[0].endswith(".npy"):
                # frames_npy: single view, return (T, H, W, C) matching pixels_vs_positions
                out = {"pixel_values": view_tensors[0], "id": sample_id}
            else:
                # existing multi-view path: stack to (V, C, T, H, W)
                videos = torch.stack(view_tensors, dim=0)
                out = {"pixel_values": videos, "id": sample_id}

            if label is not None:
                out["labels"] = label
            return out

TrackingDataset

Bases: ClassificationDataset

graph-based classification dataset built from player tracking data.

each sample is a temporal sequence of per-frame graphs where nodes represent the ball and 22 players, and edges encode spatial or tactical relationships (see build_edge_index).

supports optional preloading of all clips into memory and training-time augmentation (horizontal/vertical/team flip).

Parameters:

Name	Type	Description	Default
config		the loaded YAML configuration.	required
annotations_path		path to the annotation JSON file.	required
split		one of "train", "valid", "test".	'train'

Source code in opensportslib/datasets/classification_dataset.py

class TrackingDataset(ClassificationDataset):
    """graph-based classification dataset built from player tracking data.

    each sample is a temporal sequence of per-frame graphs where nodes
    represent the ball and 22 players, and edges encode spatial or
    tactical relationships (see build_edge_index).

    supports optional preloading of all clips into memory and 
    training-time augmentation (horizontal/vertical/team flip).

    Args:
        config: the loaded YAML configuration.
        annotations_path: path to the annotation JSON file.
        split: one of "train", "valid", "test".
    """

    def __init__(self, config, annotations_path, split="train"):
        super().__init__(config, annotations_path, processor=None, split=split)

        from opensportslib.datasets.utils.tracking import (
            FEATURE_DIM,
            NUM_OBJECTS,
            HorizontalFlip,
            TeamFlip,
            VerticalFlip,
            build_edge_index,
        )

        # storing references for the constants without repeating the import.
        self._build_edge_index = build_edge_index

        sampling_cfg = get_data_sampling(config)
        transform_cfg = get_data_transform(config)
        params_cfg = get_data_params(config)
        encoder_cfg = get_component_params_by_kind(config, "encoder")

        self.num_frames = sampling_cfg.get("num_frames")
        self.normalize = transform_cfg.get("normalize", False)
        self.edge_type = encoder_cfg.get("edge_type")
        self.k = encoder_cfg.get("k")
        self.r = encoder_cfg.get("radius", encoder_cfg.get("r"))
        self.preload_data = params_cfg.get("preload_data", False)
        objects_cfg = params_cfg.get("objects", {}) or {}
        # Fall back to the canonical layout constants rather than to these
        # attributes themselves - they have no class-level definition, so the
        # self-referencing default raised AttributeError and made this dataset
        # impossible to construct.
        self._NUM_OBJECTS = int(objects_cfg.get("num_objects", NUM_OBJECTS))
        self._FEATURE_DIM = int(objects_cfg.get("feature_dim", FEATURE_DIM))
        self._pitch_half_length = float(objects_cfg.get("pitch_half_length", 85.0))
        self._pitch_half_width = float(objects_cfg.get("pitch_half_width", 50.0))
        self._max_displacement = float(objects_cfg.get("max_displacement", 110.0))
        self._max_ball_height = float(objects_cfg.get("max_ball_height", 30.0))

        self.transforms = self._build_transforms(
            config, split, HorizontalFlip, VerticalFlip, TeamFlip
        )

        self.processed_samples = None
        if self.preload_data:
            self._preload_all_data()

    @staticmethod
    def _build_transforms(config, split, HorizontalFlip, VerticalFlip, TeamFlip):
        """assmeble the list of training-time augmentations.

        Args:
            config: the loaded YAML configuration.
            split: dataset split; augmentations are only applied during training.
            HorizontalFlip: augmentation class (passed to avoid re-importing).
            VerticalFlip: augmentation class.
            TeamFlip: augmentation class.

        Returns:
            list of callable augmentation transforms (empty for
            non-training splits).
        """
        if split != "train":
            return []

        transforms = []
        aug_config = get_data_augmentations(config)

        # augmentation flags are optional in the config; default to off.
        if aug_config.get("horizontal_flip", False):
            transforms.append(HorizontalFlip(probability=0.5))

        if aug_config.get("vertical_flip", False):
            transforms.append(VerticalFlip(probability=0.5))

        if aug_config.get("team_flip", False):
            transforms.append(TeamFlip(probability=0.5))

        return transforms

    def _preload_all_data(self):
        """parse every clip and cache features and edge indices in memory.

        edge indices are built on raw (un-augmented, un-normalized)
        features so that the graph topology is deterministic and 
        augmentation-independent.
        """
        from tqdm import tqdm

        from opensportslib.datasets.utils.tracking import (
            compute_deltas,
            parse_frame
        )

        print(f"Preloading {len(self.samples)} {self.split} samples into memory...")

        self.processed_samples = []

        for item in tqdm(self.samples, desc=f"Loading {self.split}"):
            clip_paths = item["video_paths"]
            if not clip_paths:
                continue

            clip_path = clip_paths[0]
            df = self._load_tracking_clip(clip_path)

            num_frames = len(df)
            all_features = np.zeros(
                (num_frames, self._NUM_OBJECTS, self._FEATURE_DIM), 
                dtype=np.float32
            )
            all_positions = []

            for t, (_, row) in enumerate(df.iterrows()):
                features, positions = parse_frame(
                    row,
                    num_objects=self._NUM_OBJECTS,
                    feature_dim=self._FEATURE_DIM,
                )
                all_features[t] = features
                all_positions.append(positions)

            all_features = compute_deltas(
                all_features,
                num_objects=self._NUM_OBJECTS,
            )

            # build edge indices on raw features (before any augmentation
            # or normalization) so the graph topology stays consistent
            # regardless of transforms.
            edge_indices = []
            for t in range(num_frames):
                edge_index = self.build_edge_index(
                    all_features[t],
                    all_positions[t],
                    self.edge_type,
                    self.k,
                    self.r
                )
                edge_indices.append(edge_index)

            self.processed_samples.append({
                "features": all_features,
                "positions": all_positions,
                "edge_indices": edge_indices,
                "label": item["label"],
                "id": item["id"]
            })

        print(f"Loaded {len(self.processed_samples)} {self.split} samples")

    def _load_tracking_clip(self, path):
        """read a single parquet tracking clip.

        Args:
            path: Relative path (under ``video_path``) to the parquet
                file.

        Returns:
            ``pandas.DataFrame`` with one row per frame.
        """
        import pandas as pd

        full_path = os.path.join(self.video_path, path)
        return pd.read_parquet(full_path)

    def __getitem__(self, idx):
        if self.preload_data:
            return self._getitem_preloaded(idx)
        else:
            return self._getitem_on_the_fly(idx)

    def _getitem_preloaded(self, idx):
        """return a sample from the in-memory cache.

        a copy of the feature array is made before augmentation and 
        normalization so the cached data is never mutated.
        """
        try:
            from torch_geometric.data import Data
        except ImportError as exc:
            raise ImportError(
                "torch-geometric is required for tracking_collate_fn. "
                "Run: `opensportslib setup --pyg` to install the correct version "
                "based on your system (PyTorch & CUDA compatible)."
            ) from exc

        from opensportslib.datasets.utils.tracking import normalize_features

        sample = self.processed_samples[idx]
        features = sample["features"].copy()

        for transform in self.transforms:
            features = transform(features)

        if self.normalize:
            features = normalize_features(
                features,
                pitch_half_length=self._pitch_half_length,
                pitch_half_width=self._pitch_half_width,
                max_displacement=self._max_displacement,
                max_ball_height=self._max_ball_height,
            )

        # build one PyG Data object per frame. The downstream collate function
        # (tracking_collate) uses PyG Batch.from_data_list to merge these across
        # the batch dimension.
        graphs = []
        for t in range(features.shape[0]):
            data = Data(
                x=torch.tensor(features[t], dtype=torch.float),
                edge_index=torch.tensor(
                    sample["edge_indices"][t], dtype=torch.long
                ),
            )
            graphs.append(data)

        out = {
            "graphs": graphs,
            "seq_len": len(graphs),
            "id": sample["id"]
        }
        if "label" in sample:
            out["label"] = sample["label"]
        return out

    def _getitem_on_the_fly(self, idx):
        """load, parse, and process a single sample from disk."""
        try:
            from torch_geometric.data import Data
        except ImportError as exc:
            raise ImportError(
                "torch-geometric is required for tracking_collate_fn. "
                "Run: `opensportslib setup --pyg` to install the correct version "
                "based on your system (PyTorch & CUDA compatible)."
            ) from exc

        from opensportslib.datasets.utils.tracking import (
            compute_deltas,
            normalize_features,
            parse_frame,
        )

        item = self.samples[idx]
        label = item["label"]

        clip_paths = item["video_paths"]
        if not clip_paths:
            raise ValueError(f"No tracking paths found for item {idx}")

        clip_path = clip_paths[0]
        df = self._load_tracking_clip(clip_path)

        num_frames = len(df)
        all_features = np.zeros(
            (num_frames, self._NUM_OBJECTS, self._FEATURE_DIM), 
            dtype=np.float32
        )
        all_positions = []

        for t, (_, row) in enumerate(df.iterrows()):
            features, positions = parse_frame(
                row,
                num_objects=self._NUM_OBJECTS,
                feature_dim=self._FEATURE_DIM,
            )
            all_features[t] = features
            all_positions.append(positions)

        all_features = compute_deltas(
            all_features,
            num_objects=self._NUM_OBJECTS,
        )

        # edge indices are built on raw features (before any augmentation /
        # normalization) so the graph structure is augmentation-invariant.
        edge_indices = []
        for t in range(num_frames):
            edge_index = self._build_edge_index(
                all_features[t],
                all_positions[t],
                self.edge_type,
                self.k,
                self.r
            )
            edge_indices.append(edge_index)

        for transform in self.transforms:
            all_features = transform(all_features)

        if self.normalize:
            all_features = normalize_features(
                all_features,
                pitch_half_length=self._pitch_half_length,
                pitch_half_width=self._pitch_half_width,
                max_displacement=self._max_displacement,
                max_ball_height=self._max_ball_height,
            )

        graphs = []
        for t in range(num_frames):
            data = Data(
                x=torch.tensor(all_features[t], dtype=torch.float),
                edge_index=torch.tensor(edge_indices[t], dtype=torch.long),
            )
            graphs.append(data)

        out = {
            "graphs": graphs,
            "seq_len": len(graphs),
            "id": item["id"]
        }
        if label is not None:
            out["label"] = label
        return out

HFTrackingDataset

Bases: TrackingDataset

Tracking graphs backed by indexed Hugging Face WebDataset TAR shards.

Source code in opensportslib/datasets/classification_dataset.py

class HFTrackingDataset(TrackingDataset):
    """Tracking graphs backed by indexed Hugging Face WebDataset TAR shards."""

    def __init__(self, config, split="train"):
        from opensportslib.datasets.hf_tracking import HFTarTrackingReader, prepare_hf_tracking_split

        if get_data_params(config).get("preload_data", False):
            raise ValueError("hf_webdataset requires DATA.inputs.tracking.params.preload_data=false")
        prepared = prepare_hf_tracking_split(config, split)
        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)

H5TrackingDataset

Bases: ClassificationDataset

Graph-based classification dataset for UTC-indexed player/ball H5 files.

Source code in opensportslib/datasets/classification_dataset.py

class H5TrackingDataset(ClassificationDataset):
    """Graph-based classification dataset for UTC-indexed player/ball H5 files."""

    def __init__(self, config, annotations_path, split="train"):
        super().__init__(config, annotations_path, processor=None, split=split)

        from opensportslib.datasets.utils.tracking import build_edge_index

        self._build_edge_index = build_edge_index
        sampling_cfg = get_data_sampling(config)
        transform_cfg = get_data_transform(config)
        params_cfg = get_data_params(config)
        encoder_cfg = get_component_params_by_kind(config, "encoder")
        objects_cfg = params_cfg.get("objects", {}) or {}

        self.modality = get_data_modality(config).lower()
        if self.modality == "tracking_h5":
            self.modality = params_cfg.get("h5_input_type", "player_centroids_h5")

        self.num_frames = sampling_cfg.get("num_frames")
        self.stride_ms = sampling_cfg.get("stride_ms")
        self.target_fps = sampling_cfg.get("target_fps")
        self.normalize = transform_cfg.get("normalize", False)
        self.edge_type = encoder_cfg.get("edge_type", "full")
        self.k = encoder_cfg.get("k")
        self.r = encoder_cfg.get("radius", encoder_cfg.get("r"))
        self.timestamp_tolerance_ms = float(params_cfg.get("timestamp_tolerance_ms", 50.0))
        self.missing_value = float(params_cfg.get("missing_value", -200.0))
        self._pitch_half_length = float(objects_cfg.get("pitch_half_length", 85.0))
        self._pitch_half_width = float(objects_cfg.get("pitch_half_width", 50.0))
        self._max_ball_height = float(objects_cfg.get("max_ball_height", 30.0))
        self._max_displacement = float(objects_cfg.get("max_displacement", 110.0))
        self.preload_data = params_cfg.get("preload_data", False)

        self.processed_samples = None
        if self.preload_data:
            self._preload_all_data()

    def _resolve_input_path(self, path):
        if os.path.isabs(path):
            return path
        return os.path.join(self.video_path, path)

    def _sample_input(self, item):
        inputs = item.get("inputs") or []
        for input_obj in inputs:
            if input_obj.get("type") == self.modality:
                return input_obj
        if inputs:
            return inputs[0]
        clip_paths = item.get("video_paths") or []
        if clip_paths:
            return {"type": self.modality, "path": clip_paths[0]}
        raise ValueError(f"No H5 input found for sample {item.get('id', '<unknown>')}")

    def _load_h5_frames(self, item):
        from opensportslib.datasets.utils.h5_tracking import (
            H5_FEATURE_DIM,
            H5TrackingReader,
            normalize_h5_features,
        )

        input_obj = self._sample_input(item)
        player_path = self._resolve_input_path(input_obj["path"])
        ball_path = input_obj.get("ball_path")
        if ball_path:
            ball_path = self._resolve_input_path(ball_path)

        reader = H5TrackingReader(
            player_path,
            input_obj.get("type", self.modality),
            ball_path=ball_path,
            missing_value=self.missing_value,
            timestamp_tolerance_ms=self.timestamp_tolerance_ms,
        )
        metadata = item.get("metadata", {}) or {}
        timestamps = reader.select_timestamps(
            start_utc=metadata.get("start_utc"),
            end_utc=metadata.get("end_utc"),
            num_frames=self.num_frames,
            stride_ms=self.stride_ms,
            target_fps=self.target_fps,
        )
        frames = reader.read_sequence(timestamps)

        if self.normalize and frames:
            frames = [
                frame.__class__(
                    frame.timestamp,
                    normalize_h5_features(
                        frame.features[None, :, :],
                        pitch_half_length=self._pitch_half_length,
                        pitch_half_width=self._pitch_half_width,
                        max_height=self._max_ball_height,
                        max_displacement=self._max_displacement,
                        missing_value=self.missing_value,
                    )[0],
                    frame.entity_ids,
                    frame.positions,
                )
                for frame in frames
            ]

        return frames, H5_FEATURE_DIM

    def _preload_all_data(self):
        print(f"Preloading {len(self.samples)} {self.split} H5 samples into memory...")
        self.processed_samples = [self._materialize_sample(item) for item in tqdm(self.samples)]

    def __getitem__(self, idx):
        if self.preload_data:
            return self.processed_samples[idx]
        return self._materialize_sample(self.samples[idx])

    def _materialize_sample(self, item):
        try:
            from torch_geometric.data import Data
        except ImportError as exc:
            raise ImportError(
                "torch-geometric is required for H5 tracking graph inputs. "
                "Run: `opensportslib setup --pyg` to install the correct version "
                "based on your system (PyTorch & CUDA compatible)."
            ) from exc

        frames, _ = self._load_h5_frames(item)
        graphs = []
        for frame in frames:
            edge_index = self._build_edge_index(
                frame.features,
                frame.positions,
                self.edge_type,
                self.k,
                self.r,
            )
            graphs.append(
                Data(
                    x=torch.tensor(frame.features, dtype=torch.float),
                    edge_index=torch.tensor(edge_index, dtype=torch.long),
                )
            )

        out = {
            "graphs": graphs,
            "seq_len": len(graphs),
            "id": item["id"],
        }
        label = item.get("label", None)
        if label is not None:
            out["label"] = label
        return out

build(config, annotations_path, processor=None, split='train')

construct the appropriate dataset for the configured modality.

Parameters:

Name	Type	Description	Default
config		the loaded YAML configuration.	required
annotations_path		path to the annotation JSON file.	required
processor		HuggingFace image processor (video modality only).	None
split		one of "train", "valid", "test".	'train'

Returns:

Type	Description
	a Dataset instance (VideoDataset or TrackingDataset).

Raises:

Type	Description
ValueError	if the data_modality is not recognized.

Source code in opensportslib/datasets/classification_dataset.py

def build(config, annotations_path, processor=None, split="train"):
    """construct the appropriate dataset for the configured modality.

    Args:
        config: the loaded YAML configuration.
        annotations_path: path to the annotation JSON file.
        processor: HuggingFace image processor (video modality only).
        split: one of "train", "valid", "test".

    Returns:
        a Dataset instance (VideoDataset or TrackingDataset).

    Raises:
        ValueError: if the data_modality is not recognized.
    """
    modality = get_data_modality(config).lower()

    if modality in ("tracking", "tracking_parquet"):
        from opensportslib.datasets.hf_tracking import hf_tracking_source

        if hf_tracking_source(config):
            return HFTrackingDataset(config, split=split)
        return TrackingDataset(config, annotations_path, split)
    elif modality in H5_TRACKING_MODALITIES:
        return H5TrackingDataset(config, annotations_path, split)
    elif modality in ("video", "frames_npy", "frames"):
        return VideoDataset(config, annotations_path, processor, split)
    else:
        raise ValueError(f"Unknown data_modality: {modality}")

Localization datasets

LocalizationDataset is the top-level localization entry point. It selects video, tracking, feature, and SoccerNet-oriented concrete datasets based on the configured split type and input modality. FrameReader and DatasetVideoSharedMethods are shared implementation support for video routes.

Dataset group	Concrete classes
Spotting/video	ActionSpotDataset, ActionSpotVideoDataset
Tracking spotting	TrackingActionSpotDataset, TrackingActionSpotVideoDataset
Feature JSON	FeaturefromJson, FeatureClipsfromJSON, FeatureClipChunksfromJson
SoccerNet	SoccerNetGame, SoccerNetGameClips, SoccerNetGameClipsChunks, SoccerNet, SoccerNetClips, SoccerNetClipsChunks

LocalizationDataset

Bases: Dataset

Source code in opensportslib/datasets/localization_dataset.py

class LocalizationDataset(Dataset):
    def __init__(self, config, annotations_path=None, processor=None, split="train"):
        self.config = config
        self.split = split
        self.gpu_count = get_system_gpu_count(self.config)

        execution_cfg = getattr(getattr(self.config, "TRAIN", None), "execution", None)
        repartitions = get_repartition_gpu(self.gpu_count)
        if execution_cfg is not None:
            if isinstance(execution_cfg, dict):
                execution_cfg["repartitions"] = repartitions
            else:
                execution_cfg.repartitions = repartitions

        sampling = get_data_sampling(self.config)
        transform = get_data_transform(self.config)
        normalization = transform.get("normalization", {})
        resize = transform.get("resize", {})
        params = get_data_params(self.config)
        encoder_params = get_component_params_by_kind(self.config, "encoder")
        adapter_params = get_component_params_by_kind(self.config, "adapter")
        split_cfg = get_split_cfg(self.config, split)
        split_dataloader = get_split_dataloader_cfg(self.config, split)
        self.data_cfg = SimpleNamespace(
            modality=get_runtime_modality(self.config),
            epoch_num_frames=sampling.get("epoch_num_frames"),
            clip_len=sampling.get("clip_len"),
            input_fps=sampling.get("input_fps"),
            extract_fps=sampling.get("extract_fps"),
            crop_dim=params.get("crop_dim"),
            dilate_len=params.get("dilate_len"),
            mixup=params.get("mixup"),
            imagenet_mean=normalization.get("mean"),
            imagenet_std=normalization.get("std"),
            target_height=resize.get("height"),
            target_width=resize.get("width"),
            preserve_aspect_ratio=resize.get("preserve_aspect_ratio", False),
        )
        annotation_path = annotations_path or getattr(
            split_cfg, "annotation_path", getattr(split_cfg, "path", None)
        )
        source_path = getattr(split_cfg, "source_path", getattr(split_cfg, "video_path", None))
        dataset_type = getattr(split_cfg, "type", None)
        dataset_split = getattr(split_cfg, "split", None)
        output_map = getattr(split_cfg, "output_map", ["data", "label"])
        overlap_len = getattr(split_cfg, "overlap_len", 0)
        feature_name = getattr(split_cfg, "features", encoder_params.get("encoder"))
        framerate = getattr(
            split_cfg,
            "framerate",
            sampling.get("framerate", self.data_cfg.extract_fps),
        )
        window_size = getattr(
            split_cfg,
            "window_size",
            sampling.get("window_size", encoder_params.get("window_size")),
        )
        chunk_size = getattr(
            split_cfg,
            "chunk_size",
            sampling.get("chunk_size", adapter_params.get("chunk_size")),
        )
        receptive_field = getattr(
            split_cfg,
            "receptive_field",
            sampling.get("receptive_field", adapter_params.get("receptive_field")),
        )
        chunks_per_epoch = getattr(
            split_cfg,
            "chunks_per_epoch",
            sampling.get("chunks_per_epoch", 6000),
        )

        # Keep a normalized split spec so runtime logic can stay canonical-first.
        self.cfg = SimpleNamespace(
            split_name=split,
            dataset_type=dataset_type,
            dataset_split=dataset_split,
            annotation_path=annotation_path,
            source_path=source_path,
            classes=getattr(split_cfg, "classes", get_data_classes(self.config)),
            dataloader=split_dataloader,
            output_map=output_map,
            overlap_len=overlap_len,
            features=feature_name,
            version=getattr(split_cfg, "version", 2),
            framerate=framerate,
            window_size=window_size,
            chunk_size=chunk_size,
            receptive_field=receptive_field,
            chunks_per_epoch=chunks_per_epoch,
            results=getattr(split_cfg, "results", None),
            metric=getattr(split_cfg, "metric", None),
            nms_window=getattr(split_cfg, "nms_window", None),
            infer_split=getattr(split_cfg, "infer_split", None),
            challenge=getattr(split_cfg, "challenge", None),
            path=annotation_path,
            video_path=source_path,
            type=dataset_type,
            split=dataset_split,
        )

        if split == "train":
            self.default_args = get_default_args_dataset("train", self.config)
        elif split == "valid":
            self.default_args = get_default_args_dataset("valid", self.config)
        elif split == "test":
            self.default_args = get_default_args_dataset("test", self.config)
        elif split == "valid_data_frames":
            self.default_args = get_default_args_dataset("valid_data_frames", self.config)
        #self.built_dataset = self.building_dataset(cfg=cfg, default_args=default_args)
        #self.data_loader = self.building_dataloader(self.built_dataset, cfg=cfg.dataloader, gpu=0, dali=True)


    def building_dataset(self, cfg, gpu=None, default_args=None):
        dataset_type = getattr(cfg, "dataset_type", getattr(cfg, "type", None))
        dataset_split = getattr(cfg, "dataset_split", getattr(cfg, "split", None))
        annotation_path = getattr(cfg, "annotation_path", getattr(cfg, "path", None))
        source_path = getattr(cfg, "source_path", getattr(cfg, "video_path", None))

        if dataset_type == "SoccerNetClips" or dataset_type == "SoccerNetGames":
            if dataset_split is None:
                dataset = SoccerNetGameClips(
                    path=source_path,
                    features=cfg.features,
                    version=cfg.version,
                    framerate=cfg.framerate,
                    window_size=cfg.window_size,
                )
            else:
                dataset = SoccerNetClips(
                    path=source_path,
                    features=cfg.features,
                    split=dataset_split,
                    version=cfg.version,
                    framerate=cfg.framerate,
                    window_size=cfg.window_size,
                    train=True if dataset_type == "SoccerNetClips" else False,
                )
        elif dataset_type == "SoccerNetClipsCALF" or dataset_type == "SoccerNetClipsTestingCALF":
            if dataset_split is None:
                dataset = SoccerNetGameClipsChunks(
                    path=source_path,
                    features=cfg.features,
                    framerate=cfg.framerate,
                    chunk_size=cfg.chunk_size,
                    receptive_field=cfg.receptive_field,
                )
            else:
                dataset = SoccerNetClipsChunks(
                    path=source_path,
                    features=cfg.features,
                    split=dataset_split,
                    framerate=cfg.framerate,
                    chunk_size=cfg.chunk_size,
                    receptive_field=cfg.receptive_field,
                    chunks_per_epoch=cfg.chunks_per_epoch,
                    gpu=gpu,
                    train=True if dataset_type == "SoccerNetClipsCALF" else False,
                )
        elif dataset_type == "FeatureClipsfromJSON":
            dataset = FeatureClipsfromJSON(
                path=annotation_path,
                features_dir=source_path,
                classes=cfg.classes,
                framerate=cfg.framerate,
                window_size=cfg.window_size,
            )
        elif dataset_type == "FeatureVideosfromJSON":
            dataset = FeatureClipsfromJSON(
                path=annotation_path,
                features_dir=source_path,
                classes=cfg.classes,
                framerate=cfg.framerate,
                window_size=cfg.window_size,
                train=False,
            )
            # dataset = FeatureVideosfromJSON(path=annotation_path,
            #     framerate=cfg.framerate,
            #     window_size=cfg.window_size)
        elif dataset_type == "FeatureClipChunksfromJson":
            dataset = FeatureClipChunksfromJson(
                path=annotation_path,
                features_dir=source_path,
                classes=cfg.classes,
                framerate=cfg.framerate,
                chunk_size=cfg.chunk_size,
                receptive_field=cfg.receptive_field,
                chunks_per_epoch=cfg.chunks_per_epoch,
                gpu=gpu,
            )
        elif dataset_type == "FeatureVideosChunksfromJson":
            dataset = FeatureClipChunksfromJson(
                path=annotation_path,
                features_dir=source_path,
                classes=cfg.classes,
                framerate=cfg.framerate,
                chunk_size=cfg.chunk_size,
                receptive_field=cfg.receptive_field,
                chunks_per_epoch=cfg.chunks_per_epoch,
                gpu=gpu,
                train=False,
            )
        elif dataset_type == "VideoGameWithOpencv":
            dataset_len = self.data_cfg.epoch_num_frames // self.data_cfg.clip_len
            dataset = ActionSpotDataset(
                default_args["classes"],
                annotation_path,
                source_path,
                self.data_cfg.modality,
                self.data_cfg.clip_len,
                self.data_cfg.input_fps,
                self.data_cfg.extract_fps,
                dataset_len if default_args["train"] else dataset_len // 4,
                is_eval=not default_args["train"],
                crop_dim=self.data_cfg.crop_dim,
                dilate_len=self.data_cfg.dilate_len,
                mixup=self.data_cfg.mixup,
                IMAGENET_MEAN=self.data_cfg.imagenet_mean,
                IMAGENET_STD=self.data_cfg.imagenet_std,
                TARGET_HEIGHT=self.data_cfg.target_height,
                TARGET_WIDTH=self.data_cfg.target_width,
                preserve_aspect_ratio=self.data_cfg.preserve_aspect_ratio,
            )
        elif dataset_type == "VideoGameWithOpencvVideo":
            dataset = ActionSpotVideoDataset(
                default_args["classes"],
                annotation_path,
                source_path,
                self.data_cfg.modality,
                self.data_cfg.clip_len,
                self.data_cfg.input_fps,
                self.data_cfg.extract_fps,
                crop_dim=self.data_cfg.crop_dim,
                overlap_len=getattr(cfg, "overlap_len", self.data_cfg.clip_len // 2),
                IMAGENET_MEAN=self.data_cfg.imagenet_mean,
                IMAGENET_STD=self.data_cfg.imagenet_std,
                TARGET_HEIGHT=self.data_cfg.target_height,
                TARGET_WIDTH=self.data_cfg.target_width,
                preserve_aspect_ratio=self.data_cfg.preserve_aspect_ratio,
            )
        elif dataset_type == "VideoGameWithDali":
            if not DALI_AVAILABLE:
                raise ImportError(
                    "NVIDIA DALI is required. "
                    "Install it or use another dataset type."
                )
            default_args = default_args or {}
            acc_grad_iter = default_args.get(
                "acc_grad_iter",
                getattr(getattr(self.config.TRAIN, "execution", None), "acc_grad_iter", 1),
            )
            num_epochs = default_args.get("num_epochs", getattr(self.config.TRAIN, "epochs", 1))
            repartitions = default_args.get(
                "repartitions",
                getattr(getattr(self.config.TRAIN, "execution", None), "repartitions", ([], [])),
            )
            classes = default_args.get("classes", get_data_classes(self.config))
            is_train = default_args.get("train", self.split == "train")
            loader_batch_size = cfg.dataloader.batch_size // acc_grad_iter
            dataset_len = self.data_cfg.epoch_num_frames // self.data_cfg.clip_len
            dataset = DaliDataSet(
                epochs=num_epochs,
                batch_size=loader_batch_size,
                output_map=cfg.output_map,
                devices=(
                    repartitions[0]
                    if is_train
                    else repartitions[1]
                ),
                #devices=list(range(gpu)),
                classes=classes,
                label_file=annotation_path,
                modality=self.data_cfg.modality,
                clip_len=self.data_cfg.clip_len,
                dataset_len=dataset_len if is_train else dataset_len // 4,
                video_dir=source_path,
                input_fps=self.data_cfg.input_fps,
                extract_fps=self.data_cfg.extract_fps,
                IMAGENET_MEAN=self.data_cfg.imagenet_mean,
                IMAGENET_STD=self.data_cfg.imagenet_std,
                TARGET_HEIGHT=self.data_cfg.target_height,
                TARGET_WIDTH=self.data_cfg.target_width,
                is_eval=not is_train,
                crop_dim=self.data_cfg.crop_dim,
                dilate_len=self.data_cfg.dilate_len,
                mixup=self.data_cfg.mixup,
            )
        elif dataset_type == "VideoGameWithDaliVideo":
            if not DALI_AVAILABLE:
                raise ImportError(
                    "NVIDIA DALI is required. "
                    "Install it or use another dataset type."
                )
            default_args = default_args or {}
            repartitions = default_args.get(
                "repartitions",
                getattr(getattr(self.config.TRAIN, "execution", None), "repartitions", ([], [])),
            )
            classes = default_args.get("classes", get_data_classes(self.config))
            dataset = DaliDataSetVideo(
                batch_size=cfg.dataloader.batch_size,
                output_map=cfg.output_map,
                #devices=list(range(gpu)),
                devices=repartitions[1],
                classes=classes,
                label_file=annotation_path,
                modality=self.data_cfg.modality,
                clip_len=self.data_cfg.clip_len,
                video_dir=source_path,
                input_fps=self.data_cfg.input_fps,
                extract_fps=self.data_cfg.extract_fps,
                IMAGENET_MEAN=self.data_cfg.imagenet_mean,
                IMAGENET_STD=self.data_cfg.imagenet_std,
                TARGET_HEIGHT=self.data_cfg.target_height,
                TARGET_WIDTH=self.data_cfg.target_width,
                overlap_len=cfg.overlap_len,
                crop_dim=self.data_cfg.crop_dim,
            )
        elif dataset_type == "H5OSLJsonSpotting":
            from opensportslib.models.base.rule_based import H5OSLJsonSpottingDataset
            dataset = H5OSLJsonSpottingDataset(
                annotation_path=annotation_path,
                source_path=source_path,
            )
        elif dataset_type == "TrackingActionSpot":
            # Graph topology params live on the encoder component (canonical),
            # spatial normalization on the input transform block.
            encoder_params = get_component_params_by_kind(self.config, "encoder")
            normalize = get_data_transform(self.config).get("normalize", True)
            dataset_len = self.data_cfg.epoch_num_frames // self.data_cfg.clip_len
            dataset = TrackingActionSpotDataset(
                default_args["classes"],
                annotation_path,
                source_path,
                self.data_cfg.clip_len,
                self.data_cfg.extract_fps,
                dataset_len if default_args["train"] else dataset_len // 4,
                edge_type=encoder_params.get("edge_type"),
                k=encoder_params.get("k"),
                r=encoder_params.get("radius"),
                is_eval=not default_args["train"],
                dilate_len=self.data_cfg.dilate_len,
                augmentations=_build_tracking_spotting_augmentations(self.config),
                normalize=normalize,
            )
        elif dataset_type == "TrackingActionSpotVideo":
            encoder_params = get_component_params_by_kind(self.config, "encoder")
            normalize = get_data_transform(self.config).get("normalize", True)
            dataset = TrackingActionSpotVideoDataset(
                default_args["classes"],
                annotation_path,
                source_path,
                self.data_cfg.clip_len,
                self.data_cfg.extract_fps,
                edge_type=encoder_params.get("edge_type"),
                k=encoder_params.get("k"),
                r=encoder_params.get("radius"),
                overlap_len=getattr(cfg, "overlap_len", self.data_cfg.clip_len // 2),
                normalize=normalize,
            )
        else:
            dataset = None
        return dataset


    def building_dataloader(self, dataset, cfg, gpu, dali):
        """Build a dataloader from config dict.

        Args:
            cfg (dict): Config dict. It should at least contain the key "type".
            default_args (dict | None, optional): Default initialization arguments.
                Default: None.

        Returns:
            Dataloader: The constructed dataloader.
        """
        def worker_init_fn(id):
            random.seed(id + 100 * 100)
        if dali:
            return dataset
        if getattr(dataset, "rule_based_no_dataloader", False):
            return None
        dataloader = torch.utils.data.DataLoader(
            dataset,
            batch_size=cfg.batch_size,
            shuffle=cfg.shuffle,
            num_workers=cfg.num_workers if gpu >= 0 else 0,
            pin_memory=cfg.pin_memory if gpu >= 0 else False,
            prefetch_factor=(
                getattr(cfg, "prefetch_factor", None)
            ),
            worker_init_fn=worker_init_fn,
            # None for RGB datasets (default collate); the tracking datasets
            # expose a graph collate that flattens per-frame graphs into a
            # single PyG Batch (see core.utils.data).
            collate_fn=getattr(dataset, "collate_fn", None),
        )
        return dataloader

building_dataloader(dataset, cfg, gpu, dali)

Build a dataloader from config dict.

Parameters:

Name	Type	Description	Default
cfg	dict	Config dict. It should at least contain the key "type".	required
default_args	dict | None	Default initialization arguments. Default: None.	required

Returns:

Name	Type	Description
Dataloader		The constructed dataloader.

Source code in opensportslib/datasets/localization_dataset.py

def building_dataloader(self, dataset, cfg, gpu, dali):
    """Build a dataloader from config dict.

    Args:
        cfg (dict): Config dict. It should at least contain the key "type".
        default_args (dict | None, optional): Default initialization arguments.
            Default: None.

    Returns:
        Dataloader: The constructed dataloader.
    """
    def worker_init_fn(id):
        random.seed(id + 100 * 100)
    if dali:
        return dataset
    if getattr(dataset, "rule_based_no_dataloader", False):
        return None
    dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=cfg.batch_size,
        shuffle=cfg.shuffle,
        num_workers=cfg.num_workers if gpu >= 0 else 0,
        pin_memory=cfg.pin_memory if gpu >= 0 else False,
        prefetch_factor=(
            getattr(cfg, "prefetch_factor", None)
        ),
        worker_init_fn=worker_init_fn,
        # None for RGB datasets (default collate); the tracking datasets
        # expose a graph collate that flattens per-frame graphs into a
        # single PyG Batch (see core.utils.data).
        collate_fn=getattr(dataset, "collate_fn", None),
    )
    return dataloader

FrameReader

Class used to read a video and create a clip of frames by applying some transformations.

Parameters:

Name	Type	Description	Default
modality	string	Modality of the frames.	required
crop_transform		Transformations that apply to frame for cropping.	required
img_transform		Transformations that apply to frame like GaussianBlur and Normalization.	required
same_transform	bool	Whether to apply same trasnforms to every frame of a same clip.	required
sample_fps	int	Fps at which we want to extract frames from the video.	2

Source code in opensportslib/datasets/localization_dataset.py

class FrameReader:
    """Class used to read a video and create a clip of frames by applying some transformations.

    Args:
        modality (string): Modality of the frames.
        crop_transform : Transformations that apply to frame for cropping.
        img_transform : Transformations that apply to frame like GaussianBlur and Normalization.
        same_transform (bool): Whether to apply same trasnforms to every frame of a same clip.
        sample_fps (int): Fps at which we want to extract frames from the video.
    """

    def __init__(
        self,
        modality,
        crop_transform,
        img_transform,
        same_transform,
        sample_fps=2,
        TARGET_HEIGHT=224,
        TARGET_WIDTH=398,
        preserve_aspect_ratio=False,
    ):
        self._is_flow = modality == "flow"
        self._crop_transform = crop_transform
        self._img_transform = img_transform
        self._same_transform = same_transform
        self._sample_fps = sample_fps
        self.TARGET_HEIGHT = TARGET_HEIGHT
        self.TARGET_WIDTH = TARGET_WIDTH
        self.preserve_aspect_ratio = preserve_aspect_ratio

    def _resize_frame_ocv(self, frame):
        """Resize an OpenCV frame using the configured spatial policy.

        With aspect preservation enabled, ``TARGET_HEIGHT`` fixes the output
        height and the width is derived from the source aspect ratio. If no
        target height is configured, ``TARGET_WIDTH`` fixes the output width.
        """
        import cv2

        height, width = frame.shape[:2]
        if self.preserve_aspect_ratio:
            if self.TARGET_HEIGHT is not None and self.TARGET_HEIGHT > 0:
                scale = self.TARGET_HEIGHT / height
            elif self.TARGET_WIDTH is not None and self.TARGET_WIDTH > 0:
                scale = self.TARGET_WIDTH / width
            else:
                return frame

            output_height = max(1, round(height * scale))
            output_width = max(1, round(width * scale))
        else:
            output_height = self.TARGET_HEIGHT
            output_width = self.TARGET_WIDTH

        if output_height is None or output_width is None:
            return frame
        if (height, width) == (output_height, output_width):
            return frame
        return cv2.resize(frame, (output_width, output_height))

    def adapt_frame_ocv(self, frame):
        """Apply some modifications to the frame to have the expected shape and format.

        Args:
            frame (np.array).

        Returns:
            img (torch.tensor).
        """
        import cv2
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        img = torch.from_numpy(frame).float() / 255
        img = img.permute(2, 0, 1)
        if self._is_flow:
            img = img[1:, :, :]  # GB channels contain data
        return img

    def load_frames_ocv(
        self,
        video_name,
        start,
        end,
        pad=False,
        source_start_frame=None,
        source_end_frame=None,
    ):
        """Load frames from a video to create a clip of frames.

        Args:
            video_name (string): The path of the video. This is the full path of the file if we infer a single video
            or the partial path if we infer from json file.
            start (int): Start frame at which we load the clip.
            end (int): End frame at which we finish the clip.
            pad (bool): Whether to apply padding to the clip or not.
            source_start_frame (int | None): First physical frame of a logical
                interval. Both source bounds must be provided together.
            source_end_frame (int | None): Exclusive physical end frame of a
                logical interval.
        """
        import cv2
        import torch.nn as nn

        # if self.infer:
        #     video_path = video_name
        # else:
        #     video_path = os.path.join(self._video_dir, video_name + self.extension)
        vc = cv2.VideoCapture(video_name)
        fps = vc.get(cv2.CAP_PROP_FPS)
        total_frames = int(vc.get(cv2.CAP_PROP_FRAME_COUNT))

        oh = self.TARGET_HEIGHT
        ow = self.TARGET_WIDTH

        frames = []
        rand_crop_state = None
        rand_state_backup = None
        ret = []
        requested_frames = end - start
        if requested_frames <= 0:
            vc.release()
            raise ValueError(f"Invalid requested frame range [{start}, {end})")

        stride_extract = get_stride(fps, self._sample_fps)
        bounded = source_start_frame is not None or source_end_frame is not None
        if bounded and (
            source_start_frame is None or source_end_frame is None
        ):
            vc.release()
            raise ValueError(
                "Both source_start_frame and source_end_frame are required."
            )

        n_pad_start = 0
        n_pad_end = 0
        if bounded:
            source_start_frame = int(source_start_frame)
            source_end_frame = int(source_end_frame)
            if not 0 <= source_start_frame < source_end_frame <= total_frames:
                vc.release()
                raise ValueError(
                    f"Invalid source interval [{source_start_frame}, "
                    f"{source_end_frame}) for {video_name} with "
                    f"{total_frames} frames."
                )
            logical_num_frames = int(
                math.ceil(
                    (source_end_frame - source_start_frame) / stride_extract
                )
            )
            n_pad_start = max(0, -start)
            logical_start = max(0, start)
            logical_end = min(end, logical_num_frames)
            frames_to_read = max(0, logical_end - logical_start)
            target_start_frame = (
                source_start_frame + logical_start * stride_extract
            )
            physical_stop_frame = source_end_frame
        else:
            frames_to_read = requested_frames
            target_start_frame = start * stride_extract
            physical_stop_frame = total_frames

        vc.set(cv2.CAP_PROP_POS_FRAMES, target_start_frame)
        out_frame_num = 0
        i = 0
        while out_frame_num < frames_to_read:
            if target_start_frame + i >= physical_stop_frame:
                break
            # The source MP4 is ~30 fps, but spotting uses ~5 fps. Avoid
            # retrieving and converting the five frames skipped between
            # samples; grab still advances the video decoder correctly.
            if i % stride_extract == 0:
                ret, frame = vc.read()
            else:
                ret = vc.grab()
            if ret:
                if i % stride_extract == 0:
                    frame = self._resize_frame_ocv(frame)
                    img = self.adapt_frame_ocv(frame)
                    if self._crop_transform:
                        if self._same_transform:
                            if rand_crop_state is None:
                                rand_crop_state = random.getstate()
                            else:
                                rand_state_backup = random.getstate()
                                random.setstate(rand_crop_state)

                        img = self._crop_transform(img)

                        if rand_state_backup is not None:
                            # Make sure that rand state still advances
                            random.setstate(rand_state_backup)
                            rand_state_backup = None
                    if not self._same_transform:
                        img = self._img_transform(img)
                    frames.append(img)
                    out_frame_num += 1
                i += 1
            else:
                break
        n_pad_end = requested_frames - n_pad_start - out_frame_num
        vc.release()
        if not frames:
            raise ValueError(
                f"No frames read for logical range [{start}, {end}) from "
                f"{video_name}."
            )
        # In the multicrop case, the shape is (B, T, C, H, W)
        frames = torch.stack(frames, dim=int(len(frames[0].shape) == 4))
        if self._same_transform:
            frames = self._img_transform(frames)

        # Always pad start, but only pad end if requested
        if n_pad_start > 0 or (pad and n_pad_end > 0):
            frames = nn.functional.pad(
                frames, (0, 0, 0, 0, 0, 0, n_pad_start, n_pad_end if pad else 0)
            )
        return frames

adapt_frame_ocv(frame)

Apply some modifications to the frame to have the expected shape and format.

Returns:

Type	Description
	img (torch.tensor).

Source code in opensportslib/datasets/localization_dataset.py

def adapt_frame_ocv(self, frame):
    """Apply some modifications to the frame to have the expected shape and format.

    Args:
        frame (np.array).

    Returns:
        img (torch.tensor).
    """
    import cv2
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    img = torch.from_numpy(frame).float() / 255
    img = img.permute(2, 0, 1)
    if self._is_flow:
        img = img[1:, :, :]  # GB channels contain data
    return img

load_frames_ocv(video_name, start, end, pad=False, source_start_frame=None, source_end_frame=None)

Load frames from a video to create a clip of frames.

Parameters:

Name	Type	Description	Default
video_name	string	The path of the video. This is the full path of the file if we infer a single video	required
start	int	Start frame at which we load the clip.	required
end	int	End frame at which we finish the clip.	required
pad	bool	Whether to apply padding to the clip or not.	False
source_start_frame	int | None	First physical frame of a logical interval. Both source bounds must be provided together.	None
source_end_frame	int | None	Exclusive physical end frame of a logical interval.	None

Source code in opensportslib/datasets/localization_dataset.py

def load_frames_ocv(
    self,
    video_name,
    start,
    end,
    pad=False,
    source_start_frame=None,
    source_end_frame=None,
):
    """Load frames from a video to create a clip of frames.

    Args:
        video_name (string): The path of the video. This is the full path of the file if we infer a single video
        or the partial path if we infer from json file.
        start (int): Start frame at which we load the clip.
        end (int): End frame at which we finish the clip.
        pad (bool): Whether to apply padding to the clip or not.
        source_start_frame (int | None): First physical frame of a logical
            interval. Both source bounds must be provided together.
        source_end_frame (int | None): Exclusive physical end frame of a
            logical interval.
    """
    import cv2
    import torch.nn as nn

    # if self.infer:
    #     video_path = video_name
    # else:
    #     video_path = os.path.join(self._video_dir, video_name + self.extension)
    vc = cv2.VideoCapture(video_name)
    fps = vc.get(cv2.CAP_PROP_FPS)
    total_frames = int(vc.get(cv2.CAP_PROP_FRAME_COUNT))

    oh = self.TARGET_HEIGHT
    ow = self.TARGET_WIDTH

    frames = []
    rand_crop_state = None
    rand_state_backup = None
    ret = []
    requested_frames = end - start
    if requested_frames <= 0:
        vc.release()
        raise ValueError(f"Invalid requested frame range [{start}, {end})")

    stride_extract = get_stride(fps, self._sample_fps)
    bounded = source_start_frame is not None or source_end_frame is not None
    if bounded and (
        source_start_frame is None or source_end_frame is None
    ):
        vc.release()
        raise ValueError(
            "Both source_start_frame and source_end_frame are required."
        )

    n_pad_start = 0
    n_pad_end = 0
    if bounded:
        source_start_frame = int(source_start_frame)
        source_end_frame = int(source_end_frame)
        if not 0 <= source_start_frame < source_end_frame <= total_frames:
            vc.release()
            raise ValueError(
                f"Invalid source interval [{source_start_frame}, "
                f"{source_end_frame}) for {video_name} with "
                f"{total_frames} frames."
            )
        logical_num_frames = int(
            math.ceil(
                (source_end_frame - source_start_frame) / stride_extract
            )
        )
        n_pad_start = max(0, -start)
        logical_start = max(0, start)
        logical_end = min(end, logical_num_frames)
        frames_to_read = max(0, logical_end - logical_start)
        target_start_frame = (
            source_start_frame + logical_start * stride_extract
        )
        physical_stop_frame = source_end_frame
    else:
        frames_to_read = requested_frames
        target_start_frame = start * stride_extract
        physical_stop_frame = total_frames

    vc.set(cv2.CAP_PROP_POS_FRAMES, target_start_frame)
    out_frame_num = 0
    i = 0
    while out_frame_num < frames_to_read:
        if target_start_frame + i >= physical_stop_frame:
            break
        # The source MP4 is ~30 fps, but spotting uses ~5 fps. Avoid
        # retrieving and converting the five frames skipped between
        # samples; grab still advances the video decoder correctly.
        if i % stride_extract == 0:
            ret, frame = vc.read()
        else:
            ret = vc.grab()
        if ret:
            if i % stride_extract == 0:
                frame = self._resize_frame_ocv(frame)
                img = self.adapt_frame_ocv(frame)
                if self._crop_transform:
                    if self._same_transform:
                        if rand_crop_state is None:
                            rand_crop_state = random.getstate()
                        else:
                            rand_state_backup = random.getstate()
                            random.setstate(rand_crop_state)

                    img = self._crop_transform(img)

                    if rand_state_backup is not None:
                        # Make sure that rand state still advances
                        random.setstate(rand_state_backup)
                        rand_state_backup = None
                if not self._same_transform:
                    img = self._img_transform(img)
                frames.append(img)
                out_frame_num += 1
            i += 1
        else:
            break
    n_pad_end = requested_frames - n_pad_start - out_frame_num
    vc.release()
    if not frames:
        raise ValueError(
            f"No frames read for logical range [{start}, {end}) from "
            f"{video_name}."
        )
    # In the multicrop case, the shape is (B, T, C, H, W)
    frames = torch.stack(frames, dim=int(len(frames[0].shape) == 4))
    if self._same_transform:
        frames = self._img_transform(frames)

    # Always pad start, but only pad end if requested
    if n_pad_start > 0 or (pad and n_pad_end > 0):
        frames = nn.functional.pad(
            frames, (0, 0, 0, 0, 0, 0, n_pad_start, n_pad_end if pad else 0)
        )
    return frames

ActionSpotDataset

Bases: Dataset

Class that overrides Dataset class. This class is to prepare training data using opencv. Training data consists of frames, associated labels and a boolean indicating if the clip of frames contains an event. A training sample can be mixed up with another one if mixup is used or not. In particular, a training sample contains the following informations without mixup: "frame": The frames. "contains_event": True if event occurs within these frames, False otherwise. "label": The labels associated to the frames. and the following informations with mixup: "frame": A combination of the frames of the first video and the second one. "contains_event": True if event occurs within these frames, False otherwise. "label": Rearrangement of the labels of each video. "mix_frame": Frames of the second video. "mix_weight": The weight that have been used for mixing frames and labels.

Parameters:

Name	Type	Description	Default
classes	dict	dict of class names to idx.	required
label_file	list[string] | string	Path to label json files. Can be a single file or a list or a json files.	required
video_dir	list[string] | string	Path to folders where videos are located. Can be a single folder or a list of folders. Must match the number of json files.	required
modality	string	[rgb] Modality of the frame.	required
clip_len	int	Length of a clip of frames.	required
input_fps	int	Fps of the input videos.	required
extract_fps	int	Fps at which we want to extract frames.	required
dataset_len	int	Number of clips.	required
is_eval	bool	Disable random augmentation Default: True.	True
crop_dim	int	The dimension for cropping frames. Default: None.	None
same_transform	bool	Apply the same random augmentation to each frame in a clip. Default: True.	True
dilate_len	int	Dilate ground truth labels. Default: 0.	0
mixup	bool	Whether to mixup clips of two videos or not. Default: False.	False
pad_len	int	Number of frames to pad the start and end of videos. Default: DEFAULT_PAD_LEN.	5
fg_upsample		Sample foreground explicitly. Default: -1.	-1

Source code in opensportslib/datasets/localization_dataset.py

class ActionSpotDataset(Dataset):
    """Class that overrides Dataset class. This class is to prepare training data using opencv.
    Training data consists of frames, associated labels and a boolean indicating if the clip of frames contains an event.
    A training sample can be mixed up with another one if mixup is used or not.
    In particular, a training sample contains the following informations without mixup:
        "frame": The frames.
        "contains_event": True if event occurs within these frames, False otherwise.
        "label": The labels associated to the frames.
    and the following informations with mixup:
        "frame": A combination of the frames of the first video and the second one.
        "contains_event": True if event occurs within these frames, False otherwise.
        "label": Rearrangement of the labels of each video.
        "mix_frame": Frames of the second video.
        "mix_weight": The weight that have been used for mixing frames and labels.

    Args:
        classes (dict): dict of class names to idx.
        label_file (list[string]|string): Path to label json files. Can be a single file or a list or a json files.
        video_dir (list[string]|string): Path to folders where videos are located. Can be a single folder or a list of folders. Must match the number of json files.
        modality (string): [rgb] Modality of the frame.
        clip_len (int): Length of a clip of frames.
        input_fps (int): Fps of the input videos.
        extract_fps (int): Fps at which we want to extract frames.
        dataset_len (int): Number of clips.
        is_eval (bool): Disable random augmentation
            Default: True.
        crop_dim (int): The dimension for cropping frames.
            Default: None.
        same_transform (bool): Apply the same random augmentation to each frame in a clip.
            Default: True.
        dilate_len (int): Dilate ground truth labels.
            Default: 0.
        mixup (bool): Whether to mixup clips of two videos or not.
            Default: False.
        pad_len (int): Number of frames to pad the start and end of videos.
            Default: DEFAULT_PAD_LEN.
        fg_upsample: Sample foreground explicitly.
            Default: -1.
    """

    def __init__(
        self,
        classes,
        label_file,
        video_dir,
        modality,
        clip_len,
        input_fps,
        extract_fps,
        dataset_len,
        is_eval=True,
        crop_dim=None,
        # stride=1,  # Downsample frame rate
        same_transform=True,
        dilate_len=0,
        mixup=False,
        pad_len=5,
        fg_upsample=-1,
        IMAGENET_MEAN=[0.485, 0.456, 0.406],
        IMAGENET_STD=[0.229, 0.224, 0.225],
        TARGET_HEIGHT=224,
        TARGET_WIDTH=398,
        preserve_aspect_ratio=False,
    ):
        import random
        from opensportslib.core.utils.load_annotations import annotationstoe2eformat
        from opensportslib.core.utils.video_processing import _get_deferred_rgb_transform, _get_img_transforms

        self._src_file = label_file
        self._labels, self.task_name = annotationstoe2eformat(
            label_file, video_dir, input_fps, extract_fps, False
        )
        # self._labels = load_json(label_file)
        self._class_dict = _normalize_class_dict(classes)
        self._video_idxs = {x["video"]: i for i, x in enumerate(self._labels)}
        # Sample videos weighted by their length
        num_frames = [v["num_frames"] for v in self._labels]
        self._weights_by_length = np.array(num_frames) / np.sum(num_frames)

        self._clip_len = clip_len
        assert clip_len > 0
        self._stride = 1
        assert self._stride > 0
        self._dataset_len = dataset_len
        assert dataset_len > 0
        self._pad_len = pad_len
        assert pad_len >= 0
        self._is_eval = is_eval

        # Label modifications
        self._dilate_len = dilate_len
        self._fg_upsample = fg_upsample

        # Sample based on foreground labels
        if self._fg_upsample > 0:
            self._flat_labels = []
            for i, x in enumerate(self._labels):
                for event in x["events"]:
                    if event["frame"] < x["num_frames"]:
                        self._flat_labels.append((i, event["frame"]))

        self._mixup = mixup

        self.IMAGENET_MEAN = IMAGENET_MEAN
        self.IMAGENET_STD = IMAGENET_STD
        self.TARGET_HEIGHT = TARGET_HEIGHT
        self.TARGET_WIDTH = TARGET_WIDTH
        # Try to do defer the latter half of the transforms to the GPU
        self._gpu_transform = None
        if not is_eval and same_transform:
            if modality == "rgb":
                print("=> Deferring some RGB transforms to the GPU!")
                self._gpu_transform = _get_deferred_rgb_transform(self.IMAGENET_MEAN, self.IMAGENET_STD)


        crop_transform, img_transform = _get_img_transforms(
            self.IMAGENET_MEAN,
            self.IMAGENET_STD,
            is_eval,
            crop_dim,
            modality,
            same_transform,
            defer_transform=self._gpu_transform is not None,
        )

        self._frame_reader = FrameReader(
            modality,
            crop_transform,
            img_transform,
            same_transform,
            extract_fps,
            self.TARGET_HEIGHT,
            self.TARGET_WIDTH,
            preserve_aspect_ratio,
        )

    def load_frame_gpu(self, batch, device):
        """Load frame to gpu by appliyng some transformations or not.

        Args:
            batch: Batch containing data.
            device: The device on which to load the frames.
        """
        from opensportslib.core.utils.video_processing import _load_frame_deferred
        if self._gpu_transform is None:
            frame = batch["frame"].to(device)
        else:
            frame = _load_frame_deferred(self._gpu_transform, batch, device)
        return frame

    def _sample_uniform(self):
        """Sample video metadata and a base start index uniformly based on video lengths and weights.
        Returns:
            video_meta: metadata of a video.
            base_idx: base start index for the video processing.
        """
        video_meta = random.choices(self._labels, weights=self._weights_by_length)[0]

        video_len = video_meta["num_frames"]
        base_idx = -self._pad_len * self._stride + random.randint(
            0,
            max(0, video_len - 1 + (2 * self._pad_len - self._clip_len) * self._stride),
        )
        return video_meta, base_idx

    def _sample_foreground(self):
        """Samples video metadata and a base start index focusing on foreground labels.
        Returns:
            video_meta: metadata of a video.
            base_idx: base start index for the video processing.
        """
        video_idx, frame_idx = random.choices(self._flat_labels)[0]
        video_meta = self._labels[video_idx]
        video_len = video_meta["num_frames"]

        lower_bound = max(
            -self._pad_len * self._stride, frame_idx - self._clip_len * self._stride + 1
        )
        upper_bound = min(
            video_len - 1 + (self._pad_len - self._clip_len) * self._stride, frame_idx
        )

        base_idx = (
            random.randint(lower_bound, upper_bound)
            if upper_bound > lower_bound
            else lower_bound
        )

        assert base_idx <= frame_idx
        assert base_idx + self._clip_len > frame_idx
        return video_meta, base_idx

    def _get_one(self):
        """Get a training sample for one video."""
        if self._fg_upsample > 0 and random.random() >= self._fg_upsample:
            video_meta, base_idx = self._sample_foreground()
        else:
            video_meta, base_idx = self._sample_uniform()

        labels = np.zeros(self._clip_len, np.int64)
        for event in video_meta["events"]:
            event_frame = event["frame"]

            # Index of event in label array
            label_idx = (event_frame - base_idx) // self._stride
            if (
                label_idx >= -self._dilate_len
                and label_idx < self._clip_len + self._dilate_len
            ):
                label = self._class_dict[event["label"]]
                for i in range(
                    max(0, label_idx - self._dilate_len),
                    min(self._clip_len, label_idx + self._dilate_len + 1),
                ):
                    labels[i] = label
        frames = self._frame_reader.load_frames_ocv(
            video_meta["video"],
            base_idx,
            base_idx + self._clip_len * self._stride,
            pad=True,
            source_start_frame=video_meta.get("source_start_frame"),
            source_end_frame=video_meta.get("source_end_frame"),
        )

        return {
            "frame": frames,
            "contains_event": int(np.sum(labels) > 0),
            "label": labels,
        }

    def __getitem__(self, unused):
        """Get a training sample based on one video without mixup, two otherwise."""
        ret = self._get_one()

        if self._mixup:
            mix = self._get_one()  # Sample another clip
            l = random.betavariate(0.2, 0.2)
            label_dist = np.zeros((self._clip_len, len(self._class_dict) + 1))
            label_dist[range(self._clip_len), ret["label"]] = l
            label_dist[range(self._clip_len), mix["label"]] += 1.0 - l

            if self._gpu_transform is None:
                ret["frame"] = l * ret["frame"] + (1.0 - l) * mix["frame"]
            else:
                ret["mix_frame"] = mix["frame"]
                ret["mix_weight"] = l

            ret["contains_event"] = max(ret["contains_event"], mix["contains_event"])
            ret["label"] = label_dist

        return ret

    def __len__(self):
        return self._dataset_len

    def print_info(self):
        from opensportslib.core.utils.config import _print_info_helper
        _print_info_helper(self._src_file, self._labels)

__getitem__(unused)

Get a training sample based on one video without mixup, two otherwise.

Source code in opensportslib/datasets/localization_dataset.py

def __getitem__(self, unused):
    """Get a training sample based on one video without mixup, two otherwise."""
    ret = self._get_one()

    if self._mixup:
        mix = self._get_one()  # Sample another clip
        l = random.betavariate(0.2, 0.2)
        label_dist = np.zeros((self._clip_len, len(self._class_dict) + 1))
        label_dist[range(self._clip_len), ret["label"]] = l
        label_dist[range(self._clip_len), mix["label"]] += 1.0 - l

        if self._gpu_transform is None:
            ret["frame"] = l * ret["frame"] + (1.0 - l) * mix["frame"]
        else:
            ret["mix_frame"] = mix["frame"]
            ret["mix_weight"] = l

        ret["contains_event"] = max(ret["contains_event"], mix["contains_event"])
        ret["label"] = label_dist

    return ret

load_frame_gpu(batch, device)

Load frame to gpu by appliyng some transformations or not.

Parameters:

Name	Type	Description	Default
batch		Batch containing data.	required
device		The device on which to load the frames.	required

Source code in opensportslib/datasets/localization_dataset.py

def load_frame_gpu(self, batch, device):
    """Load frame to gpu by appliyng some transformations or not.

    Args:
        batch: Batch containing data.
        device: The device on which to load the frames.
    """
    from opensportslib.core.utils.video_processing import _load_frame_deferred
    if self._gpu_transform is None:
        frame = batch["frame"].to(device)
    else:
        frame = _load_frame_deferred(self._gpu_transform, batch, device)
    return frame

ActionSpotVideoDataset

Bases: Dataset, DatasetVideoSharedMethods

Class that overrides Dataset class. This class is to prepare testing data. Testing data consists of frames, the name of the video and index of the first frame in the video. This class can process as input a json file containing metadatas of video or just a video. Args: classes (dict): dict of class names to idx. label_file (string): Can be path to label json or path of a video. video_dir (string): path to folder where videos are located. modality (string): [rgb] Modality of the frame. clip_len (int): Length of a clip of frames. input_fps (int): Fps of the input videos. extract_fps (int): Fps at which we want to extract frames. overlap_len (int): The number of overlapping frames between consecutive clips. crop_dim (int): The dimension for cropping frames. Default: None. flip (bool): Whether to flip or not the frames. Default: False. multi_crop (bool): Whether multi croping or not Default: False. skip_partial_end (bool): Whether to skip a partial number of clips at the end. Default: True. pad_len (int): Number of frames to pad the start and end of videos. Default: DEFAULT_PAD_LEN.

Source code in opensportslib/datasets/localization_dataset.py

class ActionSpotVideoDataset(Dataset, DatasetVideoSharedMethods):
    """Class that overrides Dataset class. This class is to prepare testing data.
    Testing data consists of frames, the name of the video and index of the first frame in the video.
    This class can process as input a json file containing metadatas of video or just a video.
    Args:
        classes (dict): dict of class names to idx.
        label_file (string): Can be path to label json or path of a video.
        video_dir (string): path to folder where videos are located.
        modality (string): [rgb] Modality of the frame.
        clip_len (int): Length of a clip of frames.
        input_fps (int): Fps of the input videos.
        extract_fps (int): Fps at which we want to extract frames.
        overlap_len (int): The number of overlapping frames between consecutive clips.
        crop_dim (int): The dimension for cropping frames.
            Default: None.
        flip (bool): Whether to flip or not the frames.
            Default: False.
        multi_crop (bool): Whether multi croping or not
            Default: False.
        skip_partial_end (bool): Whether to skip a partial number of clips at the end.
            Default: True.
        pad_len (int): Number of frames to pad the start and end of videos.
            Default: DEFAULT_PAD_LEN.
    """

    def __init__(
        self,
        classes,
        label_file,
        video_dir,
        modality,
        clip_len,
        input_fps,
        extract_fps,
        overlap_len=0,
        crop_dim=None,
        pad_len=5,
        flip=False,
        multi_crop=False,
        skip_partial_end=True,
        IMAGENET_MEAN=[0.485, 0.456, 0.406],
        IMAGENET_STD=[0.229, 0.224, 0.225],
        TARGET_HEIGHT=224,
        TARGET_WIDTH=398,
        preserve_aspect_ratio=False,
    ):
        from opensportslib.core.utils.load_annotations import annotationstoe2eformat, construct_labels
        from opensportslib.core.utils.video_processing import _get_img_transforms

        self._src_file = label_file
        if label_file.endswith(".json"):
            self._labels, self.task_name = annotationstoe2eformat(
                label_file, video_dir, input_fps, extract_fps, False
            )
            # self._labels = load_json(label_file)
        else:
            self._labels, _ = construct_labels(label_file, extract_fps)
        # self._labels = load_json(label_file)
        self._class_dict = _normalize_class_dict(classes)
        self._video_idxs = {x["path"]: i for i, x in enumerate(self._labels)}
        self._clip_len = clip_len
        stride = 1
        self._stride = stride
        self.IMAGENET_MEAN = IMAGENET_MEAN
        self.IMAGENET_STD = IMAGENET_STD
        self.TARGET_HEIGHT = TARGET_HEIGHT
        self.TARGET_WIDTH = TARGET_WIDTH
        crop_transform, img_transform = _get_img_transforms(
            self.IMAGENET_MEAN,
            self.IMAGENET_STD,
            is_eval=True,
            crop_dim=crop_dim,
            modality=modality,
            same_transform=True,
            multi_crop=multi_crop,
        )

        # No need to enforce same_transform since the transforms are
        # deterministic
        self._frame_reader = FrameReader(
            modality,
            crop_transform,
            img_transform,
            False,
            extract_fps,
            self.TARGET_HEIGHT,
            self.TARGET_WIDTH,
            preserve_aspect_ratio,
        )

        self._flip = flip
        self._multi_crop = multi_crop

        self._clips = []
        for l in self._labels:
            has_clip = False
            for i in range(
                -pad_len * self._stride,
                max(
                    0, l["num_frames"] - (overlap_len * stride) * int(skip_partial_end)
                ),  # Need to ensure that all clips have at least one frame
                (clip_len - overlap_len) * self._stride,
            ):
                has_clip = True
                self._clips.append(
                    (
                        l["path"],
                        l["video"],
                        i,
                        l.get("source_start_frame"),
                        l.get("source_end_frame"),
                    )
                )
            assert has_clip, l

    def __len__(self):
        return len(self._clips)

    def __getitem__(self, idx):
        """Get a dict of metadata containing the name of the video, the index of the first frame and the clip of frame.

        Args:
            idx (int): The index of the clip in the list of clips.

        Returns:
            dict :{"video","start","frame"}.
        """
        (
            video_path,
            video_name,
            start,
            source_start_frame,
            source_end_frame,
        ) = self._clips[idx]
        frames = self._frame_reader.load_frames_ocv(
            video_name,
            start,
            start + self._clip_len * self._stride,
            pad=True,
            source_start_frame=source_start_frame,
            source_end_frame=source_end_frame,
        )
        # ,stride=self._stride)

        if self._flip:
            frames = torch.stack((frames, frames.flip(-1)), dim=0)

        return {"video": video_path, "start": start // self._stride, "frame": frames}

__getitem__(idx)

Get a dict of metadata containing the name of the video, the index of the first frame and the clip of frame.

Parameters:

Name	Type	Description	Default
idx	int	The index of the clip in the list of clips.	required

Returns:

Name	Type	Description
dict		{"video","start","frame"}.

Source code in opensportslib/datasets/localization_dataset.py

def __getitem__(self, idx):
    """Get a dict of metadata containing the name of the video, the index of the first frame and the clip of frame.

    Args:
        idx (int): The index of the clip in the list of clips.

    Returns:
        dict :{"video","start","frame"}.
    """
    (
        video_path,
        video_name,
        start,
        source_start_frame,
        source_end_frame,
    ) = self._clips[idx]
    frames = self._frame_reader.load_frames_ocv(
        video_name,
        start,
        start + self._clip_len * self._stride,
        pad=True,
        source_start_frame=source_start_frame,
        source_end_frame=source_end_frame,
    )
    # ,stride=self._stride)

    if self._flip:
        frames = torch.stack((frames, frames.flip(-1)), dim=0)

    return {"video": video_path, "start": start // self._stride, "frame": frames}

TrackingActionSpotDataset

Bases: Dataset, _TrackingGraphWindowMixin

Training dataset for tracking-based action spotting.

Mirrors ActionSpotDataset (random clip_len-frame windows, dense dilated per-frame labels, uniform/foreground-weighted sampling) but slices per-frame graphs out of a cached whole-match tracking parquet (see datasets.utils.tracking.load_or_build_game_cache) instead of decoding video frames. Mixup is not supported (mixing raw node features across two differently-structured graphs is not well-defined).

Parameters:

Name	Type	Description	Default
classes	dict	dict of class names to idx.	required
label_file	str | list[str]	Path to annotation JSON file(s).	required
video_dir	str | list[str]	Root dir(s) containing the tracking parquets referenced by the annotations.	required
clip_len	int	Length of a clip, in decimated (extract_fps) frames.	required
extract_fps	float	Target fps to decimate the native stream to.	required
dataset_len	int	Number of clips per epoch.	required
edge_type	str	Graph edge-building strategy (see build_edge_index).	'positional'
k	int	neighbour count for knn-style edge strategies.	8
r	float	distance threshold (metres) for distance-style strategies.	15.0
is_eval	bool	Disable augmentation (mirrors ActionSpotDataset).	True
dilate_len	int	Dilate ground-truth labels.	0
pad_len	int	Frames to pad the start/end of a game.	5
fg_upsample	float	Sample foreground explicitly.	-1
augmentations	list	tracking augmentation transforms.	None
normalize	bool	whether to normalize spatial features.	True

Source code in opensportslib/datasets/localization_dataset.py

class TrackingActionSpotDataset(Dataset, _TrackingGraphWindowMixin):
    """Training dataset for tracking-based action spotting.

    Mirrors ActionSpotDataset (random clip_len-frame windows, dense dilated
    per-frame labels, uniform/foreground-weighted sampling) but slices
    per-frame graphs out of a cached whole-match tracking parquet (see
    datasets.utils.tracking.load_or_build_game_cache) instead of decoding
    video frames. Mixup is not supported (mixing raw node features across two
    differently-structured graphs is not well-defined).

    Args:
        classes (dict): dict of class names to idx.
        label_file (str | list[str]): Path to annotation JSON file(s).
        video_dir (str | list[str]): Root dir(s) containing the tracking
            parquets referenced by the annotations.
        clip_len (int): Length of a clip, in decimated (extract_fps) frames.
        extract_fps (float): Target fps to decimate the native stream to.
        dataset_len (int): Number of clips per epoch.
        edge_type (str): Graph edge-building strategy (see build_edge_index).
        k (int): neighbour count for knn-style edge strategies.
        r (float): distance threshold (metres) for distance-style strategies.
        is_eval (bool): Disable augmentation (mirrors ActionSpotDataset).
        dilate_len (int): Dilate ground-truth labels.
        pad_len (int): Frames to pad the start/end of a game.
        fg_upsample (float): Sample foreground explicitly.
        augmentations (list): tracking augmentation transforms.
        normalize (bool): whether to normalize spatial features.
    """

    collate_fn = staticmethod(tracking_spotting_collate_fn)

    def __init__(
        self,
        classes,
        label_file,
        video_dir,
        clip_len,
        extract_fps,
        dataset_len,
        edge_type="positional",
        k=8,
        r=15.0,
        is_eval=True,
        dilate_len=0,
        pad_len=5,
        fg_upsample=-1,
        augmentations=None,
        normalize=True,
    ):
        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
        )
        self._class_dict = _normalize_class_dict(classes)
        self._video_idxs = {x["video"]: i for i, x in enumerate(self._labels)}

        num_frames = [v["num_frames"] for v in self._labels]
        self._weights_by_length = np.array(num_frames) / np.sum(num_frames)

        self._is_eval = is_eval
        self._stride = 1
        self._dataset_len = dataset_len
        assert dataset_len > 0
        self._pad_len = pad_len
        assert pad_len >= 0

        self._dilate_len = dilate_len
        self._fg_upsample = fg_upsample

        if self._fg_upsample > 0:
            self._flat_labels = []
            for i, x in enumerate(self._labels):
                for event in x["events"]:
                    if event["frame"] < x["num_frames"]:
                        self._flat_labels.append((i, event["frame"]))

        # Like ActionSpotDataset, augmentation is disabled at eval time; the
        # dataset owns that decision (via is_eval) rather than the caller.
        self._init_tracking_window(
            clip_len, edge_type, k, r, normalize,
            None if is_eval else augmentations,
        )

    def _sample_uniform(self):
        """Sample video metadata and a base start index uniformly based on
        video lengths and weights."""
        video_meta = random.choices(self._labels, weights=self._weights_by_length)[0]

        video_len = video_meta["num_frames"]
        base_idx = -self._pad_len * self._stride + random.randint(
            0,
            max(0, video_len - 1 + (2 * self._pad_len - self._clip_len) * self._stride),
        )
        return video_meta, base_idx

    def _sample_foreground(self):
        """Sample video metadata and a base start index focusing on
        foreground labels."""
        video_idx, frame_idx = random.choices(self._flat_labels)[0]
        video_meta = self._labels[video_idx]
        video_len = video_meta["num_frames"]

        lower_bound = max(
            -self._pad_len * self._stride, frame_idx - self._clip_len * self._stride + 1
        )
        upper_bound = min(
            video_len - 1 + (self._pad_len - self._clip_len) * self._stride, frame_idx
        )

        base_idx = (
            random.randint(lower_bound, upper_bound)
            if upper_bound > lower_bound
            else lower_bound
        )
        return video_meta, base_idx

    def _get_one(self):
        """Get a training sample for one game."""
        if self._fg_upsample > 0 and random.random() >= self._fg_upsample:
            video_meta, base_idx = self._sample_foreground()
        else:
            video_meta, base_idx = self._sample_uniform()

        labels = np.zeros(self._clip_len, np.int64)
        for event in video_meta["events"]:
            event_frame = event["frame"]

            label_idx = (event_frame - base_idx) // self._stride
            if (
                label_idx >= -self._dilate_len
                and label_idx < self._clip_len + self._dilate_len
            ):
                label = self._class_dict[event["label"]]
                for i in range(
                    max(0, label_idx - self._dilate_len),
                    min(self._clip_len, label_idx + self._dilate_len + 1),
                ):
                    labels[i] = label

        window, positions = self._slice_window(video_meta, base_idx)
        graphs = self._build_graphs(window, positions)

        return {
            "graphs": graphs,
            "contains_event": int(np.sum(labels) > 0),
            "label": labels,
        }

    def __getitem__(self, unused):
        """Get a training sample based on one game."""
        return self._get_one()

    def __len__(self):
        return self._dataset_len

    def print_info(self):
        num_frames = sum(x["num_frames"] for x in self._labels)
        num_events = sum(len(x["events"]) for x in self._labels)
        print(
            "{} : {} games, {} frames ({} stride), {:0.5f}% non-bg".format(
                self._src_file,
                len(self._labels),
                num_frames,
                self._stride,
                num_events / num_frames * 100 if num_frames else 0.0,
            )
        )

__getitem__(unused)

Get a training sample based on one game.

Source code in opensportslib/datasets/localization_dataset.py

def __getitem__(self, unused):
    """Get a training sample based on one game."""
    return self._get_one()

TrackingActionSpotVideoDataset

Bases: Dataset, _TrackingGraphWindowMixin, DatasetVideoSharedMethods

Sliding-window evaluation dataset for tracking-based action spotting.

Mirrors ActionSpotVideoDataset/DatasetVideoSharedMethods - same (path, video, start) sliding-clip list with overlap_len for full-match coverage, and the same get_labels/videos/labels properties the mAP evaluator relies on - but each clip is a window of per-frame graphs sliced out of a cached whole-match tracking parquet instead of decoded video frames.

Parameters:

Name	Type	Description	Default
classes	dict	dict of class names to idx.	required
label_file	str	Path to annotation JSON file.	required
video_dir	str	Root dir containing the tracking parquets.	required
clip_len	int	Length of a clip, in decimated (extract_fps) frames.	required
extract_fps	float	Target fps to decimate the native stream to.	required
edge_type	str	Graph edge-building strategy (see build_edge_index).	'positional'
k	int	neighbour count for knn-style edge strategies.	8
r	float	distance threshold (metres) for distance-style strategies.	15.0
overlap_len	int	Overlapping frames between consecutive clips.	0
pad_len	int	Frames to pad the start/end of a game.	5
skip_partial_end	bool	Whether to skip a partial clip at the end.	True
normalize	bool	whether to normalize spatial features.	True

Source code in opensportslib/datasets/localization_dataset.py

class TrackingActionSpotVideoDataset(Dataset, _TrackingGraphWindowMixin, DatasetVideoSharedMethods):
    """Sliding-window evaluation dataset for tracking-based action spotting.

    Mirrors ActionSpotVideoDataset/DatasetVideoSharedMethods - same
    (path, video, start) sliding-clip list with overlap_len for full-match
    coverage, and the same get_labels/videos/labels properties the mAP
    evaluator relies on - but each clip is a window of per-frame graphs
    sliced out of a cached whole-match tracking parquet instead of decoded
    video frames.

    Args:
        classes (dict): dict of class names to idx.
        label_file (str): Path to annotation JSON file.
        video_dir (str): Root dir containing the tracking parquets.
        clip_len (int): Length of a clip, in decimated (extract_fps) frames.
        extract_fps (float): Target fps to decimate the native stream to.
        edge_type (str): Graph edge-building strategy (see build_edge_index).
        k (int): neighbour count for knn-style edge strategies.
        r (float): distance threshold (metres) for distance-style strategies.
        overlap_len (int): Overlapping frames between consecutive clips.
        pad_len (int): Frames to pad the start/end of a game.
        skip_partial_end (bool): Whether to skip a partial clip at the end.
        normalize (bool): whether to normalize spatial features.
    """

    collate_fn = staticmethod(tracking_spotting_eval_collate_fn)

    def __init__(
        self,
        classes,
        label_file,
        video_dir,
        clip_len,
        extract_fps,
        edge_type="positional",
        k=8,
        r=15.0,
        overlap_len=0,
        pad_len=5,
        skip_partial_end=True,
        normalize=True,
    ):
        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
        )
        self._class_dict = _normalize_class_dict(classes)
        self._video_idxs = {x["path"]: i for i, x in enumerate(self._labels)}
        self._stride = 1

        # DatasetVideoSharedMethods.augment reads these; tracking data has
        # no notion of flip/multi-crop augmentation at eval time.
        self._flip = False
        self._multi_crop = False

        self._init_tracking_window(clip_len, edge_type, k, r, normalize)

        stride = self._stride
        self._clips = []
        for l in self._labels:
            has_clip = False
            for i in range(
                -pad_len * stride,
                max(
                    0, l["num_frames"] - (overlap_len * stride) * int(skip_partial_end)
                ),  # Need to ensure that all clips have at least one frame
                (clip_len - overlap_len) * stride,
            ):
                has_clip = True
                self._clips.append((l["path"], l["video"], i))
            assert has_clip, l

    def __len__(self):
        return len(self._clips)

    def __getitem__(self, idx):
        """Get a dict of metadata containing the name of the game, the index
        of the first frame, and the clip's per-frame graphs.

        Returns:
            dict: {"video", "start", "graphs"}.
        """
        video_path, video, start = self._clips[idx]
        video_meta = self._labels[self._video_idxs[video_path]]
        window, positions = self._slice_window(video_meta, start)
        graphs = self._build_graphs(window, positions)
        return {"video": video_path, "start": start // self._stride, "graphs": graphs}

    @property
    def frame_times(self):
        """Map each video to the absolute match clock (ms) of its decimated frames.

        Unlike a decoded video, a tracking parquet's frame index is not a
        multiple of the frame period: recording starts mid-broadcast (row 0
        can be ~80 s into the clock) and the clock jumps ~60-90 s across
        half-time. process_frame_predictions uses this to report predictions
        on the same clock the annotations' position_ms uses; without it every
        prediction lands tens of seconds away from its ground truth.
        """
        return {
            v["path"]: v["frame_times"]
            for v in self._labels
            if v.get("frame_times") is not None
        }

frame_times property

Map each video to the absolute match clock (ms) of its decimated frames.

Unlike a decoded video, a tracking parquet's frame index is not a multiple of the frame period: recording starts mid-broadcast (row 0 can be ~80 s into the clock) and the clock jumps ~60-90 s across half-time. process_frame_predictions uses this to report predictions on the same clock the annotations' position_ms uses; without it every prediction lands tens of seconds away from its ground truth.

__getitem__(idx)

Get a dict of metadata containing the name of the game, the index of the first frame, and the clip's per-frame graphs.

Returns:

Name	Type	Description
dict		{"video", "start", "graphs"}.

Source code in opensportslib/datasets/localization_dataset.py

def __getitem__(self, idx):
    """Get a dict of metadata containing the name of the game, the index
    of the first frame, and the clip's per-frame graphs.

    Returns:
        dict: {"video", "start", "graphs"}.
    """
    video_path, video, start = self._clips[idx]
    video_meta = self._labels[self._video_idxs[video_path]]
    window, positions = self._slice_window(video_meta, start)
    graphs = self._build_graphs(window, positions)
    return {"video": video_path, "start": start // self._stride, "graphs": graphs}

FeaturefromJson

Bases: Dataset

Parent class that is used to modularise common operations between the classes that prepare features data. In particular, this class loads the input, creates dictionnaries of classes and has a method to process the annotations.

Parameters:

Name	Type	Description	Default
path	str | List(str)	Path of the input. Can be a json file, a features file or a list of json files (list only for training purposes).	required
features_dir	str | List(str)	Path where the features are located. Must match the number of input paths.	required
classes	path	Path of the file containing the classes. Only used if input is feature file, otherwise list of classes is in the json file.	required
framerate	int	The framerate at which the features have been extracted. Default: 2.	2

Source code in opensportslib/datasets/localization_dataset.py

class FeaturefromJson(Dataset):
    """Parent class that is used to modularise common operations between the classes that prepare features data.
    In particular, this class loads the input, creates dictionnaries of classes and has a method to process the annotations.

    Args:
        path (str|List(str)): Path of the input. Can be a json file, a features file or a list of json files (list only for training purposes).
        features_dir (str|List(str)): Path where the features are located. Must match the number of input paths.
        classes (path): Path of the file containing the classes. Only used if input is feature file, otherwise list of classes is in the json file.
        framerate (int): The framerate at which the features have been extracted.
            Default: 2.
    """

    def __init__(self, path, features_dir, classes, framerate=2):
        self.classes = classes
        self.path = path
        self.framerate = framerate
        self.features_dir = features_dir

        self.is_json = True

        if isinstance(path, list):
            self.data_json = []
            self.classes = []
            for single_path in path:
                assert os.path.isfile(single_path)
                with open(single_path) as f:
                    tmp = json.load(f)
                    self.data_json.append(tmp)
                    for task_name, task_data in tmp["labels"].items():
                        self.classes.append(task_data.get("labels", {}))
            assert all(x == self.classes[0] for x in self.classes) == True
            self.classes = self.classes[0]

        else:
            self.features_dir = [features_dir]

            assert os.path.isfile(path)
            if path.endswith(".json"):
                with open(path) as f:
                    tmp = json.load(f)
                    self.data_json = [tmp]

                for task_name, task_data in tmp["labels"].items():
                    self.classes = task_data.get("labels", {})
            else:
                self.is_json = False
                self.data_json = [
                    {
                        "data": [
                            {
                                "inputs": [{
                                    "path": path
                                }],
                                "events": [],
                            }
                        ]
                    }
                ]
                assert isinstance(self.classes, list) or os.path.isfile(self.classes)

                from opensportslib.core.utils.config import load_text
                if not isinstance(self.classes, list):
                    self.classes = load_text(self.classes)

        self.num_classes = len(self.classes)
        print(self.num_classes)
        self.event_dictionary = {cls: i_cls for i_cls, cls in enumerate(self.classes)}
        self.inverse_event_dictionary = {
            i_cls: cls for i_cls, cls in enumerate(self.classes)
        }
        logging.info("Pre-compute clips")

    def __getitem__(self, index):
        pass

    def __len__(self):
        pass

    def annotation(self, annotation):
        """Process an annotation to derive the frame number, the class index and a boolean.
        Args:
            annotation (dict): It must contains the keys "gameTime" and "label".

        Returns:
            label (int): The index of the class.
            frame (int): The index of the frame.
            cont (bool): Whether to continue in the loop or not.
        """
        # time = annotation["gameTime"]
        event = annotation["label"]

        if "position_ms" in annotation.keys():
            frame = int(self.framerate * (int(annotation["position_ms"]) / 1000))
        else:
            time = annotation["gameTime"]

            minutes = int(time[-5:-3])
            seconds = int(time[-2::])
            frame = self.framerate * (seconds + 60 * minutes)

        cont = False

        if event not in self.classes:
            cont = True
        else:
            label = self.classes.index(event)

        return label, frame, cont

annotation(annotation)

Process an annotation to derive the frame number, the class index and a boolean. Args: annotation (dict): It must contains the keys "gameTime" and "label".

Returns:

Name	Type	Description
label	int	The index of the class.
frame	int	The index of the frame.
cont	bool	Whether to continue in the loop or not.

Source code in opensportslib/datasets/localization_dataset.py

def annotation(self, annotation):
    """Process an annotation to derive the frame number, the class index and a boolean.
    Args:
        annotation (dict): It must contains the keys "gameTime" and "label".

    Returns:
        label (int): The index of the class.
        frame (int): The index of the frame.
        cont (bool): Whether to continue in the loop or not.
    """
    # time = annotation["gameTime"]
    event = annotation["label"]

    if "position_ms" in annotation.keys():
        frame = int(self.framerate * (int(annotation["position_ms"]) / 1000))
    else:
        time = annotation["gameTime"]

        minutes = int(time[-5:-3])
        seconds = int(time[-2::])
        frame = self.framerate * (seconds + 60 * minutes)

    cont = False

    if event not in self.classes:
        cont = True
    else:
        label = self.classes.index(event)

    return label, frame, cont

FeatureClipsfromJSON

Bases: FeaturefromJson

Class that inherits from FeaturefromJson to prepare features data as clips of features. This class is used for the pooling methods. The class has 2 behaviours for processing the data depending if it is for training or testing purposes.

Parameters:

Name	Type	Description	Default
path	str | List(str)	Path of the input. Can be a json file, a features file or a list of json files (list only for training purposes).	required
features_dir	str | List(str)	Path where the features are located. Must match the number of input paths.	required
classes	path	Path of the file containing the classes. Only used if input is feature file, otherwise list of classes is in the json file.	required
framerate	int	The framerate at which the features have been extracted. Default: 2.	2
window_size	int	The length of the window that will be used for the length of the clips. Default: 15.	15
train	bool	Whether we prepare data for training or for testing purposes. Default: True.	True

Source code in opensportslib/datasets/localization_dataset.py

class FeatureClipsfromJSON(FeaturefromJson):
    """Class that inherits from FeaturefromJson to prepare features data as clips of features.
    This class is used for the pooling methods.
    The class has 2 behaviours for processing the data depending if it is for training or testing purposes.

    Args:
        path (str|List(str)): Path of the input. Can be a json file, a features file or a list of json files (list only for training purposes).
        features_dir (str|List(str)): Path where the features are located. Must match the number of input paths.
        classes (path): Path of the file containing the classes. Only used if input is feature file, otherwise list of classes is in the json file.
        framerate (int): The framerate at which the features have been extracted.
            Default: 2.
        window_size (int): The length of the window that will be used for the length of the clips.
            Default: 15.
        train (bool): Whether we prepare data for training or for testing purposes.
            Default: True.
    """

    def __init__(
        self, path, features_dir, classes, framerate=2, window_size=15, train=True
    ):
        super().__init__(path, features_dir, classes, framerate)
        self.window_size_frame = window_size * framerate
        self.train = train
        self.features_clips = list()
        self.labels_clips = list()

        if self.train:
            for i, single_data_json in enumerate(self.data_json):
                if isinstance(path, list):
                    logging.info("Processing " + path[i])
                else:
                    logging.info("Processing " + path)
                # loop over videos
                for video in tqdm.tqdm(single_data_json["data"]):
                    # for video in tqdm(self.data_json["videos"]):
                    # Load features
                    features = np.load(
                        os.path.join(self.features_dir[i], video["inputs"][0]["path"])
                    )
                    features = features.reshape(-1, features.shape[-1])

                    # convert video features into clip features
                    features = feats2clip(
                        torch.from_numpy(features),
                        stride=self.window_size_frame,
                        clip_length=self.window_size_frame,
                    )

                    # Load labels
                    labels = np.zeros((features.shape[0], self.num_classes + 1))
                    labels[:, 0] = 1  # those are BG classes

                    # loop annotation for that video
                    for annotation in video.get("events", []):

                        label, frame, cont = self.annotation(annotation)

                        if cont:
                            continue

                        # if label outside temporal of view
                        if frame // self.window_size_frame >= labels.shape[0]:
                            continue

                        labels[frame // self.window_size_frame][0] = 0  # not BG anymore
                        labels[frame // self.window_size_frame][
                            label + 1
                        ] = 1  # that's my class

                    self.features_clips.append(features)
                    self.labels_clips.append(labels)

            self.features_clips = np.concatenate(self.features_clips)
            self.labels_clips = np.concatenate(self.labels_clips)

    def __getitem__(self, index):
        """
        Args:
            index (int): Index
        Returns:
            If training:
                clip_feat (np.array): clip of features.
                clip_labels (np.array): clip of labels.
            If testing:
                Name of the feature file.
                features (np.array): clip of features.
                labels (np.array): clip of labels.
        """
        if self.train:
            return self.features_clips[index, :, :], self.labels_clips[index, :]
        else:
            video = self.data_json[0]["data"][index]
            video_path = video["inputs"][0]["path"]
            # Load features
            if self.is_json:
                features = np.load(os.path.join(self.features_dir[0], video_path))
            else:
                features = np.load(os.path.join(video_path))
            features = features.reshape(-1, features.shape[-1])

            # Load labels
            labels = np.zeros((features.shape[0], self.num_classes))

            if "events" in video.keys():
                for annotation in video.get("events", []):

                    label, frame, cont = self.annotation(annotation)

                    if cont:
                        continue

                    frame = min(frame, features.shape[0] - 1)
                    labels[frame][label] = 1

            features = feats2clip(
                torch.from_numpy(features),
                stride=1,
                off=int(self.window_size_frame / 2),
                clip_length=self.window_size_frame,
            )

            return video_path, features, labels

    def __len__(self):
        if self.train:
            return len(self.features_clips)
        else:
            return len(self.data_json[0]["data"])

__getitem__(index)

Parameters:

Name	Type	Description	Default
index	int	Index	required

Returns: If training: clip_feat (np.array): clip of features. clip_labels (np.array): clip of labels. If testing: Name of the feature file. features (np.array): clip of features. labels (np.array): clip of labels.

Source code in opensportslib/datasets/localization_dataset.py

def __getitem__(self, index):
    """
    Args:
        index (int): Index
    Returns:
        If training:
            clip_feat (np.array): clip of features.
            clip_labels (np.array): clip of labels.
        If testing:
            Name of the feature file.
            features (np.array): clip of features.
            labels (np.array): clip of labels.
    """
    if self.train:
        return self.features_clips[index, :, :], self.labels_clips[index, :]
    else:
        video = self.data_json[0]["data"][index]
        video_path = video["inputs"][0]["path"]
        # Load features
        if self.is_json:
            features = np.load(os.path.join(self.features_dir[0], video_path))
        else:
            features = np.load(os.path.join(video_path))
        features = features.reshape(-1, features.shape[-1])

        # Load labels
        labels = np.zeros((features.shape[0], self.num_classes))

        if "events" in video.keys():
            for annotation in video.get("events", []):

                label, frame, cont = self.annotation(annotation)

                if cont:
                    continue

                frame = min(frame, features.shape[0] - 1)
                labels[frame][label] = 1

        features = feats2clip(
            torch.from_numpy(features),
            stride=1,
            off=int(self.window_size_frame / 2),
            clip_length=self.window_size_frame,
        )

        return video_path, features, labels

FeatureClipChunksfromJson

Bases: FeaturefromJson

Class that inherits from FeaturefromJson to prepare features data as clips of features based on a chunk approach. This class is used for the CALF method. The class has 2 behaviours for processing the data depending if it is for training or testing purposes.

Parameters:

Name	Type	Description	Default
path	str | List(str)	Path of the input. Can be a json file, a features file or a list of json files (list only for training purposes).	required
features_dir	str | List(str)	Path where the features are located. Must match the number of input paths.	required
classes	path	Path of the file containing the classes. Only used if input is feature file, otherwise list of classes is in the json file.	required
framerate	int	The framerate at which the features have been extracted. Default: 2.	2
chunk_size	int	Size of the chunk. Default: 240.	240
receptive_field	int	temporal receptive field of x seconds on both sides of the central frame in the temporal dimension of the 3D convolutions Default: 80.	80
chunks_per_epoch	int	Number of chunks per epoch. Default: 6000.	6000
gpu	bool	Whether gpu is used or not. Default: True.	True
train	bool	Whether we prepare data for training or for testing purposes. Default: True.	True

Source code in opensportslib/datasets/localization_dataset.py

class FeatureClipChunksfromJson(FeaturefromJson):
    """Class that inherits from FeaturefromJson to prepare features data as clips of features based on a chunk approach.
    This class is used for the CALF method.
    The class has 2 behaviours for processing the data depending if it is for training or testing purposes.

    Args:
        path (str|List(str)): Path of the input. Can be a json file, a features file or a list of json files (list only for training purposes).
        features_dir (str|List(str)): Path where the features are located. Must match the number of input paths.
        classes (path): Path of the file containing the classes. Only used if input is feature file, otherwise list of classes is in the json file.
        framerate (int): The framerate at which the features have been extracted.
            Default: 2.
        chunk_size (int): Size of the chunk.
            Default: 240.
        receptive_field (int):  temporal receptive field of x seconds on both sides of the central frame in the temporal dimension of the 3D convolutions
            Default: 80.
        chunks_per_epoch (int): Number of chunks per epoch.
            Default: 6000.
        gpu (bool): Whether gpu is used or not.
            Default: True.
        train (bool): Whether we prepare data for training or for testing purposes.
            Default: True.
    """

    def __init__(
        self,
        path,
        features_dir,
        classes,
        framerate=2,
        chunk_size=240,
        receptive_field=80,
        chunks_per_epoch=6000,
        gpu=True,
        train=True,
    ):
        super().__init__(path, features_dir, classes, framerate)
        self.chunk_size = chunk_size
        self.receptive_field = receptive_field
        self.chunks_per_epoch = chunks_per_epoch
        self.gpu = gpu
        self.train = train
        global K_V2
        if self.gpu >= 0:
            K_V2 = K_V2.cuda()
        self.K_parameters = K_V2 * framerate
        self.num_detections = 15

        if self.train:
            self.features_clips = list()
            self.labels_clips = list()
            self.anchors_clips = list()
            for i in np.arange(self.num_classes + 1):
                self.anchors_clips.append(list())

            game_counter = 0
            # loop over videos
            for i, single_data_json in enumerate(self.data_json):
                if isinstance(path, list):
                    logging.info("Processing " + path[i])
                else:
                    logging.info("Processing " + path)
                # loop over videos
                for video in tqdm.tqdm(single_data_json["data"]):
                    # for video in tqdm(self.data_json["videos"]):
                    # Load features
                    features = np.load(
                        os.path.join(self.features_dir[i], video["inputs"][0]["path"])
                    )

                    # Load labels
                    labels = np.zeros((features.shape[0], self.num_classes))

                    # loop annotation for that video
                    for annotation in video.get("events", []):

                        label, frame, cont = self.annotation(annotation)

                        if cont:
                            continue

                        frame = min(frame, features.shape[0] - 1)
                        labels[frame][label] = 1

                    shift_half = oneHotToShifts(labels, self.K_parameters.cpu().numpy())

                    anchors_half = getChunks_anchors(
                        shift_half,
                        game_counter,
                        self.K_parameters.cpu().numpy(),
                        self.chunk_size,
                        self.receptive_field,
                    )

                    game_counter = game_counter + 1

                    self.features_clips.append(features)
                    self.labels_clips.append(shift_half)

                    for anchor in anchors_half:
                        self.anchors_clips[anchor[2]].append(anchor)

    def __getitem__(self, index):
        """
        Args:
            index (int): Index
        Returns:
            If training:
                clip_feat (np.array): clip of features.
                clip_labels (np.array): clip of labels.
                clip_targets (np.array): clip of targets.
            If testing:
                features (np.array): clip of features.
                labels (np.array): clip of labels.
        """
        if self.train:
            # Retrieve the game index and the anchor
            class_selection = random.randint(0, self.num_classes)
            event_selection = random.randint(
                0, len(self.anchors_clips[class_selection]) - 1
            )
            game_index = self.anchors_clips[class_selection][event_selection][0]
            anchor = self.anchors_clips[class_selection][event_selection][1]

            # Compute the shift for event chunks
            if class_selection < self.num_classes:
                shift = np.random.randint(
                    -self.chunk_size + self.receptive_field, -self.receptive_field
                )
                start = anchor + shift
            # Compute the shift for non-event chunks
            else:
                start = random.randint(anchor[0], anchor[1] - self.chunk_size)
            if start < 0:
                start = 0
            if start + self.chunk_size >= self.features_clips[game_index].shape[0]:
                start = self.features_clips[game_index].shape[0] - self.chunk_size - 1

            # Extract the clips
            clip_feat = self.features_clips[game_index][start : start + self.chunk_size]
            clip_labels = self.labels_clips[game_index][start : start + self.chunk_size]

            # Put loss to zero outside receptive field
            clip_labels[0 : int(np.ceil(self.receptive_field / 2)), :] = -1
            clip_labels[-int(np.ceil(self.receptive_field / 2)) :, :] = -1

            # Get the spotting target
            clip_targets = getTimestampTargets(
                np.array([clip_labels]), self.num_detections
            )[0]

            return (
                torch.from_numpy(clip_feat),
                torch.from_numpy(clip_labels),
                torch.from_numpy(clip_targets),
            )
        else:
            video = self.data_json[0]["data"][index]
            video_path = video["inputs"][0]["path"]
            # Load features
            if self.is_json:
                features = np.load(os.path.join(self.features_dir[0], video_path))
            else:
                features = np.load(os.path.join(video_path))

            # Load labels
            labels = np.zeros((features.shape[0], self.num_classes))

            if "annotations" in video.keys():
                for annotation in video.get("events",[]):

                    label, frame, cont = self.annotation(annotation)
                    if cont:
                        continue

                    value = 1
                    if "visibility" in annotation.keys():
                        if annotation["visibility"] == "not shown":
                            value = -1

                    frame = min(frame, features.shape[0] - 1)
                    labels[frame][label] = value

            features = feats2clip(
                torch.from_numpy(features),
                stride=self.chunk_size - self.receptive_field,
                clip_length=self.chunk_size,
                modif_last_index=True,
            )

            return features, torch.from_numpy(labels)

    def __len__(self):
        if self.train:
            return self.chunks_per_epoch
        else:
            return len(self.data_json[0]["data"])

__getitem__(index)

Parameters:

Name	Type	Description	Default
index	int	Index	required

Returns: If training: clip_feat (np.array): clip of features. clip_labels (np.array): clip of labels. clip_targets (np.array): clip of targets. If testing: features (np.array): clip of features. labels (np.array): clip of labels.

Source code in opensportslib/datasets/localization_dataset.py

def __getitem__(self, index):
    """
    Args:
        index (int): Index
    Returns:
        If training:
            clip_feat (np.array): clip of features.
            clip_labels (np.array): clip of labels.
            clip_targets (np.array): clip of targets.
        If testing:
            features (np.array): clip of features.
            labels (np.array): clip of labels.
    """
    if self.train:
        # Retrieve the game index and the anchor
        class_selection = random.randint(0, self.num_classes)
        event_selection = random.randint(
            0, len(self.anchors_clips[class_selection]) - 1
        )
        game_index = self.anchors_clips[class_selection][event_selection][0]
        anchor = self.anchors_clips[class_selection][event_selection][1]

        # Compute the shift for event chunks
        if class_selection < self.num_classes:
            shift = np.random.randint(
                -self.chunk_size + self.receptive_field, -self.receptive_field
            )
            start = anchor + shift
        # Compute the shift for non-event chunks
        else:
            start = random.randint(anchor[0], anchor[1] - self.chunk_size)
        if start < 0:
            start = 0
        if start + self.chunk_size >= self.features_clips[game_index].shape[0]:
            start = self.features_clips[game_index].shape[0] - self.chunk_size - 1

        # Extract the clips
        clip_feat = self.features_clips[game_index][start : start + self.chunk_size]
        clip_labels = self.labels_clips[game_index][start : start + self.chunk_size]

        # Put loss to zero outside receptive field
        clip_labels[0 : int(np.ceil(self.receptive_field / 2)), :] = -1
        clip_labels[-int(np.ceil(self.receptive_field / 2)) :, :] = -1

        # Get the spotting target
        clip_targets = getTimestampTargets(
            np.array([clip_labels]), self.num_detections
        )[0]

        return (
            torch.from_numpy(clip_feat),
            torch.from_numpy(clip_labels),
            torch.from_numpy(clip_targets),
        )
    else:
        video = self.data_json[0]["data"][index]
        video_path = video["inputs"][0]["path"]
        # Load features
        if self.is_json:
            features = np.load(os.path.join(self.features_dir[0], video_path))
        else:
            features = np.load(os.path.join(video_path))

        # Load labels
        labels = np.zeros((features.shape[0], self.num_classes))

        if "annotations" in video.keys():
            for annotation in video.get("events",[]):

                label, frame, cont = self.annotation(annotation)
                if cont:
                    continue

                value = 1
                if "visibility" in annotation.keys():
                    if annotation["visibility"] == "not shown":
                        value = -1

                frame = min(frame, features.shape[0] - 1)
                labels[frame][label] = value

        features = feats2clip(
            torch.from_numpy(features),
            stride=self.chunk_size - self.receptive_field,
            clip_length=self.chunk_size,
            modif_last_index=True,
        )

        return features, torch.from_numpy(labels)

SoccerNetGame

Bases: Dataset

Parent class that is used to modularise common operations between the classes that prepare features of a single game of the soccernet dataset.

Parameters:

Name	Type	Description	Default
path	string	Path of the game.	required
features	str	Name of the features. Default: ResNET_PCA512.npy.	'ResNET_PCA512.npy'
version	int	Version of the dataset. Default: 1.	1
framerate	int	The framerate at which the features have been extracted. Default: 2.	2

Source code in opensportslib/datasets/localization_dataset.py

class SoccerNetGame(Dataset):
    """Parent class that is used to modularise common operations between the classes that prepare features of a single game of the soccernet dataset.

    Args:
        path (string): Path of the game.
        features (str): Name of the features.
            Default: ResNET_PCA512.npy.
        version (int): Version of the dataset.
            Default: 1.
        framerate (int): The framerate at which the features have been extracted.
            Default: 2.
    """

    def __init__(self, path, features="ResNET_PCA512.npy", version=1, framerate=2):
        self.path = path
        self.framerate = framerate
        self.version = version
        self.features = features
        self.listGames = [self.path]
        if version == 1:
            self.dict_event = EVENT_DICTIONARY_V1
            self.num_classes = 3
        elif version == 2:
            self.dict_event = EVENT_DICTIONARY_V2
            self.num_classes = 17

    def __getitem__(self, index):
        pass

    def __len__(self):
        return 2

    def load_features(self):
        self.feat_half1 = np.load(os.path.join(self.path, "1_" + self.features))
        self.feat_half2 = np.load(os.path.join(self.path, "2_" + self.features))

SoccerNetGameClips

Bases: SoccerNetGame

Class that inherits from SoccerNetGame to prepare features data of a single game as clips of features. This class is used for the pooling methods.

Parameters:

Name	Type	Description	Default
path	string	Path of the game.	required
features	str	Name of the features. Default: ResNET_PCA512.npy.	'ResNET_PCA512.npy'
version	int	Version of the dataset. Default: 1.	1
framerate	int	The framerate at which the features have been extracted. Default: 2.	2
window_size	int	The length of the window that will be used for the length of the clips. Default: 15.	15

Source code in opensportslib/datasets/localization_dataset.py

class SoccerNetGameClips(SoccerNetGame):
    """Class that inherits from SoccerNetGame to prepare features data of a single game as clips of features.
    This class is used for the pooling methods.

    Args:
        path (string): Path of the game.
        features (str): Name of the features.
            Default: ResNET_PCA512.npy.
        version (int): Version of the dataset.
            Default: 1.
        framerate (int): The framerate at which the features have been extracted.
            Default: 2.
        window_size (int): The length of the window that will be used for the length of the clips.
            Default: 15.
    """

    def __init__(
        self, path, features="ResNET_PCA512.npy", version=1, framerate=2, window_size=15
    ):

        super().__init__(path, features, version, framerate)
        self.window_size_frame = window_size * self.framerate

    def __getitem__(self, index):
        """
        Args:
            index (int): Index
        Returns:
            path: path of the game.
            clip_feat_1 (np.array): clip of features for the first half.
            clip_feat_2 (np.array): clip of features for the second half.
            empty list
            empty list
        """
        # Load features
        feat_half1, feat_half2 = self.load_features()

        feat_half1 = feats2clip(
            torch.from_numpy(feat_half1),
            stride=1,
            off=int(self.window_size_frame / 2),
            clip_length=self.window_size_frame,
        )

        feat_half2 = feats2clip(
            torch.from_numpy(feat_half2),
            stride=1,
            off=int(self.window_size_frame / 2),
            clip_length=self.window_size_frame,
        )

        return self.path, feat_half1, feat_half2, [], []

    def load_features(self):
        super().load_features()
        self.feat_half1 = self.feat_half1.reshape(-1, self.feat_half1.shape[-1])
        self.feat_half2 = self.feat_half2.reshape(-1, self.feat_half2.shape[-1])
        return self.feat_half1, self.feat_half2

__getitem__(index)

Parameters:

Name	Type	Description	Default
index	int	Index	required

Returns: path: path of the game. clip_feat_1 (np.array): clip of features for the first half. clip_feat_2 (np.array): clip of features for the second half. empty list empty list

Source code in opensportslib/datasets/localization_dataset.py

def __getitem__(self, index):
    """
    Args:
        index (int): Index
    Returns:
        path: path of the game.
        clip_feat_1 (np.array): clip of features for the first half.
        clip_feat_2 (np.array): clip of features for the second half.
        empty list
        empty list
    """
    # Load features
    feat_half1, feat_half2 = self.load_features()

    feat_half1 = feats2clip(
        torch.from_numpy(feat_half1),
        stride=1,
        off=int(self.window_size_frame / 2),
        clip_length=self.window_size_frame,
    )

    feat_half2 = feats2clip(
        torch.from_numpy(feat_half2),
        stride=1,
        off=int(self.window_size_frame / 2),
        clip_length=self.window_size_frame,
    )

    return self.path, feat_half1, feat_half2, [], []

SoccerNetGameClipsChunks

Bases: SoccerNetGame

Class that inherits from SoccerNetGame to prepare features data of a single game as clips of features based on a chunk approach. This class is used for the CALF method.

Parameters:

Name	Type	Description	Default
path	string	Path of the game.	required
features	str	Name of the features. Default: ResNET_PCA512.npy.	'ResNET_PCA512.npy'
version	int	Version of the dataset. Default: 1.	required
framerate	int	The framerate at which the features have been extracted. Default: 2.	2
chunk_size	int	Size of the chunk. Default: 240.	240
receptive_field	int	temporal receptive field of x seconds on both sides of the central frame in the temporal dimension of the 3D convolutions Default: 80.	80

Source code in opensportslib/datasets/localization_dataset.py

class SoccerNetGameClipsChunks(SoccerNetGame):
    """Class that inherits from SoccerNetGame to prepare features data of a single game as clips of features based on a chunk approach.
    This class is used for the CALF method.

    Args:
        path (string): Path of the game.
        features (str): Name of the features.
            Default: ResNET_PCA512.npy.
        version (int): Version of the dataset.
            Default: 1.
        framerate (int): The framerate at which the features have been extracted.
            Default: 2.
        chunk_size (int): Size of the chunk.
            Default: 240.
        receptive_field (int):  temporal receptive field of x seconds on both sides of the central frame in the temporal dimension of the 3D convolutions
            Default: 80.
    """

    def __init__(
        self,
        path,
        features="ResNET_PCA512.npy",
        framerate=2,
        chunk_size=240,
        receptive_field=80,
    ):
        super().__init__(path, features, 2, framerate)
        self.chunk_size = chunk_size
        self.receptive_field = receptive_field

    def __getitem__(self, index):
        """
        Args:
            index (int): Index
        Returns:
            clip_feat_1 (np.array): clip of features for the first half.
            clip_feat_2 (np.array): clip of features for the second half.
            clip_label_1 (np.array): clip of labels for the first half.
            clip_label_2 (np.array): clip of labels for the second half.
        """
        # Load features
        feat_half1, feat_half2 = self.load_features()

        # Load labels
        label_half1, label_half2 = self.load_labels(feat_half1, feat_half2)

        feat_half1 = feats2clip(
            torch.from_numpy(feat_half1),
            stride=self.chunk_size - self.receptive_field,
            clip_length=self.chunk_size,
            modif_last_index=True,
        )

        feat_half2 = feats2clip(
            torch.from_numpy(feat_half2),
            stride=self.chunk_size - self.receptive_field,
            clip_length=self.chunk_size,
            modif_last_index=True,
        )

        return (
            feat_half1,
            feat_half2,
            torch.from_numpy(label_half1),
            torch.from_numpy(label_half2),
        )

    def load_features(
        self,
    ):
        super().load_features()
        return self.feat_half1, self.feat_half2

    def load_labels(self, feat_half1, feat_half2):
        self.label_half1 = np.zeros((feat_half1.shape[0], self.num_classes))
        self.label_half2 = np.zeros((feat_half2.shape[0], self.num_classes))
        return self.label_half1, self.label_half2

__getitem__(index)

Parameters:

Name	Type	Description	Default
index	int	Index	required

Returns: clip_feat_1 (np.array): clip of features for the first half. clip_feat_2 (np.array): clip of features for the second half. clip_label_1 (np.array): clip of labels for the first half. clip_label_2 (np.array): clip of labels for the second half.

Source code in opensportslib/datasets/localization_dataset.py

def __getitem__(self, index):
    """
    Args:
        index (int): Index
    Returns:
        clip_feat_1 (np.array): clip of features for the first half.
        clip_feat_2 (np.array): clip of features for the second half.
        clip_label_1 (np.array): clip of labels for the first half.
        clip_label_2 (np.array): clip of labels for the second half.
    """
    # Load features
    feat_half1, feat_half2 = self.load_features()

    # Load labels
    label_half1, label_half2 = self.load_labels(feat_half1, feat_half2)

    feat_half1 = feats2clip(
        torch.from_numpy(feat_half1),
        stride=self.chunk_size - self.receptive_field,
        clip_length=self.chunk_size,
        modif_last_index=True,
    )

    feat_half2 = feats2clip(
        torch.from_numpy(feat_half2),
        stride=self.chunk_size - self.receptive_field,
        clip_length=self.chunk_size,
        modif_last_index=True,
    )

    return (
        feat_half1,
        feat_half2,
        torch.from_numpy(label_half1),
        torch.from_numpy(label_half2),
    )

SoccerNet

Bases: Dataset

Parent class that is used to modularise common operations between the classes that prepare features of a split of a soccernet dataset.

Parameters:

Name	Type	Description	Default
path	string	Data root of the feature files.	required
features	str	Name of the features. Default: ResNET_PCA512.npy.	'ResNET_PCA512.npy'
split	List(string)	List of splits. Default: ["train"].	['train']
version	int	Version of the dataset. Default: 1.	1
framerate	int	The framerate at which the features have been extracted. Default: 2.	2

Source code in opensportslib/datasets/localization_dataset.py

class SoccerNet(Dataset):
    """Parent class that is used to modularise common operations between the classes that prepare features of a split of a soccernet dataset.

    Args:
        path (string): Data root of the feature files.
        features (str): Name of the features.
            Default: ResNET_PCA512.npy.
        split (List(string)): List of splits.
            Default: ["train"].
        version (int): Version of the dataset.
            Default: 1.
        framerate (int): The framerate at which the features have been extracted.
            Default: 2.
    """

    def __init__(
        self,
        path,
        features="ResNET_PCA512.npy",
        split=["train"],
        version=1,
        framerate=2,
    ):
        self.path = path
        # split=["train"] if clips else ["test"]
        self.listGames = getListGames(split)
        self.features = features
        self.framerate = framerate
        self.split = split
        self.version = version

        if version == 1:
            self.dict_event = EVENT_DICTIONARY_V1
            self.num_classes = 3
            self.labels = "Labels.json"
        elif version == 2:
            self.dict_event = EVENT_DICTIONARY_V2
            self.num_classes = 17
            self.labels = "Labels-v2.json"

    def __getitem__(self, index):
        pass

    def __len__(self):
        return len(self.listGames)

    def load_features(self, index=0, game=""):
        """Load features from files.

        Args:
            index (int): Used for testing purpose to retrieve the game based on the index.
                Default: 0.
            game (string): Used for training purpose, this is the name of the game.
                Default: "".
        """
        if self.train:
            game = game
        else:
            game = self.listGames[index].replace(" ", "_")
        self.feat_half1 = np.load(os.path.join(self.path, game, "1_" + self.features))
        self.feat_half2 = np.load(os.path.join(self.path, game, "2_" + self.features))

    def load_labels(self, feat_half1, feat_half2, number_classes):
        # Load labels
        self.label_half1 = np.zeros((feat_half1.shape[0], number_classes))
        self.label_half2 = np.zeros((feat_half2.shape[0], number_classes))

    def annotation(self, annotation):
        """Process an annotation to derive the frame number, the class index and a boolean.
        Args:
            annotation (dict): It must contains the keys "gameTime" and "label".

        Returns:
            label (int): The index of the class.
            half (int): Which half of the game.
            frame (int): The index of the frame.
            cont (bool): Whether to continue in the loop or not.
        """
        time = annotation["gameTime"]
        event = annotation["label"]

        half = int(time[0])

        minutes = int(time[-5:-3])
        seconds = int(time[-2::])
        frame = self.framerate * (seconds + 60 * minutes)

        cont = False
        if self.version == 1:
            if "card" in event:
                label = 0
            elif "subs" in event:
                label = 1
            elif "soccer" in event:
                label = 2
            else:
                cont = True
        elif self.version == 2:
            if event not in self.dict_event:
                cont = True
            else:
                label = self.dict_event[event]

        return label, half, frame, cont

annotation(annotation)

Process an annotation to derive the frame number, the class index and a boolean. Args: annotation (dict): It must contains the keys "gameTime" and "label".

Returns:

Name	Type	Description
label	int	The index of the class.
half	int	Which half of the game.
frame	int	The index of the frame.
cont	bool	Whether to continue in the loop or not.

Source code in opensportslib/datasets/localization_dataset.py

def annotation(self, annotation):
    """Process an annotation to derive the frame number, the class index and a boolean.
    Args:
        annotation (dict): It must contains the keys "gameTime" and "label".

    Returns:
        label (int): The index of the class.
        half (int): Which half of the game.
        frame (int): The index of the frame.
        cont (bool): Whether to continue in the loop or not.
    """
    time = annotation["gameTime"]
    event = annotation["label"]

    half = int(time[0])

    minutes = int(time[-5:-3])
    seconds = int(time[-2::])
    frame = self.framerate * (seconds + 60 * minutes)

    cont = False
    if self.version == 1:
        if "card" in event:
            label = 0
        elif "subs" in event:
            label = 1
        elif "soccer" in event:
            label = 2
        else:
            cont = True
    elif self.version == 2:
        if event not in self.dict_event:
            cont = True
        else:
            label = self.dict_event[event]

    return label, half, frame, cont

load_features(index=0, game='')

Load features from files.

Parameters:

Name	Type	Description	Default
index	int	Used for testing purpose to retrieve the game based on the index. Default: 0.	0
game	string	Used for training purpose, this is the name of the game. Default: "".	''

Source code in opensportslib/datasets/localization_dataset.py

def load_features(self, index=0, game=""):
    """Load features from files.

    Args:
        index (int): Used for testing purpose to retrieve the game based on the index.
            Default: 0.
        game (string): Used for training purpose, this is the name of the game.
            Default: "".
    """
    if self.train:
        game = game
    else:
        game = self.listGames[index].replace(" ", "_")
    self.feat_half1 = np.load(os.path.join(self.path, game, "1_" + self.features))
    self.feat_half2 = np.load(os.path.join(self.path, game, "2_" + self.features))

SoccerNetClips

Bases: SoccerNet

Class that inherits from SoccerNetGame to prepare features data of a split as clips of features. This class is used for the pooling methods.

Parameters:

Name	Type	Description	Default
path	string	Data root of the feature files.	required
features	str	Name of the features. Default: ResNET_PCA512.npy.	'ResNET_PCA512.npy'
split	List(string)	List of splits. Default: ["train"].	['train']
version	int	Version of the dataset. Default: 1.	1
framerate	int	The framerate at which the features have been extracted. Default: 2.	2
window_size	int	The length of the window that will be used for the length of the clips. Default: 15.	15
train	bool	Whether training or testing. Default: True.	True

Source code in opensportslib/datasets/localization_dataset.py

class SoccerNetClips(SoccerNet):
    """Class that inherits from SoccerNetGame to prepare features data of a split as clips of features.
    This class is used for the pooling methods.

    Args:
        path (string): Data root of the feature files.
        features (str): Name of the features.
            Default: ResNET_PCA512.npy.
        split (List(string)): List of splits.
            Default: ["train"].
        version (int): Version of the dataset.
            Default: 1.
        framerate (int): The framerate at which the features have been extracted.
            Default: 2.
        window_size (int): The length of the window that will be used for the length of the clips.
            Default: 15.
        train (bool): Whether training or testing.
            Default: True.
    """

    def __init__(
        self,
        path,
        features="ResNET_PCA512.npy",
        split=["train"],
        version=1,
        framerate=2,
        window_size=15,
        train=True,
    ):

        super().__init__(path, features, split, version, framerate)
        self.window_size_frame = window_size * self.framerate
        self.train = train

        logging.info("Checking/Download features and labels locally")
        downloader = SoccerNetDownloader(path)

        # if train :
        #     downloader.downloadGames(files=[self.labels, f"1_{self.features}", f"2_{self.features}"], split=split, verbose=False,randomized=True)
        # else :
        #     for s in split:
        #         if s == "challenge":
        #             downloader.downloadGames(files=[f"1_{self.features}", f"2_{self.features}"], split=[s], verbose=False,randomized=True)
        #         else:
        #             downloader.downloadGames(files=[self.labels, f"1_{self.features}", f"2_{self.features}"], split=[s], verbose=False,randomized=True)

        if train:
            logging.info("Pre-compute clips")

            self.game_feats = list()
            self.game_labels = list()

            for game in tqdm.tqdm(self.listGames):
                game = game.replace(" ", "_")
                # Load features
                feat_half1, feat_half2 = self.load_features(game=game)

                feat_half1 = feats2clip(
                    torch.from_numpy(feat_half1),
                    stride=self.window_size_frame,
                    clip_length=self.window_size_frame,
                )
                feat_half2 = feats2clip(
                    torch.from_numpy(feat_half2),
                    stride=self.window_size_frame,
                    clip_length=self.window_size_frame,
                )

                # Load labels
                labels = json.load(open(os.path.join(self.path, game, self.labels)))

                label_half1, label_half2 = self.load_labels(feat_half1, feat_half2)

                for annotation in labels["annotations"]:

                    label, half, frame, cont = self.annotation(annotation)
                    if cont:
                        continue

                    # if label outside temporal of view
                    if (
                        half == 1
                        and frame // self.window_size_frame >= label_half1.shape[0]
                    ):
                        continue
                    if (
                        half == 2
                        and frame // self.window_size_frame >= label_half2.shape[0]
                    ):
                        continue

                    if half == 1:
                        label_half1[frame // self.window_size_frame][
                            0
                        ] = 0  # not BG anymore
                        label_half1[frame // self.window_size_frame][
                            label + 1
                        ] = 1  # that's my class

                    if half == 2:
                        label_half2[frame // self.window_size_frame][
                            0
                        ] = 0  # not BG anymore
                        label_half2[frame // self.window_size_frame][
                            label + 1
                        ] = 1  # that's my class
                self.game_feats.append(feat_half1)
                self.game_feats.append(feat_half2)
                self.game_labels.append(label_half1)
                self.game_labels.append(label_half2)

            self.game_feats = np.concatenate(self.game_feats)
            self.game_labels = np.concatenate(self.game_labels)

    def __getitem__(self, index):
        """
        Args:
            index (int): Index
        Returns:
            if train :
                clip_feat (np.array): clip of features.
                clip_labels (np.array): clip of labels for the segmentation.
            if testing:
                name of the game.
                clip_feat_1 (np.array): clip of features for the first half.
                clip_feat_2 (np.array): clip of features for the second half.
                clip_label_1 (np.array): clip of labels for the first half.
                clip_label_2 (np.array): clip of labels for the second half.
        """
        if self.train:
            return self.game_feats[index, :, :], self.game_labels[index, :]
        else:
            # Load features
            feat_half1, feat_half2 = self.load_features(index=index)

            # Load labels
            label_half1, label_half2 = self.load_labels(feat_half1, feat_half2)

            # check if annotation exists
            if os.path.exists(
                os.path.join(self.path, self.listGames[index], self.labels)
            ):
                labels = json.load(
                    open(os.path.join(self.path, self.listGames[index], self.labels))
                )

                for annotation in labels["annotations"]:

                    label, half, frame, cont = self.annotation(annotation)
                    if cont:
                        continue

                    value = 1
                    if "visibility" in annotation.keys():
                        if annotation["visibility"] == "not shown":
                            value = -1

                    if half == 1:
                        frame = min(frame, feat_half1.shape[0] - 1)
                        label_half1[frame][label] = value

                    if half == 2:
                        frame = min(frame, feat_half2.shape[0] - 1)
                        label_half2[frame][label] = value

            feat_half1 = feats2clip(
                torch.from_numpy(feat_half1),
                stride=1,
                off=int(self.window_size_frame / 2),
                clip_length=self.window_size_frame,
            )

            feat_half2 = feats2clip(
                torch.from_numpy(feat_half2),
                stride=1,
                off=int(self.window_size_frame / 2),
                clip_length=self.window_size_frame,
            )

            return (
                self.listGames[index],
                feat_half1,
                feat_half2,
                label_half1,
                label_half2,
            )

    def __len__(self):
        if self.train:
            return len(self.game_feats)
        else:
            return super().__len__()

    def load_features(self, index=0, game=""):
        super().load_features(index, game)
        self.feat_half1 = self.feat_half1.reshape(-1, self.feat_half1.shape[-1])
        self.feat_half2 = self.feat_half2.reshape(-1, self.feat_half2.shape[-1])
        return self.feat_half1, self.feat_half2

    def load_labels(self, feat_half1, feat_half2):
        super().load_labels(
            feat_half1,
            feat_half2,
            self.num_classes + 1 if self.train else self.num_classes,
        )
        if self.train:
            self.label_half1[:, 0] = 1  # those are BG classes
            self.label_half2[:, 0] = 1  # those are BG classes
        return self.label_half1, self.label_half2

__getitem__(index)

Parameters:

Name	Type	Description	Default
index	int	Index	required

Returns: if train : clip_feat (np.array): clip of features. clip_labels (np.array): clip of labels for the segmentation. if testing: name of the game. clip_feat_1 (np.array): clip of features for the first half. clip_feat_2 (np.array): clip of features for the second half. clip_label_1 (np.array): clip of labels for the first half. clip_label_2 (np.array): clip of labels for the second half.

Source code in opensportslib/datasets/localization_dataset.py

def __getitem__(self, index):
    """
    Args:
        index (int): Index
    Returns:
        if train :
            clip_feat (np.array): clip of features.
            clip_labels (np.array): clip of labels for the segmentation.
        if testing:
            name of the game.
            clip_feat_1 (np.array): clip of features for the first half.
            clip_feat_2 (np.array): clip of features for the second half.
            clip_label_1 (np.array): clip of labels for the first half.
            clip_label_2 (np.array): clip of labels for the second half.
    """
    if self.train:
        return self.game_feats[index, :, :], self.game_labels[index, :]
    else:
        # Load features
        feat_half1, feat_half2 = self.load_features(index=index)

        # Load labels
        label_half1, label_half2 = self.load_labels(feat_half1, feat_half2)

        # check if annotation exists
        if os.path.exists(
            os.path.join(self.path, self.listGames[index], self.labels)
        ):
            labels = json.load(
                open(os.path.join(self.path, self.listGames[index], self.labels))
            )

            for annotation in labels["annotations"]:

                label, half, frame, cont = self.annotation(annotation)
                if cont:
                    continue

                value = 1
                if "visibility" in annotation.keys():
                    if annotation["visibility"] == "not shown":
                        value = -1

                if half == 1:
                    frame = min(frame, feat_half1.shape[0] - 1)
                    label_half1[frame][label] = value

                if half == 2:
                    frame = min(frame, feat_half2.shape[0] - 1)
                    label_half2[frame][label] = value

        feat_half1 = feats2clip(
            torch.from_numpy(feat_half1),
            stride=1,
            off=int(self.window_size_frame / 2),
            clip_length=self.window_size_frame,
        )

        feat_half2 = feats2clip(
            torch.from_numpy(feat_half2),
            stride=1,
            off=int(self.window_size_frame / 2),
            clip_length=self.window_size_frame,
        )

        return (
            self.listGames[index],
            feat_half1,
            feat_half2,
            label_half1,
            label_half2,
        )

SoccerNetClipsChunks

Bases: SoccerNet

Class that inherits from SoccerNetGame to prepare features data of a split as clips of features with a chunk approach. This class is used for the CALF method.

Parameters:

Name	Type	Description	Default
path	str	Data root of the feature files.	required
features	str	Name of the features. Default: ResNET_PCA512.npy.	'ResNET_PCA512.npy'
split	string	split. Default: "train".	'train'
version	int	Version of the dataset. Default: 1.	required
framerate	int	The framerate at which the features have been extracted. Default: 2.	2
chunk_size	int	Size of the chunk. Default: 240.	240
receptive_field	int	temporal receptive field of x seconds on both sides of the central frame in the temporal dimension of the 3D convolutions Default: 80.	80
chunks_per_epoch	int	Number of chunks per epoch. Default: 6000.	6000
gpu	bool	Whether gpu is used or not. Default: True.	True
train	bool	Whether we prepare data for training or for testing purposes. Default: True.	True

Source code in opensportslib/datasets/localization_dataset.py

class SoccerNetClipsChunks(SoccerNet):
    """Class that inherits from SoccerNetGame to prepare features data of a split as clips of features with a chunk approach.
    This class is used for the CALF method.

    Args:
        path (str): Data root of the feature files.
        features (str): Name of the features.
            Default: ResNET_PCA512.npy.
        split (string): split.
            Default: "train".
        version (int): Version of the dataset.
            Default: 1.
        framerate (int): The framerate at which the features have been extracted.
            Default: 2.
        chunk_size (int): Size of the chunk.
            Default: 240.
        receptive_field (int):  temporal receptive field of x seconds on both sides of the central frame in the temporal dimension of the 3D convolutions
            Default: 80.
        chunks_per_epoch (int): Number of chunks per epoch.
            Default: 6000.
        gpu (bool): Whether gpu is used or not.
            Default: True.
        train (bool): Whether we prepare data for training or for testing purposes.
            Default: True.
    """

    def __init__(
        self,
        path,
        features="ResNET_PCA512.npy",
        split="train",
        framerate=2,
        chunk_size=240,
        receptive_field=80,
        chunks_per_epoch=6000,
        gpu=True,
        train=True,
    ):
        super().__init__(path, features, split, 2, framerate)
        self.chunk_size = chunk_size
        self.receptive_field = receptive_field
        self.chunks_per_epoch = chunks_per_epoch
        self.gpu = gpu
        self.train = train
        global K_V2
        if self.gpu >= 0:
            K_V2 = K_V2.cuda()
        self.K_parameters = K_V2 * framerate
        self.num_detections = 15

        logging.info("Checking/Download features and labels locally")
        downloader = SoccerNetDownloader(path)

        # if train:
        #     downloader.downloadGames(files=[self.labels, f"1_{self.features}", f"2_{self.features}"], split=split, verbose=False)
        # else:
        #     for s in split:
        #         if s == "challenge":
        #             downloader.downloadGames(files=[f"1_{self.features}", f"2_{self.features}"], split=[s], verbose=False)
        #         else:
        #             downloader.downloadGames(files=[self.labels, f"1_{self.features}", f"2_{self.features}"], split=[s], verbose=False)

        if train:
            logging.info("Pre-compute clips")

            self.game_feats = list()
            self.game_labels = list()
            self.game_anchors = list()
            for i in np.arange(self.num_classes + 1):
                self.game_anchors.append(list())

            game_counter = 0
            for game in tqdm.tqdm(self.listGames):
                game = game.replace(" ", "_")
                # Load features
                feat_half1, feat_half2 = self.load_features(game=game)
                # Load labels
                labels = json.load(open(os.path.join(self.path, game, self.labels)))

                label_half1, label_half2 = self.load_labels(feat_half1, feat_half2)

                for annotation in labels["annotations"]:

                    label, half, frame, cont = self.annotation(annotation)
                    if cont:
                        continue
                    if half == 1:
                        frame = min(frame, feat_half1.shape[0] - 1)
                        label_half1[frame][label] = 1

                    if half == 2:
                        frame = min(frame, feat_half2.shape[0] - 1)
                        label_half2[frame][label] = 1

                shift_half1 = oneHotToShifts(
                    label_half1, self.K_parameters.cpu().numpy()
                )
                shift_half2 = oneHotToShifts(
                    label_half2, self.K_parameters.cpu().numpy()
                )

                anchors_half1 = getChunks_anchors(
                    shift_half1,
                    game_counter,
                    self.K_parameters.cpu().numpy(),
                    self.chunk_size,
                    self.receptive_field,
                )

                game_counter = game_counter + 1

                anchors_half2 = getChunks_anchors(
                    shift_half2,
                    game_counter,
                    self.K_parameters.cpu().numpy(),
                    self.chunk_size,
                    self.receptive_field,
                )

                game_counter = game_counter + 1
                self.game_feats.append(feat_half1)
                self.game_feats.append(feat_half2)
                self.game_labels.append(shift_half1)
                self.game_labels.append(shift_half2)

                for anchor in anchors_half1:
                    self.game_anchors[anchor[2]].append(anchor)
                for anchor in anchors_half2:
                    self.game_anchors[anchor[2]].append(anchor)

    def __getitem__(self, index):
        """
        Args:
            index (int): Index
        Returns:
            If training:
                clip_feat (np.array): clip of features.
                clip_labels (np.array): clip of labels.
                clip_targets (np.array): clip of targets.
            If testing:
                clip_feat_1 (np.array): clip of features for the first half.
                clip_feat_2 (np.array): clip of features for the second half.
                clip_labels_1 (np.array): clip of labels for the first half.
                clip_labels_2 (np.array): clip of labels for the second half.

        """
        if self.train:
            # Retrieve the game index and the anchor
            class_selection = random.randint(0, self.num_classes)
            event_selection = random.randint(
                0, len(self.game_anchors[class_selection]) - 1
            )
            game_index = self.game_anchors[class_selection][event_selection][0]
            anchor = self.game_anchors[class_selection][event_selection][1]

            # Compute the shift for event chunks
            if class_selection < self.num_classes:
                shift = np.random.randint(
                    -self.chunk_size + self.receptive_field, -self.receptive_field
                )
                start = anchor + shift
            # Compute the shift for non-event chunks
            else:
                start = random.randint(anchor[0], anchor[1] - self.chunk_size)
            if start < 0:
                start = 0
            if start + self.chunk_size >= self.game_feats[game_index].shape[0]:
                start = self.game_feats[game_index].shape[0] - self.chunk_size - 1

            # Extract the clips
            clip_feat = self.game_feats[game_index][start : start + self.chunk_size]
            clip_labels = self.game_labels[game_index][start : start + self.chunk_size]

            # Put loss to zero outside receptive field
            clip_labels[0 : int(np.ceil(self.receptive_field / 2)), :] = -1
            clip_labels[-int(np.ceil(self.receptive_field / 2)) :, :] = -1

            # Get the spotting target
            clip_targets = getTimestampTargets(
                np.array([clip_labels]), self.num_detections
            )[0]

            return (
                torch.from_numpy(clip_feat),
                torch.from_numpy(clip_labels),
                torch.from_numpy(clip_targets),
            )
        else:
            # Load features
            feat_half1, feat_half2 = self.load_features(index=index)

            # Load labels
            label_half1, label_half2 = self.load_labels(feat_half1, feat_half2)

            # check if annotation exists
            if os.path.exists(
                os.path.join(
                    self.path, self.listGames[index].replace(" ", "_"), self.labels
                )
            ):
                labels = json.load(
                    open(
                        os.path.join(
                            self.path,
                            self.listGames[index].replace(" ", "_"),
                            self.labels,
                        )
                    )
                )

                for annotation in labels["annotations"]:

                    label, half, frame, cont = self.annotation(annotation)
                    if cont:
                        continue

                    value = 1
                    if "visibility" in annotation.keys():
                        if annotation["visibility"] == "not shown":
                            value = -1

                    if half == 1:
                        frame = min(frame, feat_half1.shape[0] - 1)
                        label_half1[frame][label] = value

                    if half == 2:
                        frame = min(frame, feat_half2.shape[0] - 1)
                        label_half2[frame][label] = value

            feat_half1 = feats2clip(
                torch.from_numpy(feat_half1),
                stride=self.chunk_size - self.receptive_field,
                clip_length=self.chunk_size,
                modif_last_index=True,
            )

            feat_half2 = feats2clip(
                torch.from_numpy(feat_half2),
                stride=self.chunk_size - self.receptive_field,
                clip_length=self.chunk_size,
                modif_last_index=True,
            )

            return (
                feat_half1,
                feat_half2,
                torch.from_numpy(label_half1),
                torch.from_numpy(label_half2),
            )

    def __len__(self):
        if self.train:
            return self.chunks_per_epoch
        else:
            return super().__len__()

    def load_features(self, index=0, game=""):
        super().load_features(index, game)
        return self.feat_half1, self.feat_half2

    def load_labels(self, feat_half1, feat_half2):
        super().load_labels(feat_half1, feat_half2, self.num_classes)
        return self.label_half1, self.label_half2

__getitem__(index)

Parameters:

Name	Type	Description	Default
index	int	Index	required

Returns: If training: clip_feat (np.array): clip of features. clip_labels (np.array): clip of labels. clip_targets (np.array): clip of targets. If testing: clip_feat_1 (np.array): clip of features for the first half. clip_feat_2 (np.array): clip of features for the second half. clip_labels_1 (np.array): clip of labels for the first half. clip_labels_2 (np.array): clip of labels for the second half.

Source code in opensportslib/datasets/localization_dataset.py

def __getitem__(self, index):
    """
    Args:
        index (int): Index
    Returns:
        If training:
            clip_feat (np.array): clip of features.
            clip_labels (np.array): clip of labels.
            clip_targets (np.array): clip of targets.
        If testing:
            clip_feat_1 (np.array): clip of features for the first half.
            clip_feat_2 (np.array): clip of features for the second half.
            clip_labels_1 (np.array): clip of labels for the first half.
            clip_labels_2 (np.array): clip of labels for the second half.

    """
    if self.train:
        # Retrieve the game index and the anchor
        class_selection = random.randint(0, self.num_classes)
        event_selection = random.randint(
            0, len(self.game_anchors[class_selection]) - 1
        )
        game_index = self.game_anchors[class_selection][event_selection][0]
        anchor = self.game_anchors[class_selection][event_selection][1]

        # Compute the shift for event chunks
        if class_selection < self.num_classes:
            shift = np.random.randint(
                -self.chunk_size + self.receptive_field, -self.receptive_field
            )
            start = anchor + shift
        # Compute the shift for non-event chunks
        else:
            start = random.randint(anchor[0], anchor[1] - self.chunk_size)
        if start < 0:
            start = 0
        if start + self.chunk_size >= self.game_feats[game_index].shape[0]:
            start = self.game_feats[game_index].shape[0] - self.chunk_size - 1

        # Extract the clips
        clip_feat = self.game_feats[game_index][start : start + self.chunk_size]
        clip_labels = self.game_labels[game_index][start : start + self.chunk_size]

        # Put loss to zero outside receptive field
        clip_labels[0 : int(np.ceil(self.receptive_field / 2)), :] = -1
        clip_labels[-int(np.ceil(self.receptive_field / 2)) :, :] = -1

        # Get the spotting target
        clip_targets = getTimestampTargets(
            np.array([clip_labels]), self.num_detections
        )[0]

        return (
            torch.from_numpy(clip_feat),
            torch.from_numpy(clip_labels),
            torch.from_numpy(clip_targets),
        )
    else:
        # Load features
        feat_half1, feat_half2 = self.load_features(index=index)

        # Load labels
        label_half1, label_half2 = self.load_labels(feat_half1, feat_half2)

        # check if annotation exists
        if os.path.exists(
            os.path.join(
                self.path, self.listGames[index].replace(" ", "_"), self.labels
            )
        ):
            labels = json.load(
                open(
                    os.path.join(
                        self.path,
                        self.listGames[index].replace(" ", "_"),
                        self.labels,
                    )
                )
            )

            for annotation in labels["annotations"]:

                label, half, frame, cont = self.annotation(annotation)
                if cont:
                    continue

                value = 1
                if "visibility" in annotation.keys():
                    if annotation["visibility"] == "not shown":
                        value = -1

                if half == 1:
                    frame = min(frame, feat_half1.shape[0] - 1)
                    label_half1[frame][label] = value

                if half == 2:
                    frame = min(frame, feat_half2.shape[0] - 1)
                    label_half2[frame][label] = value

        feat_half1 = feats2clip(
            torch.from_numpy(feat_half1),
            stride=self.chunk_size - self.receptive_field,
            clip_length=self.chunk_size,
            modif_last_index=True,
        )

        feat_half2 = feats2clip(
            torch.from_numpy(feat_half2),
            stride=self.chunk_size - self.receptive_field,
            clip_length=self.chunk_size,
            modif_last_index=True,
        )

        return (
            feat_half1,
            feat_half2,
            torch.from_numpy(label_half1),
            torch.from_numpy(label_half2),
        )

VQA datasets

VQADataset converts an OSL JSON VQA split into rows used by the configured VQA trainer and evaluator. It requires question/reference-answer data in the manifest; model-specific feature generation belongs to the VQA model/trainer, not this dataset route.

VQA dataset adapter for canonical OpenSportsLib JSON annotations.

VQADataset

Bases: Dataset

Flatten VQA annotations into single question-answer training samples.

Source code in opensportslib/datasets/vqa_dataset.py

class VQADataset(Dataset):
    """Flatten VQA annotations into single question-answer training samples."""

    def __init__(self, config, annotation_file: str | None = None, split: str = "train"):
        self.config = config
        self.split = split
        self._train_execution = self._as_dict(getattr(getattr(config, "TRAIN", None), "execution", None))
        self._prompt_cfg = self._as_dict(self._train_execution.get("prompt"))
        self._view_policy = str(self._train_execution.get("view_sampling_policy", "random_train_deterministic_eval")).lower()
        self._rng = random.Random(42)
        self.annotation_path = annotation_file or get_split_annotation_path(config, split)
        if not self.annotation_path:
            raise ValueError(
                f"Missing annotation path for split '{split}'. "
                f"Expected DATA.common.splits.{split}.annotation_path."
            )
        self.annotation_path = os.path.abspath(os.path.expanduser(self.annotation_path))
        with open(self.annotation_path, encoding="utf-8") as f:
            payload = json.load(f)

        data = payload.get("data", [])
        label_space = []
        payload_labels = payload.get("labels", {})
        if isinstance(payload_labels, dict):
            action_spec = payload_labels.get("action", {})
            if isinstance(action_spec, dict):
                label_space = [str(label).strip() for label in action_spec.get("labels", []) if str(label).strip()]
        source_root = get_split_source_path(config, split) or ""
        source_root = (
            os.path.abspath(os.path.expanduser(source_root))
            if source_root
            else os.path.dirname(self.annotation_path)
        )
        common = self._as_dict(self._as_dict(getattr(config, "DATA", None)).get("common"))
        feature_index_path = str(common.get("feature_index") or "").strip()
        prediction_index_path = str(common.get("prediction_index") or "").strip()
        backend = get_vqa_backend(config)
        self.backend = backend
        self.native_vl = backend == "qwen_vl_native_infer"
        feature_backend = str(self._train_execution.get("feature_backend", "xvars_clip")).lower()
        feature_source = get_vqa_feature_source(config, default="indexed")
        if not self.native_vl and feature_backend != "xvars_clip":
            raise ValueError(f"Unsupported VQA feature backend '{feature_backend}'. Expected 'xvars_clip'.")
        strict_feature_index = (not self.native_vl) and feature_source in {"indexed", ""}
        fallback_feature_index = (not self.native_vl) and feature_source in {"indexed_or_raw", "indexed_or_raw_clip"}
        if strict_feature_index and not feature_index_path:
            raise ValueError("Missing required config key DATA.common.feature_index for VQA xvars_clip mode.")
        self.feature_source = feature_source
        self.feature_mode = get_vqa_xvars_feature_mode(config, default="strict_xvars")
        self.expected_feature_tokens = get_xvars_train_video_token_len(config)
        if self.native_vl:
            # Native Qwen VL consumes raw video or frames and never uses X-VARS indexes.
            self.feature_index = {}
            self.prediction_index = {}
        else:
            self.feature_index = self._load_feature_index(
                feature_index_path,
                split=split,
                strict=strict_feature_index,
                allow_missing=fallback_feature_index,
            )
            self.prediction_index = (
                load_prediction_index(os.path.abspath(os.path.expanduser(prediction_index_path)), split=split)
                if prediction_index_path
                else {}
            )

        self.samples: list[dict[str, Any]] = []
        for item in data:
            item_id = item.get("id")
            item_id_str = str(item_id)
            labels = item.get("labels", {})
            metadata = item.get("metadata", {})
            allowed_labels = [str(label).strip() for label in item.get("allowed_labels", []) if str(label).strip()] or label_space
            video_path = self._resolve_video_path(item, source_root)
            frame_paths = self._resolve_frame_paths(item, source_root)

            feature_candidates = self.feature_index.get(item_id_str, [])
            if strict_feature_index and not feature_candidates:
                raise ValueError(
                    f"Missing feature index entry for sample id '{item_id_str}'. "
                    "Provide DATA.common.feature_index mapping with feature_paths or feature_dir/path."
                )
            pred_row = self.prediction_index.get(item_id_str, {})
            qa_rows = self._iter_qa_rows(item)
            for qa in qa_rows:
                question = qa["question"]
                refs = qa["references"]
                prior_prediction_text = build_prediction_prior_text(
                    pred_row,
                    adapter=self._prompt_cfg.get("prediction_prior_adapter"),
                    fields=self._prompt_cfg.get("prediction_prior_fields"),
                )
                self.samples.append(
                    {
                        "id": qa.get("sample_id", item_id),
                        "question": question,
                        "references": refs,
                        "video_path": video_path,
                        "frame_paths": frame_paths,
                        "video_frames": list(item.get("video_frames") or []),
                        "feature_candidates": list(feature_candidates),
                        "feature_source": self.feature_source,
                        "labels": labels,
                        "metadata": metadata,
                        "prediction": pred_row,
                        "prior_prediction_text": prior_prediction_text,
                        "ground_truth_label": qa["ground_truth_label"],
                        "allowed_labels": list(allowed_labels),
                    }
                )
        if self.native_vl:
            for sample in self.samples:
                if sample.get("video_path") or sample.get("frame_paths") or sample.get("video_frames"):
                    continue
                raise ValueError(
                    f"Native Qwen VL sample '{sample.get('id')}' is missing visual input. "
                    "Expected video_path, frame_paths, or video_frames."
                )

    def _load_feature_index(
        self,
        feature_index_path: str,
        *,
        split: str,
        strict: bool,
        allow_missing: bool,
    ) -> dict[str, list[str]]:
        if not feature_index_path:
            return {}

        resolved_path = os.path.abspath(os.path.expanduser(feature_index_path))
        if allow_missing and not os.path.exists(resolved_path):
            logger.warning(
                "VQA feature index unavailable; falling back to raw-video extraction | "
                "feature_source=%s | split=%s | feature_index=%s",
                self.feature_source,
                split,
                resolved_path,
            )
            return {}

        try:
            return load_feature_index(resolved_path, split=split)
        except Exception:
            if strict or not allow_missing:
                raise
            logger.warning(
                "VQA feature index unreadable; falling back to raw-video extraction | "
                "feature_source=%s | split=%s | feature_index=%s",
                self.feature_source,
                split,
                resolved_path,
                exc_info=True,
            )
            return {}

    def __len__(self) -> int:
        return len(self.samples)

    def __getitem__(self, idx: int) -> dict[str, Any]:
        sample = dict(self.samples[idx])
        if self.native_vl:
            sample["selected_feature_path"] = None
            sample["video_spatio_temporal_features"] = None
            return sample
        if self.feature_source == "raw_video":
            sample["selected_feature_path"] = None
            sample["video_spatio_temporal_features"] = None
            return sample
        feature_path = self._choose_feature_path(sample.get("feature_candidates") or [])
        if not feature_path or not os.path.exists(feature_path):
            if self.feature_source in {"raw_video", "auto", "indexed_or_raw", "indexed_or_raw_clip"}:
                sample["selected_feature_path"] = None
                sample["video_spatio_temporal_features"] = None
                return sample
            raise FileNotFoundError(
                f"Missing CLIP feature file for sample '{sample.get('id')}'. "
                f"Selected path: {feature_path!r}"
            )
        with open(feature_path, "rb") as f:
            raw = pickle.load(f)
        features = validate_xvars_feature_tensor(
            torch.as_tensor(raw, dtype=torch.float32),
            expected_tokens=self.expected_feature_tokens,
            context=f"X-VARS features for sample '{sample.get('id')}'",
        )
        sample["selected_feature_path"] = feature_path
        sample["video_spatio_temporal_features"] = features
        return sample

    def _choose_feature_path(self, candidates: list[str]) -> str:
        existing = [p for p in candidates if p and os.path.exists(p)]
        if not existing:
            return ""
        if self.split == "train" and self._view_policy == "random_train_deterministic_eval":
            return existing[self._rng.randint(0, len(existing) - 1)]
        return existing[0]

    @staticmethod
    def _resolve_video_path(item: dict[str, Any], source_root: str) -> str | None:
        direct_path = item.get("video_path")
        if direct_path:
            path = str(direct_path)
            return os.path.join(source_root, path) if source_root and not os.path.isabs(path) else path

        inputs = item.get("inputs", [])
        for inp in inputs:
            input_type = str(inp.get("type", "")).lower()
            if input_type not in {"video", "frames_npy", "frames"}:
                continue
            rel = inp.get("path")
            if rel:
                rel = str(rel)
                return os.path.join(source_root, rel) if source_root and not os.path.isabs(rel) else rel
        return None

    @staticmethod
    def _resolve_frame_paths(item: dict[str, Any], source_root: str) -> list[str]:
        frame_paths = item.get("frame_paths") or []
        if frame_paths:
            out = []
            for path in frame_paths:
                path = str(path)
                out.append(os.path.join(source_root, path) if source_root and not os.path.isabs(path) else path)
            return out

        inputs = item.get("inputs", [])
        for inp in inputs:
            input_type = str(inp.get("type", "")).lower()
            if input_type not in {"image", "images", "frame", "frames"}:
                continue
            path = inp.get("path")
            if path:
                path = str(path)
                return [os.path.join(source_root, path) if source_root and not os.path.isabs(path) else path]
            paths = inp.get("paths") or []
            out = []
            for subpath in paths:
                subpath = str(subpath)
                out.append(os.path.join(source_root, subpath) if source_root and not os.path.isabs(subpath) else subpath)
            if out:
                return out
        return []

    @staticmethod
    def _iter_qa_rows(item: dict[str, Any]) -> list[dict[str, Any]]:
        rows: list[dict[str, Any]] = []
        direct_question = str(item.get("question", "")).strip()
        if direct_question:
            refs = item.get("references", []) or []
            if isinstance(refs, str):
                refs = [refs]
            ground_truth_label = str(item.get("ground_truth_label", "")).strip()
            if not ground_truth_label and refs:
                ground_truth_label = str(refs[0]).strip()
            rows.append(
                {
                    "sample_id": item.get("id"),
                    "question": direct_question,
                    "references": [str(ref).strip() for ref in refs if str(ref).strip()],
                    "ground_truth_label": ground_truth_label,
                }
            )
            return rows

        direct_questions = item.get("questions", []) or []
        if isinstance(direct_questions, list) and direct_questions:
            for idx, question in enumerate(direct_questions):
                question = str(question).strip()
                if not question:
                    continue
                rows.append(
                    {
                        "sample_id": f"{item.get('id')}:{idx}",
                        "question": question,
                        "references": [],
                        "ground_truth_label": str(item.get("ground_truth_label", "")).strip(),
                    }
                )
            if rows:
                return rows

        for qa in item.get("answers", []):
            question = str(qa.get("question", "")).strip()
            refs = qa.get("answers", []) or []
            if isinstance(refs, str):
                refs = [refs]
            if not question:
                continue
            ground_truth_label = str(item.get("ground_truth_label", "")).strip()
            if not ground_truth_label and refs:
                ground_truth_label = str(refs[0]).strip()
            rows.append(
                {
                    "sample_id": item.get("id"),
                    "question": question,
                    "references": [str(ref).strip() for ref in refs if str(ref).strip()],
                    "ground_truth_label": ground_truth_label,
                }
            )
        return rows

    @staticmethod
    def _as_dict(obj: Any) -> dict[str, Any]:
        if obj is None:
            return {}
        if isinstance(obj, dict):
            return obj
        if hasattr(obj, "__dict__"):
            return vars(obj)
        return {}

Tracking and HDF5 utilities

Tracking utilities normalize positions, build graph edge indices, cache games, and apply spatial transforms. HDF5 utilities read timestamped tracking rows, normalize their features, discover games, and write manifests. write_h5_manifest writes a file; readers and cache builders perform local filesystem I/O.

tracking utilities for player-position graph data.

provides constants, edge-building strategies, spatial augmentations, and frame-level feature extraction for the SN-GAR tracking modality.

feature vector layout (per object, per frame)::

Index  Field
-----  ----------------------------------
0      x (pitch coordinates)
1      y 
2      is_ball (one-hot entity type)
3      is_home
4      is_away
5      dx (velocity delta, computed across consecutive frames)
6      dy 
7      z (ball height; MISSING_VALUE for players)

HorizontalFlip

randomly negate x-coordinates and dx-velocities (pitch length axis).

Parameters:

Name	Type	Description	Default
probability		chance of applying the flip per call.	0.5

Source code in opensportslib/datasets/utils/tracking.py

class HorizontalFlip:
    """randomly negate x-coordinates and dx-velocities (pitch length axis).

    Args:
        probability: chance of applying the flip per call.
    """

    def __init__(self, probability=0.5):
        self.probability = probability

    def __call__(self, features):
        """apply the transform.

        Args:
            features: numpy.ndarray of shape
                (num_frames, NUM_OBJECTS, FEATURE_DIM).

        Returns:
            the (possibly flipped) feature array. a copy is made when
            the flip is applied; the input is never modified.
        """
        if random.random() < self.probability:
            features_flipped = features.copy()
            valid_mask = features_flipped[:, :, 0] != MISSING_VALUE
            features_flipped[valid_mask, 0] *= -1  # flip x
            features_flipped[valid_mask, 5] *= -1  # flip dx
            return features_flipped
        return features

__call__(features)

apply the transform.

Parameters:

Name	Type	Description	Default
features		numpy.ndarray of shape (num_frames, NUM_OBJECTS, FEATURE_DIM).	required

Returns:

Type	Description
	the (possibly flipped) feature array. a copy is made when
	the flip is applied; the input is never modified.

Source code in opensportslib/datasets/utils/tracking.py

def __call__(self, features):
    """apply the transform.

    Args:
        features: numpy.ndarray of shape
            (num_frames, NUM_OBJECTS, FEATURE_DIM).

    Returns:
        the (possibly flipped) feature array. a copy is made when
        the flip is applied; the input is never modified.
    """
    if random.random() < self.probability:
        features_flipped = features.copy()
        valid_mask = features_flipped[:, :, 0] != MISSING_VALUE
        features_flipped[valid_mask, 0] *= -1  # flip x
        features_flipped[valid_mask, 5] *= -1  # flip dx
        return features_flipped
    return features

VerticalFlip

randomly negate y-coordinates and dy-velocities (pitch width axis).

Parameters:

Name	Type	Description	Default
probability		chance of applying the flip per call.	0.5

Source code in opensportslib/datasets/utils/tracking.py

class VerticalFlip:
    """randomly negate y-coordinates and dy-velocities (pitch width axis).

    Args:
        probability: chance of applying the flip per call.
    """

    def __init__(self, probability=0.5):
        self.probability = probability

    def __call__(self, features):
        """apply the transform.

        Args:
            features: numpy.ndarray of shape
                (num_frames, NUM_OBJECTS, FEATURE_DIM).

        Returns:
            the (possibly flipped) feature array.
        """
        if random.random() < self.probability:
            features_flipped = features.copy()
            valid_mask = features_flipped[:, :, 0] != MISSING_VALUE
            features_flipped[valid_mask, 1] *= -1  # flip y
            features_flipped[valid_mask, 6] *= -1  # flip dy
            return features_flipped
        return features

__call__(features)

apply the transform.

Parameters:

Name	Type	Description	Default
features		numpy.ndarray of shape (num_frames, NUM_OBJECTS, FEATURE_DIM).	required

Returns:

Type	Description
	the (possibly flipped) feature array.

Source code in opensportslib/datasets/utils/tracking.py

def __call__(self, features):
    """apply the transform.

    Args:
        features: numpy.ndarray of shape
            (num_frames, NUM_OBJECTS, FEATURE_DIM).

    Returns:
        the (possibly flipped) feature array.
    """
    if random.random() < self.probability:
        features_flipped = features.copy()
        valid_mask = features_flipped[:, :, 0] != MISSING_VALUE
        features_flipped[valid_mask, 1] *= -1  # flip y
        features_flipped[valid_mask, 6] *= -1  # flip dy
        return features_flipped
    return features

TeamFlip

randomly swap the home and away one-hot team flags.

this is a label-preserving augmentation: swapping team identity does not change the group activity class.

Parameters:

Name	Type	Description	Default
probability		chance of applying the swap per call.	0.5

Source code in opensportslib/datasets/utils/tracking.py

class TeamFlip:
    """randomly swap the home and away one-hot team flags.

    this is a label-preserving augmentation: swapping team identity
    does not change the group activity class.

    Args:
        probability: chance of applying the swap per call.
    """

    def __init__(self, probability=0.5):
        self.probability = probability

    def __call__(self, features):
        """apply the transform.

        Args:
            features: numpy.ndarray of shape
                (num_frames, NUM_OBJECTS, FEATURE_DIM).

        Returns:
            the (possibly swapped) feature array.
        """
        if random.random() < self.probability:
            features_flipped = features.copy()
            home_col = features_flipped[:, :, 3].copy()
            features_flipped[:, :, 3] = features_flipped[:, :, 4]
            features_flipped[:, :, 4] = home_col
            return features_flipped
        return features

__call__(features)

apply the transform.

Parameters:

Name	Type	Description	Default
features		numpy.ndarray of shape (num_frames, NUM_OBJECTS, FEATURE_DIM).	required

Returns:

Type	Description
	the (possibly swapped) feature array.

Source code in opensportslib/datasets/utils/tracking.py

def __call__(self, features):
    """apply the transform.

    Args:
        features: numpy.ndarray of shape
            (num_frames, NUM_OBJECTS, FEATURE_DIM).

    Returns:
        the (possibly swapped) feature array.
    """
    if random.random() < self.probability:
        features_flipped = features.copy()
        home_col = features_flipped[:, :, 3].copy()
        features_flipped[:, :, 3] = features_flipped[:, :, 4]
        features_flipped[:, :, 4] = home_col
        return features_flipped
    return features

parse_frame(row, num_objects=NUM_OBJECTS, feature_dim=FEATURE_DIM)

parse a single parquet row into per-object features and positions.

Each row is expected to contain JSON-encoded columns balls, homePlayers, and awayPlayers following the SN-GAR tracking format.

Parameters:

Name	Type	Description	Default
row		A single row (pandas.Series) from a tracking parquet file.	required

Returns:

Type	Description
	A tuple (features, positions) where features is a
	numpy.ndarray of shape (num_objects, feature_dim) and
	positions is a list of position-group strings (e.g.
	"GK", "DEF"). Unobserved slots are filled with
	MISSING_VALUE and an empty string respectively.

Source code in opensportslib/datasets/utils/tracking.py

def parse_frame(row, num_objects=NUM_OBJECTS, feature_dim=FEATURE_DIM):
    """parse a single parquet row into per-object features and positions.

    Each row is expected to contain JSON-encoded columns balls,
    homePlayers, and awayPlayers following the SN-GAR
    tracking format.

    Args:
        row: A single row (pandas.Series) from a tracking parquet
            file.

    Returns:
        A tuple (features, positions) where features is a
        numpy.ndarray of shape (num_objects, feature_dim) and
        positions is a list of position-group strings (e.g.
        "GK", "DEF").  Unobserved slots are filled with
        MISSING_VALUE and an empty string respectively.
    """
    features = np.full(
        (num_objects, feature_dim), MISSING_VALUE, dtype=np.float32,
    )
    positions = [""] * num_objects

    obj_idx = _BALL_SLOT

    # -- ball (always slot 0) --
    ball_str = row.get("balls", "null")
    if pd.notna(ball_str) and ball_str not in ("null", ""):
        try:
            ball_list = json.loads(ball_str)
            if ball_list:
                ball = ball_list[0]
                x, y = ball.get("x"), ball.get("y")
                z = ball.get("z", 0)
                if x is not None and y is not None:
                    features[obj_idx] = [
                        float(x), float(y),
                        1, 0, 0,          # one-hot: ball
                        0, 0,             # dx, dy (filled later)
                        float(z),
                    ]
                    positions[obj_idx] = "BALL"
        except (json.JSONDecodeError, TypeError):
            pass
    obj_idx = _HOME_SLOT_START

    # -- home players (slots 1-11) --
    home_str = row.get("homePlayers", "[]")
    if pd.notna(home_str) and home_str not in ("null", ""):
        try:
            home_players = json.loads(home_str)
            home_players = sorted(
                home_players,
                key=lambda p: int(p.get("jerseyNum", 0)),
            )[:11]

            for player in home_players:
                x, y = player.get("x"), player.get("y")
                if x is not None and y is not None:
                    features[obj_idx] = [
                        float(x), float(y),
                        0, 1, 0,          # one-hot: home
                        0, 0,             # dx, dy
                        MISSING_VALUE,    # z unused for players
                    ]
                    positions[obj_idx] = player.get("positionGroup", "")
                obj_idx += 1

            # advance past any unfilled home slots.
            while obj_idx < _HOME_SLOT_END:
                obj_idx += 1
        except (json.JSONDecodeError, TypeError):
            obj_idx = _HOME_SLOT_END
    else:
        obj_idx = _HOME_SLOT_END

    # -- away players (slots 12-22) --
    away_str = row.get("awayPlayers", "[]")
    if pd.notna(away_str) and away_str not in ("null", ""):
        try:
            away_players = json.loads(away_str)
            away_players = sorted(
                away_players,
                key=lambda p: int(p.get("jerseyNum", 0)),
            )[:11]

            for player in away_players:
                x, y = player.get("x"), player.get("y")
                if x is not None and y is not None:
                    features[obj_idx] = [
                        float(x), float(y),
                        0, 0, 1,          # one-hot: away
                        0, 0,             # dx, dy
                        MISSING_VALUE,    # z unused for players
                    ]
                    positions[obj_idx] = player.get("positionGroup", "")
                obj_idx += 1
        except (json.JSONDecodeError, TypeError):
            pass

    return features, positions

compute_deltas(all_features, num_objects=NUM_OBJECTS)

compute per-object velocity deltas across consecutive frames.

For each object that is observed in both frame t and frame t - 1, the displacement (dx, dy) is written into feature indices 5 and 6. frame 0 retains zero deltas.

Parameters:

Name	Type	Description	Default
all_features		numpy.ndarray of shape (num_frames, NUM_OBJECTS, FEATURE_DIM).	required

Returns:

Type	Description
	the same array, modified in-place, with velocity deltas
	populated.

Source code in opensportslib/datasets/utils/tracking.py

def compute_deltas(all_features, num_objects=NUM_OBJECTS):
    """compute per-object velocity deltas across consecutive frames.

    For each object that is observed in both frame t and frame t - 1,
    the displacement (dx, dy) is written into feature indices 5 and 6.
    frame 0 retains zero deltas.

    Args:
        all_features: numpy.ndarray of shape
            (num_frames, NUM_OBJECTS, FEATURE_DIM).

    Returns:
        the same array, modified in-place, with velocity deltas
        populated.
    """
    for t in range(1, all_features.shape[0]):
        for obj in range(num_objects):
            curr_valid = all_features[t, obj, 0] != MISSING_VALUE
            prev_valid = all_features[t - 1, obj, 0] != MISSING_VALUE
            if curr_valid and prev_valid:
                all_features[t, obj, 5] = (
                    all_features[t, obj, 0] - all_features[t - 1, obj, 0]
                )
                all_features[t, obj, 6] = (
                    all_features[t, obj, 1] - all_features[t - 1, obj, 1]
                )
    return all_features

normalize_features(features, pitch_half_length=PITCH_HALF_LENGTH, pitch_half_width=PITCH_HALF_WIDTH, max_displacement=MAX_DISPLACEMENT, max_ball_height=MAX_BALL_HEIGHT)

normalize spatial features to roughly [-1, 1].

observed coordinates are divided by the known pitch / displacement bounds. unobserved slots are set to MISSING_VALUE_NORMALIZED so the model receives a distinct, out-of-range sentinel that cannot be confused with a valid normalized value.

Parameters:

Name	Type	Description	Default
features		numpy.ndarray of shape (num_frames, NUM_OBJECTS, FEATURE_DIM).	required

Returns:

Type	Description
	a new array with the same shape containing normalized values.
	the input is not modified.

Source code in opensportslib/datasets/utils/tracking.py

def normalize_features(
    features,
    pitch_half_length=PITCH_HALF_LENGTH,
    pitch_half_width=PITCH_HALF_WIDTH,
    max_displacement=MAX_DISPLACEMENT,
    max_ball_height=MAX_BALL_HEIGHT,
):
    """normalize spatial features to roughly [-1, 1].

    observed coordinates are divided by the known pitch / displacement
    bounds.  unobserved slots are set to MISSING_VALUE_NORMALIZED
    so the model receives a distinct, out-of-range sentinel that
    cannot be confused with a valid normalized value.

    Args:
        features: numpy.ndarray of shape
            (num_frames, NUM_OBJECTS, FEATURE_DIM).

    Returns:
        a new array with the same shape containing normalized values.
        the input is not modified.
    """
    features_norm = features.copy()
    valid_mask = features_norm[:, :, 0] != MISSING_VALUE

    features_norm[valid_mask, 0] /= pitch_half_length
    features_norm[valid_mask, 1] /= pitch_half_width
    features_norm[valid_mask, 5] /= max_displacement
    features_norm[valid_mask, 6] /= max_displacement
    features_norm[valid_mask, 7] /= max_ball_height

    # write the normalized sentinel into every spatial channel of
    # missing objects so downstream layers can distinguish "absent"
    # from "at the origin".
    for ch in (0, 1, 5, 6, 7):
        features_norm[~valid_mask, ch] = MISSING_VALUE_NORMALIZED

    return features_norm

build_edge_index(node_features, node_positions, edge_type, k=8, r=15.0)

build a graph edge index for a single frame.

supports several connectivity strategies that trade off density against spatial or tactical priors.

Parameters:

Name	Type	Description	Default
node_features		numpy.ndarray of shape (NUM_OBJECTS, FEATURE_DIM).	required
node_positions		List of position-group strings (length NUM_OBJECTS).	required
edge_type		One of "none", "full", "knn", "distance", "ball_knn", "ball_distance", or "positional".	required
k		Number of neighbours for knn / ball_knn strategies.	8
r		Distance threshold (metres) for distance / ball_distance strategies.	15.0

Returns:

Type	Description
	numpy.ndarray of shape (2, num_edges) in COO format,
	compatible with PyTorch Geometric.

Source code in opensportslib/datasets/utils/tracking.py

def build_edge_index(node_features, node_positions, edge_type, k=8, r=15.0):
    """build a graph edge index for a single frame.

    supports several connectivity strategies that trade off density
    against spatial or tactical priors.

    Args:
        node_features: numpy.ndarray of shape
            (NUM_OBJECTS, FEATURE_DIM).
        node_positions: List of position-group strings (length
            NUM_OBJECTS).
        edge_type: One of "none", "full", "knn",
            "distance", "ball_knn", "ball_distance",
            or "positional".
        k: Number of neighbours for knn / ball_knn strategies.
        r: Distance threshold (metres) for distance / ball_distance
            strategies.

    Returns:
        numpy.ndarray of shape (2, num_edges) in COO format,
        compatible with PyTorch Geometric.
    """
    if edge_type == "none":
        return np.zeros((2, 0), dtype=np.int64)

    if edge_type == "full":
        return _build_full_edges(node_features)

    if edge_type == "knn":
        return _build_knn_edges(node_features, k)

    if edge_type == "distance":
        return _build_distance_edges(node_features, threshold=r)

    if edge_type == "ball_knn":
        return _build_ball_knn_edges(node_features, k)

    if edge_type == "ball_distance":
        return _build_ball_distance_edges(node_features, threshold=r)

    if edge_type == "positional":
        return _build_positional_edges(node_features, node_positions)

    raise ValueError(f"Unknown edge_type: {edge_type}")

build_game_cache(parquet_path)

parse a whole-match tracking parquet once and cache compact arrays.

writes three sidecar files next to parquet_path: .features.npy (num_frames, NUM_OBJECTS, FEATURE_DIM) float32 .times.npy (num_frames,) float64, videoTimeMs per row .positions.npy (num_frames, NUM_OBJECTS) int8 position codes

velocity deltas (dx, dy) are intentionally left at zero here rather than computed on this native-rate array: a training/eval window may subsample the native stream down to a lower extract_fps, and deltas must reflect the gap between sampled frames, not between raw native rows. callers should slice a window out of the cached array first, then run compute_deltas on that (already-decimated) subsequence. the array is not normalized and has no augmentation applied either - both also happen per-window at train/eval time so they stay config-dependent.

Parameters:

Name	Type	Description	Default
parquet_path		path to the whole-match tracking parquet.	required

Returns:

Type	Description
	a tuple (features, times, positions); see load_or_build_game_cache.

Source code in opensportslib/datasets/utils/tracking.py

def build_game_cache(parquet_path):
    """parse a whole-match tracking parquet once and cache compact arrays.

    writes three sidecar files next to `parquet_path`:
        <path>.features.npy   (num_frames, NUM_OBJECTS, FEATURE_DIM) float32
        <path>.times.npy      (num_frames,) float64, videoTimeMs per row
        <path>.positions.npy  (num_frames, NUM_OBJECTS) int8 position codes

    velocity deltas (dx, dy) are intentionally left at zero here rather
    than computed on this native-rate array: a training/eval window may
    subsample the native stream down to a lower extract_fps, and deltas
    must reflect the gap between *sampled* frames, not between raw
    native rows. callers should slice a window out of the cached array
    first, then run compute_deltas on that (already-decimated)
    subsequence. the array is not normalized and has no augmentation
    applied either - both also happen per-window at train/eval time so
    they stay config-dependent.

    Args:
        parquet_path: path to the whole-match tracking parquet.

    Returns:
        a tuple (features, times, positions); see load_or_build_game_cache.
    """
    df = pd.read_parquet(
        parquet_path,
        columns=["videoTimeMs", "balls", "homePlayers", "awayPlayers"],
    )

    num_frames = len(df)
    all_features = np.zeros((num_frames, NUM_OBJECTS, FEATURE_DIM), dtype=np.float32)
    all_positions = np.zeros((num_frames, NUM_OBJECTS), dtype=np.int8)

    for t, (_, row) in enumerate(df.iterrows()):
        features, positions = parse_frame(row)
        all_features[t] = features
        all_positions[t] = [_POSITION_CODES.get(p, 0) for p in positions]

    times = df["videoTimeMs"].to_numpy(dtype=np.float64)

    features_path, times_path, positions_path = _cache_paths(parquet_path)
    _atomic_save(features_path, all_features)
    _atomic_save(times_path, times)
    _atomic_save(positions_path, all_positions)

    return all_features, times, all_positions

load_or_build_game_cache(parquet_path, mmap=True)

load the cached per-frame features/times/positions for a game, building the cache (see build_game_cache) on first access.

Parameters:

Name	Type	Description	Default
parquet_path		path to the whole-match tracking parquet.	required
mmap		if True, load cached arrays memory-mapped (cheap random window slicing without pulling the whole game into RAM).	True

Returns:

Name	Type	Description
features		numpy.ndarray (num_frames, NUM_OBJECTS, FEATURE_DIM) float32, raw (no velocity deltas, unnormalized) - callers compute deltas after slicing/decimating a window.
times		numpy.ndarray (num_frames,) float64, videoTimeMs per row.
positions		numpy.ndarray (num_frames, NUM_OBJECTS) int8 position codes (see _POSITION_CODES / positions_to_strings).

Source code in opensportslib/datasets/utils/tracking.py

def load_or_build_game_cache(parquet_path, mmap=True):
    """load the cached per-frame features/times/positions for a game,
    building the cache (see build_game_cache) on first access.

    Args:
        parquet_path: path to the whole-match tracking parquet.
        mmap: if True, load cached arrays memory-mapped (cheap random
            window slicing without pulling the whole game into RAM).

    Returns:
        features: numpy.ndarray (num_frames, NUM_OBJECTS, FEATURE_DIM)
            float32, raw (no velocity deltas, unnormalized) - callers
            compute deltas after slicing/decimating a window.
        times: numpy.ndarray (num_frames,) float64, videoTimeMs per row.
        positions: numpy.ndarray (num_frames, NUM_OBJECTS) int8 position
            codes (see _POSITION_CODES / positions_to_strings).
    """
    cache_paths = _cache_paths(parquet_path)
    features_path, times_path, positions_path = cache_paths

    # Rebuild whenever the parquet is newer than its cache. Without this an
    # existence-only check silently serves stale features after the source
    # parquet is regenerated (e.g. create_action_spotting_dataset.py
    # --overwrite), which trains and evaluates on the previous dataset with
    # no visible error.
    if all(os.path.exists(p) for p in cache_paths):
        parquet_mtime = os.path.getmtime(parquet_path)
        if all(os.path.getmtime(p) >= parquet_mtime for p in cache_paths):
            mmap_mode = "r" if mmap else None
            features = np.load(features_path, mmap_mode=mmap_mode)
            times = np.load(times_path, mmap_mode=mmap_mode)
            positions = np.load(positions_path, mmap_mode=mmap_mode)
            return features, times, positions

    return build_game_cache(parquet_path)

positions_to_strings(position_codes)

decode a (NUM_OBJECTS,) array of int position codes back to strings.

Parameters:

Name	Type	Description	Default
position_codes		iterable of int codes (see _POSITION_CODES).	required

Returns:

Type	Description
	list of position-group strings, e.g. ["BALL", "GK", "DEF", ...].

Source code in opensportslib/datasets/utils/tracking.py

def positions_to_strings(position_codes):
    """decode a (NUM_OBJECTS,) array of int position codes back to strings.

    Args:
        position_codes: iterable of int codes (see _POSITION_CODES).

    Returns:
        list of position-group strings, e.g. ["BALL", "GK", "DEF", ...].
    """
    return [_POSITION_NAMES.get(int(c), "") for c in position_codes]

slice_window(features_cache, positions_cache, base_idx, clip_len, stride)

slice one decimated feature window out of a whole-match cache.

Picks clip_len frames starting at decimated index base_idx, taking every stride-th native row (native_idx = (base_idx + i) * stride). base_idx may be negative or run past the end of the match (start/end padding, as ActionSpotDataset does for video); out-of-range frames are left as MISSING_VALUE sentinels, which the edge-building / normalization utilities already treat as absent nodes. Velocity deltas are computed after this decimation (see compute_deltas) so dx/dy reflect the gap between sampled frames, not raw native rows.

Parameters:

Name	Type	Description	Default
features_cache		(native_num_frames, NUM_OBJECTS, FEATURE_DIM) array from load_or_build_game_cache (raw, no deltas, unnormalized).	required
positions_cache		(native_num_frames, NUM_OBJECTS) int8 position codes from the same cache.	required
base_idx		window start, in decimated (extract_fps) frame units.	required
clip_len		number of decimated frames in the window.	required
stride		native rows per decimated frame (from the game's metadata).	required

Returns:

Name	Type	Description
window		(clip_len, NUM_OBJECTS, FEATURE_DIM) float32 array with velocity deltas filled in.
positions		list of clip_len per-frame position-group string lists.

Source code in opensportslib/datasets/utils/tracking.py

def slice_window(features_cache, positions_cache, base_idx, clip_len, stride):
    """slice one decimated feature window out of a whole-match cache.

    Picks `clip_len` frames starting at decimated index `base_idx`, taking
    every `stride`-th native row (native_idx = (base_idx + i) * stride).
    `base_idx` may be negative or run past the end of the match (start/end
    padding, as ActionSpotDataset does for video); out-of-range frames are
    left as MISSING_VALUE sentinels, which the edge-building / normalization
    utilities already treat as absent nodes. Velocity deltas are computed
    *after* this decimation (see compute_deltas) so dx/dy reflect the gap
    between sampled frames, not raw native rows.

    Args:
        features_cache: (native_num_frames, NUM_OBJECTS, FEATURE_DIM) array
            from load_or_build_game_cache (raw, no deltas, unnormalized).
        positions_cache: (native_num_frames, NUM_OBJECTS) int8 position
            codes from the same cache.
        base_idx: window start, in decimated (extract_fps) frame units.
        clip_len: number of decimated frames in the window.
        stride: native rows per decimated frame (from the game's metadata).

    Returns:
        window: (clip_len, NUM_OBJECTS, FEATURE_DIM) float32 array with
            velocity deltas filled in.
        positions: list of `clip_len` per-frame position-group string lists.
    """
    native_num_frames = features_cache.shape[0]

    window = np.full(
        (clip_len, NUM_OBJECTS, FEATURE_DIM), MISSING_VALUE, dtype=np.float32
    )
    positions = [[""] * NUM_OBJECTS for _ in range(clip_len)]

    for i in range(clip_len):
        native_idx = (base_idx + i) * stride
        if 0 <= native_idx < native_num_frames:
            window[i] = features_cache[native_idx]
            positions[i] = positions_to_strings(positions_cache[native_idx])

    window = compute_deltas(window)
    return window, positions

Utilities for UTC-indexed H5 player and ball tracking inputs.

H5Frame dataclass

Source code in opensportslib/datasets/utils/h5_tracking.py

@dataclass(frozen=True)
class H5Frame:
    timestamp: np.datetime64
    features: np.ndarray
    entity_ids: list[str]
    positions: list[str]

H5TrackingReader

Lazy reader for columnar player tracking H5 files keyed by UTC.

Source code in opensportslib/datasets/utils/h5_tracking.py

class H5TrackingReader:
    """Lazy reader for columnar player tracking H5 files keyed by UTC."""

    def __init__(
        self,
        player_path: str | Path,
        modality: str,
        ball_path: str | Path | None = None,
        *,
        missing_value: float = MISSING_VALUE,
        timestamp_tolerance_ms: float = 50.0,
    ):
        self.player_path = str(player_path)
        self.ball_path = str(ball_path) if ball_path else None
        self.modality = str(modality).lower()
        self.missing_value = float(missing_value)
        self.timestamp_tolerance = np.timedelta64(int(timestamp_tolerance_ms * 1000), "us")

        (
            self.player_timestamps,
            self.player_index,
            self.player_unique_timestamps,
        ) = _build_timestamp_index(self.player_path)

        self.ball_timestamps = None
        self.ball_index = {}
        self.ball_unique_timestamps = np.array([], dtype="datetime64[us]")
        if self.ball_path is not None:
            (
                self.ball_timestamps,
                self.ball_index,
                self.ball_unique_timestamps,
            ) = _build_timestamp_index(self.ball_path)

    def select_timestamps(
        self,
        *,
        start_utc=None,
        end_utc=None,
        num_frames: int | None = None,
        stride_ms: float | None = None,
        target_fps: float | None = None,
    ) -> np.ndarray:
        if self.player_unique_timestamps.size == 0:
            raise ValueError(f"Player H5 input has no timestamps: {self.player_path}")
        start = parse_utc(start_utc) if start_utc else self.player_unique_timestamps[0]
        end = parse_utc(end_utc) if end_utc else self.player_unique_timestamps[-1]

        if self.ball_path and self.ball_unique_timestamps.size:
            if start_utc is None:
                start = max(start, self.ball_unique_timestamps[0])
            if end_utc is None:
                end = min(end, self.ball_unique_timestamps[-1])

        candidates = self.player_unique_timestamps[
            (self.player_unique_timestamps >= start) & (self.player_unique_timestamps <= end)
        ]
        if candidates.size == 0:
            raise ValueError(
                f"No player H5 timestamps found in requested UTC range: {start} to {end}"
            )

        stride_us = None
        if stride_ms:
            stride_us = int(float(stride_ms) * 1000)
        elif target_fps:
            stride_us = int(1_000_000 / float(target_fps))

        if stride_us and stride_us > 0:
            selected = []
            next_allowed = candidates[0]
            stride = np.timedelta64(stride_us, "us")
            for timestamp in candidates:
                if timestamp >= next_allowed:
                    selected.append(timestamp)
                    next_allowed = timestamp + stride
            candidates = np.asarray(selected, dtype="datetime64[us]")

        if num_frames is not None and int(num_frames) > 0 and candidates.size > int(num_frames):
            if stride_us:
                candidates = candidates[: int(num_frames)]
            else:
                indices = np.linspace(0, candidates.size - 1, int(num_frames))
                candidates = candidates[np.round(indices).astype(np.int64)]

        return candidates

    def read_sequence(self, timestamps: Iterable[np.datetime64]) -> list[H5Frame]:
        frames = [self._read_player_frame(timestamp) for timestamp in timestamps]
        if self.ball_path is not None:
            frames = [self._append_ball(frame) for frame in frames]
        return _compute_deltas(frames, missing_value=self.missing_value)

    def _row_indices(self, timestamp: np.datetime64) -> np.ndarray:
        key = timestamp.astype("datetime64[us]").astype("int64").item()
        return self.player_index.get(key, np.asarray([], dtype=np.int64))

    def _read_player_frame(self, timestamp: np.datetime64) -> H5Frame:
        indices = self._row_indices(timestamp)
        if indices.size == 0:
            return H5Frame(timestamp, np.empty((0, H5_FEATURE_DIM), dtype=np.float32), [], [])

        with h5py.File(self.player_path, "r") as f:
            if self.modality == "player_joints_h5":
                features, entity_ids, positions = self._read_joint_nodes(f, indices)
            else:
                features, entity_ids, positions = self._read_centroid_nodes(f, indices)
        return H5Frame(timestamp, features, entity_ids, positions)

    def _read_centroid_nodes(self, f: h5py.File, indices: np.ndarray):
        features = np.full((len(indices), H5_FEATURE_DIM), self.missing_value, dtype=np.float32)
        features[:, 3] = 0.0
        is_home_values = _read_column(f, "is_home", indices, default=-1)

        features[:, 0] = _read_column(f, "x", indices, default=self.missing_value).astype(np.float32)
        features[:, 1] = _read_column(f, "y", indices, default=self.missing_value).astype(np.float32)
        features[:, 2] = self.missing_value
        features[:, 4] = np.asarray(is_home_values == 1, dtype=np.float32)
        features[:, 5] = np.asarray(is_home_values == 0, dtype=np.float32)
        features[:, 6:9] = 0.0
        features[:, 9] = 0.0

        player_ids = _read_column(f, "player_id", indices, default="")
        jersey = _read_column(f, "jersey_number", indices, default="")
        entity_ids = []
        for row_offset, (player_id, jersey_number, is_home) in enumerate(
            zip(player_ids, jersey, is_home_values)
        ):
            player_id = _decode_scalar(player_id)
            jersey_number = _decode_scalar(jersey_number)
            side = "home" if is_home == 1 else "away" if is_home == 0 else "unknown"
            entity_ids.append(player_id or f"{side}:{jersey_number}:{row_offset}")
        positions = [_decode_scalar(v) for v in _read_column(f, "role_name", indices, default="")]
        return features, entity_ids, positions

    def _read_joint_nodes(self, f: h5py.File, indices: np.ndarray):
        joint_names = sorted(
            key[:-2]
            for key in f.keys()
            if key.endswith("_x") and f"{key[:-2]}_y" in f and f"{key[:-2]}_z" in f
        )
        is_home_values = _read_column(f, "is_home", indices, default=-1)
        player_ids = _read_column(f, "player_id", indices, default="")
        jersey = _read_column(f, "jersey_number", indices, default="")
        role_names = _read_column(f, "role_name", indices, default="")

        node_count = len(indices) * len(joint_names)
        features = np.full((node_count, H5_FEATURE_DIM), self.missing_value, dtype=np.float32)
        entity_ids = []
        positions = []
        out_idx = 0
        for row_offset, row_idx in enumerate(indices):
            player_id = _decode_scalar(player_ids[row_offset])
            jersey_number = _decode_scalar(jersey[row_offset])
            is_home = is_home_values[row_offset]
            side = "home" if is_home == 1 else "away" if is_home == 0 else "unknown"
            player_key = player_id or f"{side}:{jersey_number}:{row_offset}"
            for joint_name in joint_names:
                features[out_idx, 0] = f[f"{joint_name}_x"][row_idx]
                features[out_idx, 1] = f[f"{joint_name}_y"][row_idx]
                features[out_idx, 2] = f[f"{joint_name}_z"][row_idx]
                features[out_idx, 3] = 0.0
                features[out_idx, 4] = 1.0 if is_home == 1 else 0.0
                features[out_idx, 5] = 1.0 if is_home == 0 else 0.0
                features[out_idx, 6:9] = 0.0
                features[out_idx, 9] = 1.0
                entity_ids.append(f"{player_key}:{joint_name}")
                positions.append(_decode_scalar(role_names[row_offset]))
                out_idx += 1
        return features, entity_ids, positions

    def _append_ball(self, frame: H5Frame) -> H5Frame:
        ball_features = np.full((1, H5_FEATURE_DIM), self.missing_value, dtype=np.float32)
        ball_features[0, 3] = 1.0
        ball_features[0, 4:6] = 0.0
        ball_features[0, 6:9] = 0.0
        ball_features[0, 9] = 0.0

        nearest = self._nearest_ball_timestamp(frame.timestamp)
        if nearest is not None:
            key = nearest.astype("datetime64[us]").astype("int64").item()
            row_idx = self.ball_index[key][0]
            with h5py.File(self.ball_path, "r") as f:
                ball_features[0, 0] = f["x"][row_idx] if "x" in f else self.missing_value
                ball_features[0, 1] = f["y"][row_idx] if "y" in f else self.missing_value
                ball_features[0, 2] = f["z"][row_idx] if "z" in f else self.missing_value

        features = np.concatenate([ball_features, frame.features], axis=0)
        return H5Frame(
            timestamp=frame.timestamp,
            features=features,
            entity_ids=["BALL"] + frame.entity_ids,
            positions=["BALL"] + frame.positions,
        )

    def _nearest_ball_timestamp(self, timestamp: np.datetime64) -> np.datetime64 | None:
        if self.ball_unique_timestamps.size == 0:
            return None
        pos = np.searchsorted(self.ball_unique_timestamps, timestamp)
        candidates = []
        if pos < self.ball_unique_timestamps.size:
            candidates.append(self.ball_unique_timestamps[pos])
        if pos > 0:
            candidates.append(self.ball_unique_timestamps[pos - 1])
        if not candidates:
            return None
        nearest = min(candidates, key=lambda candidate: abs(candidate - timestamp))
        if abs(nearest - timestamp) <= self.timestamp_tolerance:
            return nearest
        return None

parse_utc(value)

Parse UTC-ish H5/JSON timestamp values into microsecond datetime64.

Source code in opensportslib/datasets/utils/h5_tracking.py

def parse_utc(value) -> np.datetime64:
    """Parse UTC-ish H5/JSON timestamp values into microsecond datetime64."""
    if isinstance(value, np.datetime64):
        return value.astype("datetime64[us]")
    if isinstance(value, (bytes, np.bytes_)):
        value = value.decode("utf-8")
    value = str(value).strip()
    if not value:
        raise ValueError("Empty UTC timestamp")
    if value.endswith("Z"):
        value = value[:-1] + "+00:00"
    dt = datetime.fromisoformat(value)
    if dt.tzinfo is not None:
        dt = dt.astimezone(timezone.utc).replace(tzinfo=None)
    return np.datetime64(dt, "us")

normalize_h5_features(features, *, pitch_half_length=85.0, pitch_half_width=50.0, max_height=30.0, max_displacement=110.0, missing_value=MISSING_VALUE)

Normalize H5 XYZ graph features without mutating the input.

Source code in opensportslib/datasets/utils/h5_tracking.py

def normalize_h5_features(
    features: np.ndarray,
    *,
    pitch_half_length: float = 85.0,
    pitch_half_width: float = 50.0,
    max_height: float = 30.0,
    max_displacement: float = 110.0,
    missing_value: float = MISSING_VALUE,
) -> np.ndarray:
    """Normalize H5 XYZ graph features without mutating the input."""
    out = features.copy()
    valid_mask = out[:, :, 0] != missing_value
    out[valid_mask, 0] /= pitch_half_length
    out[valid_mask, 1] /= pitch_half_width
    out[valid_mask, 2] /= max_height
    out[valid_mask, 6:9] /= max_displacement
    for ch in (0, 1, 2, 6, 7, 8):
        out[~valid_mask, ch] = MISSING_VALUE_NORMALIZED
    return out

find_h5_games(path, joints_filename=JOINTS_FILENAME, ball_filename=BALL_FILENAME)

List game directories holding both tracking files.

path is either one game directory or a directory of game directories.

Source code in opensportslib/datasets/utils/h5_tracking.py

def find_h5_games(
    path,
    joints_filename: str = JOINTS_FILENAME,
    ball_filename: str = BALL_FILENAME,
) -> list[Path]:
    """List game directories holding both tracking files.

    `path` is either one game directory or a directory of game directories.
    """
    def holds_both(directory: Path) -> bool:
        try:
            return (directory / joints_filename).exists() and (directory / ball_filename).exists()
        except OSError:
            return False

    path = Path(path).resolve()
    if holds_both(path):
        return [path]
    return sorted(d for d in path.iterdir() if d.is_dir() and holds_both(d))

write_h5_manifest(path, games, label='header', head_name='action', joints_filename=JOINTS_FILENAME, ball_filename=BALL_FILENAME)

Write the OSL JSON manifest a rule-based H5 spotter reads its inputs from.

Input paths are absolute, so a config published without the data beside it still resolves.

Source code in opensportslib/datasets/utils/h5_tracking.py

def write_h5_manifest(
    path,
    games: Iterable[Path],
    label: str = "header",
    head_name: str = "action",
    joints_filename: str = JOINTS_FILENAME,
    ball_filename: str = BALL_FILENAME,
) -> Path:
    """Write the OSL JSON manifest a rule-based H5 spotter reads its inputs from.

    Input paths are absolute, so a config published without the data beside it
    still resolves.
    """
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps({
        "version": "2.0",
        "task": "action_spotting",
        "labels": {head_name: {"type": "single_label", "labels": [label]}},
        "data": [
            {
                "id": game.name,
                "inputs": [{
                    "type": "player_joints_h5",
                    "path": str(game / joints_filename),
                    "ball_path": str(game / ball_filename),
                }],
            }
            for game in games
        ],
    }, indent=2))
    return path

Hugging Face-backed inputs

Hub preparation is optional. It can download or materialize local files and therefore requires a configured source, network access, and authentication for gated repositories.

Stage OSL JSON manifests and their selected media from a Hugging Face dataset.

PreparedHFJsonSplit dataclass

Source code in opensportslib/datasets/hf_json.py

@dataclass(frozen=True)
class PreparedHFJsonSplit:
    annotation_path: Path
    source_path: Path
    revision: str

hf_json_source(config)

Source code in opensportslib/datasets/hf_json.py

def hf_json_source(config) -> dict:
    source = get_input_cfg(config).get("source") or {}
    return source if str(source.get("format", "")).lower() == "hf_json" else {}

prepare_hf_json_split(config, split)

Download one split's selected media, then publish its local manifest.

The localization video and tracking loaders inspect each file during construction. Stage only the configured input type while keeping its manifest type and other metadata intact. Hub blobs are linked into a small split directory.

Source code in opensportslib/datasets/hf_json.py

def prepare_hf_json_split(config, split: str) -> PreparedHFJsonSplit:
    """Download one split's selected media, then publish its local manifest.

    The localization video and tracking loaders inspect each file during
    construction. Stage only the configured input type while keeping its
    manifest type and other metadata intact. Hub blobs are linked into a
    small split directory.
    """
    source = hf_json_source(config)
    if not source:
        raise ValueError("Source format must be 'hf_json'.")
    if split not in {"train", "valid", "test"}:
        raise ValueError(f"Unsupported Hugging Face JSON split: {split!r}")

    repo_id = str(source.get("repo_id") or "").strip()
    branch = str(source.get("revision") or "main").strip()
    cache_dir = Path(str(source.get("cache_dir") or "~/.cache/opensportslib/hf_json")).expanduser()
    media_type = str(source.get("input_type") or "video_mp4")
    if not repo_id or not branch:
        raise ValueError("hf_json requires source.repo_id and source.revision.")

    try:
        from huggingface_hub import HfApi, hf_hub_download

        revision = str(source.get("resolved_revision") or "").strip()
        if not revision:
            revision = HfApi(token=True).repo_info(
                repo_id=repo_id, repo_type="dataset", revision=branch
            ).sha
        if not revision:
            raise ValueError("The Hub did not return a commit SHA.")
    except Exception as exc:
        raise RuntimeError(
            f"Cannot access Hugging Face dataset {repo_id!r} at {branch!r}. "
            "Request access and run `hf auth login`."
        ) from exc

    _set_resolved_revision(config, revision)
    safe_repo = re.sub(r"[^A-Za-z0-9_.-]+", "--", repo_id)
    root = cache_dir / safe_repo / revision
    root.mkdir(parents=True, exist_ok=True)
    manifest_name = _safe_path(str(source.get("annotation_pattern") or "annotations_{split}.json").format(split=split))
    annotation_path = root / f"selected_{media_type}_{split}.json"
    prepared = PreparedHFJsonSplit(annotation_path, root, revision)
    hub_cache = cache_dir / "hub"

    with (root / f".{split}.lock").open("w") as lock:
        fcntl.flock(lock, fcntl.LOCK_EX)
        try:
            cached_manifest = annotation_path.is_file()
            if cached_manifest:
                payload = json.loads(annotation_path.read_text(encoding="utf-8"))
                # Older caches could contain an input relabeled to ``video``.
                # Refresh those manifests so the selected source type is kept.
                valid_cache = isinstance(payload, dict) and isinstance(payload.get("data"), list)
                if valid_cache:
                    valid_cache = all(
                        isinstance(sample, dict)
                        and len([
                            item for item in sample.get("inputs", [])
                            if isinstance(item, dict) and item.get("type") == media_type
                        ]) == 1
                        for sample in payload["data"]
                    )
                if not valid_cache:
                    cached_manifest = False
            if not cached_manifest:
                remote_manifest = hf_hub_download(
                    repo_id=repo_id, filename=manifest_name, repo_type="dataset",
                    revision=revision, token=True, cache_dir=str(hub_cache),
                )
                payload = json.loads(Path(remote_manifest).read_text(encoding="utf-8"))
            if not isinstance(payload, dict) or not isinstance(payload.get("data"), list):
                raise ValueError(f"{manifest_name} is not an OSL JSON manifest")
            if not payload["data"]:
                raise ValueError(f"{manifest_name} contains no samples")

            filenames = set()
            for sample in payload["data"]:
                selected = [item for item in sample.get("inputs", []) if item.get("type") == media_type]
                if len(selected) != 1:
                    raise ValueError(
                        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)
                sample["inputs"] = [item]

            def stage_media(filename: str) -> None:
                destination = root / filename
                if destination.is_file():
                    return
                downloaded = Path(hf_hub_download(
                    repo_id=repo_id, filename=filename, repo_type="dataset",
                    revision=revision, token=True, cache_dir=str(hub_cache),
                ))
                destination.parent.mkdir(parents=True, exist_ok=True)
                temporary = destination.with_name(f".{destination.name}.{os.getpid()}.tmp")
                try:
                    temporary.symlink_to(downloaded)
                    os.replace(temporary, destination)
                finally:
                    temporary.unlink(missing_ok=True)

            with ThreadPoolExecutor(max_workers=min(4, len(filenames))) as executor:
                list(executor.map(stage_media, sorted(filenames)))

            if not cached_manifest:
                temporary = annotation_path.with_name(f".{annotation_path.name}.{os.getpid()}.tmp")
                try:
                    temporary.write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8")
                    os.replace(temporary, annotation_path)
                finally:
                    temporary.unlink(missing_ok=True)
        except Exception as exc:
            raise RuntimeError(
                f"Failed to stage {split!r} from {repo_id}@{revision}: {exc}"
            ) from exc
    return prepared

Metadata and TAR staging for Hugging Face tracking datasets.

PreparedTrackingSplit dataclass

Source code in opensportslib/datasets/hf_tracking.py

@dataclass(frozen=True)
class PreparedTrackingSplit:
    annotations_path: Path
    index_path: Path
    repo_id: str
    revision: str
    split: str
    cache_dir: Path

HFTarTrackingReader

Read one Parquet member without extracting it to the filesystem.

Source code in opensportslib/datasets/hf_tracking.py

class HFTarTrackingReader:
    """Read one Parquet member without extracting it to the filesystem."""

    def __init__(self, prepared: PreparedTrackingSplit):
        self.prepared = prepared
        with prepared.index_path.open(encoding="utf-8") as handle:
            index = json.load(handle)
        self.members = index["members"]
        # SN-GAR training samples randomly across 13 shards. An eight-handle
        # cache repeatedly evicts and reindexes TARs (about 0.18 s per reopen).
        # Keep all of its shards open, while bounding descriptors for larger sets.
        self._max_open_tars = min(32, max(8, len({item[0] for item in self.members.values()})))
        self._pid = os.getpid()
        self._shard_paths = index.get("shard_paths", {}).copy()
        self._open_tars = OrderedDict()

    def __getstate__(self):
        state = self.__dict__.copy()
        state["_pid"] = None
        state["_open_tars"] = OrderedDict()
        return state

    def __del__(self):
        for tar in getattr(self, "_open_tars", {}).values():
            tar.close()

    def _reset_after_fork(self):
        if self._pid != os.getpid():
            for tar in self._open_tars.values():
                tar.close()
            self._pid = os.getpid()
            self._open_tars = OrderedDict()

    def _tar(self, shard_name: str):
        self._reset_after_fork()
        if shard_name in self._open_tars:
            self._open_tars.move_to_end(shard_name)
            return self._open_tars[shard_name]
        if shard_name not in self._shard_paths:
            from huggingface_hub import hf_hub_download

            try:
                logger.info("Caching tracking TAR shard %s/%s", self.prepared.split, shard_name)
                self._shard_paths[shard_name] = hf_hub_download(
                    repo_id=self.prepared.repo_id,
                    repo_type="dataset",
                    revision=self.prepared.revision,
                    filename=f"{self.prepared.split}/shards/{shard_name}",
                    cache_dir=str(self.prepared.cache_dir / "hub"),
                    token=True,
                )
            except Exception as exc:
                raise RuntimeError(f"Cannot download tracking TAR shard {shard_name!r}") from exc
        tar = tarfile.open(self._shard_paths[shard_name], mode="r:")
        self._open_tars[shard_name] = tar
        if len(self._open_tars) > self._max_open_tars:
            self._open_tars.popitem(last=False)[1].close()
        return tar

    def read_parquet(self, path: str):
        import pandas as pd

        if path not in self.members:
            raise FileNotFoundError(f"Tracking clip is absent from Hub manifest: {path}")
        shard_name, member_name = self.members[path]
        tar = self._tar(shard_name)
        try:
            member = tar.getmember(member_name)
            stream = tar.extractfile(member)
            if stream is None:
                raise KeyError(member_name)
            with stream:
                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))

hf_tracking_source(config)

Source code in opensportslib/datasets/hf_tracking.py

def hf_tracking_source(config) -> dict:
    source = get_input_cfg(config).get("source") or {}
    return source if str(source.get("format", "")).lower() == "hf_webdataset" else {}

prepare_hf_tracking_split(config, split)

Index metadata and cache all required TAR shards for one split.

Source code in opensportslib/datasets/hf_tracking.py

def prepare_hf_tracking_split(config, split: str) -> PreparedTrackingSplit:
    """Index metadata and cache all required TAR shards for one split."""
    source = hf_tracking_source(config)
    if not source:
        raise ValueError("Tracking source format must be 'hf_webdataset'.")
    if split not in {"train", "valid", "test"}:
        raise ValueError(f"Unsupported Hugging Face tracking split: {split!r}")

    repo_id = str(source.get("repo_id") or "").strip()
    revision = str(source.get("revision") or "tracking").strip()
    cache_dir = Path(str(source.get("cache_dir") or "~/.cache/opensportslib/hf_tracking")).expanduser()
    if not repo_id or not revision:
        raise ValueError("Hugging Face tracking source requires repo_id and revision.")

    try:
        from huggingface_hub import HfApi

        resolved = str(source.get("resolved_revision") or "").strip()
        if not resolved:
            resolved = HfApi(token=True).repo_info(
                repo_id=repo_id, revision=revision, repo_type="dataset"
            ).sha
        if not resolved:
            raise ValueError("The Hub did not return a commit SHA.")
    except Exception as exc:
        raise RuntimeError(
            f"Cannot access gated dataset {repo_id!r} at {revision!r}. "
            "Request access and run `hf auth login`."
        ) from exc

    _set_resolved_revision(config, resolved)
    safe_repo = re.sub(r"[^A-Za-z0-9_.-]+", "--", repo_id)
    directory = cache_dir / safe_repo / resolved / split
    directory.mkdir(parents=True, exist_ok=True)
    annotations_path = directory / "annotations.json"
    index_path = directory / "index.json"
    prepared = PreparedTrackingSplit(
        annotations_path, index_path, repo_id, resolved, split, cache_dir
    )

    with (directory / ".prepare.lock").open("w") as lock:
        fcntl.flock(lock, fcntl.LOCK_EX)
        if annotations_path.is_file() and index_path.is_file():
            _stage_shards(prepared)
            return prepared
        logger.info("Staging %s tracking metadata from %s@%s", split, repo_id, resolved)

        try:
            from datasets import load_dataset
        except ImportError as exc:
            raise ImportError(
                    "Hugging Face tracking requires `datasets`; install with "
                    "`python -m pip install -e .`."
            ) from exc

        def rows(filename: str):
            uri = f"hf://datasets/{repo_id}@{resolved}/{split}/{filename}"
            return load_dataset(
                "parquet",
                data_files={split: uri},
                split=split,
                streaming=True,
                token=True,
                cache_dir=str(cache_dir / "datasets"),
            )

        try:
            samples = []
            locations = {}
            header = None
            for row in rows("metadata.parquet"):
                if header is None:
                    header = json.loads(row["header"])
                sample = json.loads(row["sample_payload"])
                if not isinstance(sample, dict):
                    raise ValueError("metadata.parquet contains a non-object sample_payload")
                sample_id = str(row["sample_id"])
                if sample_id in locations:
                    raise ValueError(f"Duplicate sample ID in metadata: {sample_id}")
                locations[sample_id] = (str(row["shard_name"]), sample)
                samples.append(sample)

            if not samples or not isinstance(header, dict):
                raise ValueError(f"{split}/metadata.parquet has no samples or header")

            manifest = {}
            for row in rows("shard_manifest.parquet"):
                if row.get("file_role") not in (None, "", "primary") or row.get("input_type") != "tracking_parquet":
                    continue
                key = (str(row["sample_id"]), str(row["relative_path"]))
                if row.get("status") == "ok" and row.get("wds_member"):
                    manifest[key] = (str(row["shard_name"]), str(row["wds_member"]))

            members = {}
            for sample_id, (shard_name, sample) in locations.items():
                for input_item in sample.get("inputs", []):
                    if input_item.get("type") != "tracking_parquet":
                        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
            if not members:
                raise ValueError(f"{split} contains no tracking_parquet inputs")

            _write_json_atomic(annotations_path, {**header, "data": samples})
            _write_json_atomic(index_path, {"members": members})
            logger.info(
                "Indexed %d %s samples across %d TAR shards",
                len(samples), split, len({shard for shard, _ in members.values()}),
            )
        except Exception as exc:
            raise RuntimeError(
                f"Could not stage {split} tracking metadata from {repo_id}@{resolved}: {exc}"
            ) from exc
        _stage_shards(prepared)

    return prepared