pvio.torch_tools — PyTorch Integration

Dataset and DataLoader for efficient parallel frame loading with PyTorch.

pvio.torch_tools.VideoCollectionDataset

VideoCollectionDataset(
    videos: list[Video],
    *,
    transform: Callable | None = None,
    use_cached_video_metadata: bool = True,
    n_frame_counting_workers: int = -1,
    progress_bar: bool | None = None,
)

Bases: IterableDataset

Iterable dataset that yields frames from a list of Video objects.

Each frame is yielded as a CHW float32 tensor with values in [0, 1]. Call .assign_workers() (done automatically by VideoCollectionDataLoader) before iterating.

Parameters:

Name	Type	Description	Default
videos	list[Video]	Video objects to iterate. Videos are set up here; do not call .setup() on them beforehand.	required
transform	Callable | None	Applied to each tensor before yielding. The following are already applied before the transform: (i) numpy array → torch tensor, (ii) HWC → CHW layout, and (iii) uint8 [0, 255] → float32 [0, 1].	None
use_cached_video_metadata	bool	Use cached metadata for EncodedVideo objects if available. Set to False to force re-reading.	True
n_frame_counting_workers	int	Parallel workers for pre-loading EncodedVideo metadata. -1 uses all available cores.	-1
progress_bar	bool | None	Show a progress bar during metadata loading. If None, shows one only when stderr is a TTY.	None

Source code in src/pvio/torch_tools.py

def __init__(
    self,
    videos: list[Video],
    *,
    transform: Callable | None = None,
    use_cached_video_metadata: bool = True,
    n_frame_counting_workers: int = -1,
    progress_bar: bool | None = None,
) -> None:
    r"""Iterable dataset that yields frames from a list of Video objects.

    Each frame is yielded as a CHW float32 tensor with values in [0, 1].
    Call `.assign_workers()` (done automatically by VideoCollectionDataLoader)
    before iterating.

    Args:
        videos (list[Video]): Video objects to iterate. Videos are set up here;
            do not call `.setup()` on them beforehand.
        transform (Callable | None): Applied to each tensor before yielding.
            The following are already applied before the transform: (i)
            numpy array → torch tensor, (ii) HWC → CHW layout, and
            (iii) uint8 ``[0, 255]`` → float32 ``[0, 1]``.
        use_cached_video_metadata (bool): Use cached metadata for EncodedVideo
            objects if available. Set to False to force re-reading.
        n_frame_counting_workers (int): Parallel workers for pre-loading
            EncodedVideo metadata. -1 uses all available cores.
        progress_bar (bool | None): Show a progress bar during metadata loading.
            If None, shows one only when stderr is a TTY.
    """
    self.videos = videos
    self.transform = transform
    self.progress_bar = stderr.isatty() if progress_bar is None else progress_bar

    # Load video metadata before calling `.setup()` on the videos
    # This prevents redundant metadata loading if multiple Video objects are linked
    # to the same actual video
    unique_video_paths = set(
        vid.path for vid in videos if isinstance(vid, EncodedVideo)
    )
    pool = Parallel(n_jobs=n_frame_counting_workers)
    n_workers_effective = pool._effective_n_jobs()
    logger.info(
        f"Loading metadata for {len(unique_video_paths)} unique videos using "
        f"{n_frame_counting_workers} workers (effectively {n_workers_effective})."
    )
    _pbar = self.progress_bar and len(unique_video_paths) > n_workers_effective
    # Run metadata loading in parallel - no need to keep the results, we just want
    # to generate the metadata cache files
    _ = pool(
        delayed(get_video_metadata)(
            path, cache_metadata=True, use_cached_metadata=use_cached_video_metadata
        )
        for path in tqdm(
            unique_video_paths,
            file=stderr,
            disable=not _pbar,
            desc="Loading metadata",
        )
    )

    # Set up all videos and build an index of number of frames per video, using
    # the caches that we've just generated
    for video in self.videos:
        if isinstance(video, EncodedVideo) and not video.use_cached_metadata:
            logger.warning(
                "For efficiency, `use_cached_metadata` must be set to True for "
                "all EncodedVideo objects when they are managed by "
                "VideoCollectionDataset. Overriding to True."
            )
            video.use_cached_metadata = True
        video.setup()

    # Build index of number of frames per video
    self.n_frames_by_video = [len(video) for video in self.videos]
    self.n_frames_total = sum(self.n_frames_by_video)

    # Initialize attributes for worker assignment - to be filled in .assign_workers
    # The i-th worker's assignments are self.worker_assignments[i], which is a list
    # of (video_id, start_vir_frame_id, end_vir_frame_id) tuples.
    self.worker_assignments: list[list[tuple[int, int, int]]] = []

assign_workers

assign_workers(
    n_loading_workers: int, min_frames_per_worker: int = 300
) -> None

Distribute frames across workers for balanced parallel loading.

Frames are assigned in contiguous ranges within each video to minimise seeking overhead. If the per-worker frame count would fall below min_frames_per_worker, the effective worker count is reduced accordingly.

Parameters:

Name	Type	Description	Default
n_loading_workers	int	Number of workers to distribute frames across.	required
min_frames_per_worker	int	Lower bound on frames per worker. Workers below this threshold are merged to avoid excessive overhead.	300

Source code in src/pvio/torch_tools.py

def assign_workers(
    self, n_loading_workers: int, min_frames_per_worker: int = 300
) -> None:
    """Distribute frames across workers for balanced parallel loading.

    Frames are assigned in contiguous ranges within each video to minimise
    seeking overhead. If the per-worker frame count would fall below
    `min_frames_per_worker`, the effective worker count is reduced accordingly.

    Args:
        n_loading_workers (int): Number of workers to distribute frames across.
        min_frames_per_worker (int): Lower bound on frames per worker. Workers
            below this threshold are merged to avoid excessive overhead.
    """
    # Make an array of (video_id, vir_frame_id) pairs. frame_id is virtual (i.e.
    # index 0 corresponds to physical frame_id=frame_range[0])
    frame_specs_all = -np.ones((self.n_frames_total, 2), dtype=np.int32)
    start_global_frame_id = 0
    for video_id, n_frames in enumerate(self.n_frames_by_video):
        end_global_frame_id = start_global_frame_id + n_frames
        local_vfids = np.arange(n_frames)
        frame_specs_all[start_global_frame_id:end_global_frame_id, 0] = video_id
        frame_specs_all[start_global_frame_id:end_global_frame_id, 1] = local_vfids
        start_global_frame_id = end_global_frame_id
    if start_global_frame_id != self.n_frames_total:
        raise RuntimeError(
            f"Frame index mismatch after building frame_specs: "
            f"got {start_global_frame_id}, expected {self.n_frames_total}"
        )
    if np.any(frame_specs_all == -1):
        raise RuntimeError("frame_specs_all contains uninitialized entries")

    # Dynamically balance load among workers
    n_frames_per_worker = int(np.ceil(self.n_frames_total / n_loading_workers))
    logger.info(
        f"Assigning {self.n_frames_total} total frames from "
        f"{len(self.videos)} videos to {n_loading_workers} loading workers."
    )
    n_workers_effective = n_loading_workers
    if n_frames_per_worker < min_frames_per_worker:
        n_frames_per_worker = min_frames_per_worker
        n_workers_effective = int(
            np.ceil(self.n_frames_total / n_frames_per_worker)
        )
        logger.info(
            f"`n_frames_per_worker` is less than `min_frames_per_worker` "
            f"({min_frames_per_worker}). This will result in many workers working "
            f"on not so much data, leading to high overhead. "
            f"Increasing `n_frames_per_worker` to {n_frames_per_worker} and "
            f"reducing `n_loading_workers` to {n_workers_effective}."
        )

    # Initialize worker assignments (original number of workers! Unused workers get
    # empty assignments)
    self.worker_assignments = [[] for _ in range(n_loading_workers)]

    for worker_id in range(n_workers_effective):
        start_global_frame_id = worker_id * n_frames_per_worker
        end_global_frame_id = min(
            start_global_frame_id + n_frames_per_worker, self.n_frames_total
        )
        my_specs = frame_specs_all[start_global_frame_id:end_global_frame_id, :]
        # Convert to list of (video_id, start_vir_frame_id, end_vir_frame_id)
        if my_specs.shape[0] == 0:
            continue
        unique_video_ids = np.unique(my_specs[:, 0])
        for video_id in unique_video_ids:
            vir_frame_ids_local = my_specs[my_specs[:, 0] == video_id, 1]
            # Cast to Python int: these flow into the yielded batch dicts, and the
            # public API documents video_indices/frame_indices as lists of int (not
            # numpy.int32, which would surprise strict isinstance checks / JSON).
            start_vir_frame_id = int(vir_frame_ids_local[0])
            end_vir_frame_id = int(vir_frame_ids_local[-1]) + 1  # exclusive
            self.worker_assignments[worker_id].append(
                (int(video_id), start_vir_frame_id, end_vir_frame_id)
            )
    _nframes_total_check = 0
    for worker_chunks in self.worker_assignments:
        for _, start_vir_frame_id, end_vir_frame_id in worker_chunks:
            _nframes_total_check += end_vir_frame_id - start_vir_frame_id
    if _nframes_total_check != self.n_frames_total:
        raise RuntimeError(
            f"Frame count mismatch: assigned {_nframes_total_check} frames "
            f"but expected {self.n_frames_total}"
        )

pvio.torch_tools.VideoCollectionDataLoader

VideoCollectionDataLoader(
    dataset: VideoCollectionDataset,
    min_frames_per_worker: int = 300,
    **kwargs: Any,
)

Bases: DataLoader

DataLoader for VideoCollectionDataset with automatic worker assignment.

Each worker is assigned a contiguous range of frames to minimise seeking overhead. Batch dicts contain:

• frames: Tensor of shape (batch_size, C, H, W)
• video_indices: list of int — index into the videos list
• frame_indices: list of int — virtual frame index within that video

Custom batch_sampler and collate_fn are not supported.

Wrap a VideoCollectionDataset in a DataLoader with automatic worker assignment.

Parameters:

Name	Type	Description	Default
dataset	VideoCollectionDataset	The dataset to load from.	required
min_frames_per_worker	int	Minimum frames per worker; the effective worker count is reduced if this threshold would otherwise be breached.	300
**kwargs	Any	Forwarded to torch.utils.data.DataLoader.	{}

Source code in src/pvio/torch_tools.py

def __init__(
    self,
    dataset: VideoCollectionDataset,
    min_frames_per_worker: int = 300,
    **kwargs: Any,
) -> None:
    """Wrap a VideoCollectionDataset in a DataLoader with automatic worker assignment.

    Args:
        dataset (VideoCollectionDataset): The dataset to load from.
        min_frames_per_worker (int): Minimum frames per worker; the effective
            worker count is reduced if this threshold would otherwise be breached.
        **kwargs: Forwarded to torch.utils.data.DataLoader.
    """
    if not isinstance(dataset, VideoCollectionDataset):
        raise ValueError(
            "VideoCollectionDataLoader only works with VideoCollectionDataset."
        )
    if kwargs.get("batch_sampler") is not None:
        raise ValueError(
            "VideoCollectionDataLoader does not support custom batch samplers."
        )
    if kwargs.get("collate_fn") is not None:
        raise ValueError(
            "VideoCollectionDataLoader must use the built-in collate function."
        )

    kwargs["collate_fn"] = self._collate
    super().__init__(dataset, **kwargs)

    num_workers = self.num_workers
    if num_workers == 0:
        num_workers = 1

    self.dataset.assign_workers(
        n_loading_workers=num_workers, min_frames_per_worker=min_frames_per_worker
    )

pvio.torch_tools.SimpleVideoCollectionLoader

SimpleVideoCollectionLoader(
    videos: list[Path | str | Video],
    *,
    transform: Callable | None = None,
    buffer_size: int = 64,
    frame_id_regex: str
    | Pattern
    | None = "frame\\D*(\\d+)(?!\\d)",
    use_cached_video_metadata: bool = True,
    n_frame_counting_workers: int = -1,
    progress_bar: bool | None = None,
    min_frames_per_worker: int = 300,
    device: str | None = None,
    **kwargs: Any,
)

Bases: VideoCollectionDataLoader

Create a VideoCollectionDataset and VideoCollectionDataLoader in one call.

Each entry in videos may be a pre-constructed :class:Video object or a path (str / Path). Paths pointing to files become :class:EncodedVideo; paths pointing to directories become :class:ImageDirVideo.

The decode workflow is selected automatically: on a machine with a CUDA GPU, file-backed videos decode on the GPU (NVDEC) and iteration runs in the main process (num_workers is forced to 0, since CUDA cannot be used in forked workers); on a CPU-only machine, decoding uses the requested number of CPU workers as before. Pass device="cpu" to opt out and keep multi-worker CPU decoding even when a GPU is present.

Parameters:

Name	Type	Description	Default
videos	list[Path | str | Video]	Video sources.	required
transform	Callable | None	Applied to each CHW float32 frame tensor before it is yielded.	None
buffer_size	int	Decode-buffer size forwarded to :class:EncodedVideo for file-path entries.	64
frame_id_regex	str | Pattern | None	Regex forwarded to :class:ImageDirVideo for directory-path entries.	'frame\\D*(\\d+)(?!\\d)'
use_cached_video_metadata	bool	Use cached metadata when available. Set to False to force fresh reads.	True
n_frame_counting_workers	int	Workers for parallel metadata loading. -1 uses all available cores.	-1
progress_bar	bool | None	Show a progress bar during metadata loading. Defaults to True when stderr is a TTY.	None
min_frames_per_worker	int	Minimum frames per worker; see :meth:VideoCollectionDataset.assign_workers.	300
device	str | None	Decode device for file-backed videos, forwarded to :class:EncodedVideo. None (default) auto-selects the GPU when available; "cpu"/"cuda" force a choice.	None
**kwargs	Any	Forwarded to :class:~torch.utils.data.DataLoader.	{}

Source code in src/pvio/torch_tools.py

def __init__(
    self,
    videos: list[Path | str | Video],
    *,
    transform: Callable | None = None,
    buffer_size: int = 64,
    frame_id_regex: str | re.Pattern | None = r"frame\D*(\d+)(?!\d)",
    use_cached_video_metadata: bool = True,
    n_frame_counting_workers: int = -1,
    progress_bar: bool | None = None,
    min_frames_per_worker: int = 300,
    device: str | None = None,
    **kwargs: Any,
) -> None:
    """Create a VideoCollectionDataset and VideoCollectionDataLoader in one call.

    Each entry in *videos* may be a pre-constructed :class:`Video` object
    or a path (str / Path). Paths pointing to files become
    :class:`EncodedVideo`; paths pointing to directories become
    :class:`ImageDirVideo`.

    The decode workflow is selected automatically: on a machine with a CUDA
    GPU, file-backed videos decode on the GPU (NVDEC) and iteration runs in
    the main process (``num_workers`` is forced to 0, since CUDA cannot be
    used in forked workers); on a CPU-only machine, decoding uses the
    requested number of CPU workers as before. Pass ``device="cpu"`` to opt
    out and keep multi-worker CPU decoding even when a GPU is present.

    Args:
        videos: Video sources.
        transform: Applied to each CHW float32 frame tensor before it is
            yielded.
        buffer_size: Decode-buffer size forwarded to :class:`EncodedVideo`
            for file-path entries.
        frame_id_regex: Regex forwarded to :class:`ImageDirVideo` for
            directory-path entries.
        use_cached_video_metadata: Use cached metadata when available. Set
            to ``False`` to force fresh reads.
        n_frame_counting_workers: Workers for parallel metadata loading.
            ``-1`` uses all available cores.
        progress_bar: Show a progress bar during metadata loading.
            Defaults to ``True`` when stderr is a TTY.
        min_frames_per_worker: Minimum frames per worker; see
            :meth:`VideoCollectionDataset.assign_workers`.
        device: Decode device for file-backed videos, forwarded to
            :class:`EncodedVideo`. ``None`` (default) auto-selects the GPU
            when available; ``"cpu"``/``"cuda"`` force a choice.
        **kwargs: Forwarded to :class:`~torch.utils.data.DataLoader`.
    """
    logger.info(
        "Checking requested videos and creating Video objects from paths if needed"
    )
    video_objects = self._resolve_videos(
        videos,
        buffer_size=buffer_size,
        frame_id_regex=frame_id_regex,
        use_cached_video_metadata=use_cached_video_metadata,
        device=device,
    )

    # Resolve num_workers (and pin_memory) once, up front. On the GPU decode
    # path iteration must run in the main process — CUDA cannot be initialised
    # in forked DataLoader workers — so any CUDA-backed video forces
    # num_workers=0. Otherwise a negative spec follows the joblib convention
    # (-1 = all cores) and 0 is left as the standard main-process default
    # (VideoCollectionDataLoader still calls assign_workers(1) for it).
    gpu_decode = any(
        isinstance(v, EncodedVideo) and str(v.device).startswith("cuda")
        for v in video_objects
    )
    requested_workers = kwargs.get("num_workers", 0)  # 0 is the DataLoader default
    if gpu_decode:
        if requested_workers not in (0, None):
            logger.info(
                "GPU decoding selected; forcing num_workers=0 so iteration runs "
                "in the main process (was %s). CUDA cannot be used in forked "
                "DataLoader workers. Pass device='cpu' to keep CPU multi-worker "
                "loading.",
                requested_workers,
            )
        kwargs["num_workers"] = 0
        # Frames are already GPU-resident, so host-memory pinning is both
        # pointless and impossible (pin_memory cannot pin CUDA tensors).
        if kwargs.get("pin_memory"):
            logger.info(
                "GPU decoding selected; disabling pin_memory (frames are already "
                "on the GPU, so there is nothing to pin)."
            )
        kwargs["pin_memory"] = False
    elif requested_workers not in (0, None):
        kwargs["num_workers"] = _resolve_n_workers_spec(requested_workers)

    logger.info(f"Creating VideoCollectionDataset with {len(video_objects)} videos")
    dataset = VideoCollectionDataset(
        video_objects,
        transform=transform,
        use_cached_video_metadata=use_cached_video_metadata,
        n_frame_counting_workers=n_frame_counting_workers,
        progress_bar=progress_bar,
    )

    logger.info("Creating VideoCollectionDataLoader")
    super().__init__(dataset, min_frames_per_worker=min_frames_per_worker, **kwargs)