parallel-video-io 0.1.3__tar.gz

parallel_video_io-0.1.3/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Sibo Wang-Chen
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

parallel_video_io-0.1.3/PKG-INFO

@@ -0,0 +1,174 @@
+Metadata-Version: 2.4
+Name: parallel-video-io
+Version: 0.1.3
+Summary: Tools for reading and writing videos, and loading them efficiently with PyTorch.
+License: MIT
+License-File: LICENSE
+Author: Sibo Wang-Chen
+Author-email: sibo.wang@epfl.ch
+Requires-Python: >=3.10,<=3.13
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: imageio (>=2.37.0,<3)
+Requires-Dist: imageio-ffmpeg (==0.6.0)
+Requires-Dist: joblib (>=1.5.1,<2)
+Requires-Dist: numpy (>=2.0,<3)
+Requires-Dist: pytest (>=8.4.2,<9)
+Requires-Dist: torch (>=2.8,<3)
+Requires-Dist: torchcodec (==0.7.0)
+Requires-Dist: tqdm (>=4.67,<5)
+Project-URL: Repository, https://github.com/sibocw/parallel-video-io
+Description-Content-Type: text/markdown
+
+# parallel-video-io
+
+Tools for reading and writing videos and for efficient frame-level loading with PyTorch.
+
+This repository provides small, focused utilities around video I/O and a PyTorch-friendly iterable dataset + dataloader that make it easy to stream frames from many videos or directories of image frames in parallel.
+
+## Key features
+- Read frames from videos (random access or sequential) using imageio/ffmpeg.
+- Write sequences of numpy frames to H.264 MP4 files with sane defaults.
+- PyTorch-compatible `VideoCollectionDataset` and `VideoCollectionDataLoader` that provide a simple iterator that uses multiple processes to load data from different videos under the hood. This is especially handy for running trained deep learning models on many videos in production.
+
+## Table of contents
+- [Installation](#installation)
+- [Quick examples](#quick-examples)
+	- [Reading video metadata](#reading-video-metadata)
+	- [Reading video frames](#reading-video-frames)
+	- [Writing a video](#writing-a-video)
+	- [Using the PyTorch dataset and dataloader](#using-the-pytorch-dataset-and-dataloader)
+- [Testing](#testing)
+- [Notes & troubleshooting](#notes--troubleshooting)
+
+## Installation
+
+This project targets Python >= 3.10. The library's runtime dependencies are listed in `pyproject.toml` (torch, imageio, imageio-ffmpeg, torchcodec, joblib, tqdm, numpy, pytest).
+
+If you're using pip in a development environment, install editable with:
+
+```bash
+pip install -e .
+```
+
+Or with Poetry:
+
+```bash
+poetry install
+```
+
+You can include this package as a dependency for your project by including the following in your `pyproject.toml`:
+
+```toml
+[project]
+# ... other stuff
+dependencies = [
+    # ...
+    "parallel-video-io @ git+https://github.com/sibocw/parallel-video-io.git",
+]
+```
+
+Make sure `ffmpeg` is available on your `$PATH` (required by imageio-ffmpeg).
+
+## Quick examples
+
+These examples use NumPy arrays for frames in (height, width, channels) order and uint8 dtype.
+
+### Reading video metadata
+
+```python
+from pvio.video_io import get_video_metadata, check_num_frames
+
+# To get the number of frames in a video
+n_frames = check_num_frames("example.mp4")
+print(n_frames)  # this is an integer frame count
+
+# To get more information
+# Note that this function actually caches these information in a JSON file. To control
+# whether you want to save the cache file or disregard existing cache files, set the
+# `cache_metadata` (default True) and `use_cached_metadata` (default True) arguments.
+meta = get_video_metadata("example.mp4")
+print(meta)  # meta is a dictionary containing the keys "n_frames", "frame_size", "fps"
+```
+
+### Reading video frames
+
+```python
+from pvio.video_io import read_frames_from_video
+
+# You can read a whole video
+frames, fps = read_frames_from_video("example.mp4")
+
+# ... or just some frames
+frames, fps = read_frames_from_video("example.mp4", frame_indices=[0, 5])
+```
+
+### Writing a video
+
+```python
+import numpy as np
+from pvio.video_io import write_frames_to_video
+
+# Create dummy 32x32 RGB frames (H, W, C)
+frames = [np.full((32, 32, 3), fill_value=i, dtype=np.uint8) for i in range(10)]
+
+# Save them to file
+# There are more complex video writing parameters that can be tuned - see the docstring
+# for details.
+write_frames_to_video("example.mp4", frames, fps=25.0)
+```
+
+Notes: the writer verifies that all frames share the same (height, width). FFmpeg can
+automatically resize frames to meet codec alignment requirements; for deterministic
+results, use dimensions divisible by 16.
+
+### Using the PyTorch dataset and dataloader
+
+The `VideoCollectionDataset` iterates frames either from video files or from directories containing individual image frames.
+
+```python
+from pvio.torch import VideoCollectionDataset, VideoCollectionDataLoader
+
+# Initialize Dataset from video files
+paths = ["/path/to/video1.mp4", "/path/to/video2.mp4"]
+ds = VideoCollectionDataset(paths)
+# ... or from directories containing individual frames as images
+paths = ["/path/to/frames_dir1", "/path/to/frames_dir2"]
+# To control sorting of frame files within each dir, use the `frame_sorting` argument
+# (see docstring for details)
+ds = VideoCollectionDataset(paths, as_image_dirs=True)
+
+# Wrap in the special DataLoader
+# (you can add other DataLoader keyword arguments if you wish)
+loader = VideoCollectionDataLoader(ds, batch_size=8, num_workers=4)
+
+# Now you can iterate over all frames from all videos in a single iterator. Behind the
+# scenes, these frames are fetched in parallel (each worker handles one video at a time)
+for batch in loader:
+	frames = batch["frames"]  # torch.Tensor: B x C x H x W
+	video_paths = batch["video_paths"]  # list of Path or str, depending on input
+	frame_indices = batch["frame_indices"]  # list of int
+```
+
+When loading from video files (as_image_dirs=False), the dataset uses `torchcodec`'s `VideoDecoder` to decode frames and `get_video_metadata` to build per-video frame counts; you may want to enable caching if you index many large files.
+
+## Testing
+
+The test suite uses pytest. Run it from the repository root:
+
+```bash
+pytest tests
+```
+
+There are a few tests that write small MP4 files using imageio/ffmpeg; ensure `ffmpeg` is available in the environment where tests run.
+
+## Notes & troubleshooting
+
+- FFmpeg macroblock constraints: some ffmpeg builds require frame dimensions to be divisible by 16. If you see a warning about `macro_block_size=16` and unexpected resizing, choose frame sizes divisible by 16 in production pipelines.
+- If you plan to decode many large videos, enabling metadata caching (the package writes a `.metadata.json` next to each video when `get_video_metadata` is called) will speed up repeated indexing.
+- The PyTorch loader expects the dataset passed to `VideoCollectionDataLoader` to be an instance of `VideoCollectionDataset` and enforces the built-in collate function.
+

parallel_video_io-0.1.3/README.md

@@ -0,0 +1,147 @@
+# parallel-video-io
+
+Tools for reading and writing videos and for efficient frame-level loading with PyTorch.
+
+This repository provides small, focused utilities around video I/O and a PyTorch-friendly iterable dataset + dataloader that make it easy to stream frames from many videos or directories of image frames in parallel.
+
+## Key features
+- Read frames from videos (random access or sequential) using imageio/ffmpeg.
+- Write sequences of numpy frames to H.264 MP4 files with sane defaults.
+- PyTorch-compatible `VideoCollectionDataset` and `VideoCollectionDataLoader` that provide a simple iterator that uses multiple processes to load data from different videos under the hood. This is especially handy for running trained deep learning models on many videos in production.
+
+## Table of contents
+- [Installation](#installation)
+- [Quick examples](#quick-examples)
+	- [Reading video metadata](#reading-video-metadata)
+	- [Reading video frames](#reading-video-frames)
+	- [Writing a video](#writing-a-video)
+	- [Using the PyTorch dataset and dataloader](#using-the-pytorch-dataset-and-dataloader)
+- [Testing](#testing)
+- [Notes & troubleshooting](#notes--troubleshooting)
+
+## Installation
+
+This project targets Python >= 3.10. The library's runtime dependencies are listed in `pyproject.toml` (torch, imageio, imageio-ffmpeg, torchcodec, joblib, tqdm, numpy, pytest).
+
+If you're using pip in a development environment, install editable with:
+
+```bash
+pip install -e .
+```
+
+Or with Poetry:
+
+```bash
+poetry install
+```
+
+You can include this package as a dependency for your project by including the following in your `pyproject.toml`:
+
+```toml
+[project]
+# ... other stuff
+dependencies = [
+    # ...
+    "parallel-video-io @ git+https://github.com/sibocw/parallel-video-io.git",
+]
+```
+
+Make sure `ffmpeg` is available on your `$PATH` (required by imageio-ffmpeg).
+
+## Quick examples
+
+These examples use NumPy arrays for frames in (height, width, channels) order and uint8 dtype.
+
+### Reading video metadata
+
+```python
+from pvio.video_io import get_video_metadata, check_num_frames
+
+# To get the number of frames in a video
+n_frames = check_num_frames("example.mp4")
+print(n_frames)  # this is an integer frame count
+
+# To get more information
+# Note that this function actually caches these information in a JSON file. To control
+# whether you want to save the cache file or disregard existing cache files, set the
+# `cache_metadata` (default True) and `use_cached_metadata` (default True) arguments.
+meta = get_video_metadata("example.mp4")
+print(meta)  # meta is a dictionary containing the keys "n_frames", "frame_size", "fps"
+```
+
+### Reading video frames
+
+```python
+from pvio.video_io import read_frames_from_video
+
+# You can read a whole video
+frames, fps = read_frames_from_video("example.mp4")
+
+# ... or just some frames
+frames, fps = read_frames_from_video("example.mp4", frame_indices=[0, 5])
+```
+
+### Writing a video
+
+```python
+import numpy as np
+from pvio.video_io import write_frames_to_video
+
+# Create dummy 32x32 RGB frames (H, W, C)
+frames = [np.full((32, 32, 3), fill_value=i, dtype=np.uint8) for i in range(10)]
+
+# Save them to file
+# There are more complex video writing parameters that can be tuned - see the docstring
+# for details.
+write_frames_to_video("example.mp4", frames, fps=25.0)
+```
+
+Notes: the writer verifies that all frames share the same (height, width). FFmpeg can
+automatically resize frames to meet codec alignment requirements; for deterministic
+results, use dimensions divisible by 16.
+
+### Using the PyTorch dataset and dataloader
+
+The `VideoCollectionDataset` iterates frames either from video files or from directories containing individual image frames.
+
+```python
+from pvio.torch import VideoCollectionDataset, VideoCollectionDataLoader
+
+# Initialize Dataset from video files
+paths = ["/path/to/video1.mp4", "/path/to/video2.mp4"]
+ds = VideoCollectionDataset(paths)
+# ... or from directories containing individual frames as images
+paths = ["/path/to/frames_dir1", "/path/to/frames_dir2"]
+# To control sorting of frame files within each dir, use the `frame_sorting` argument
+# (see docstring for details)
+ds = VideoCollectionDataset(paths, as_image_dirs=True)
+
+# Wrap in the special DataLoader
+# (you can add other DataLoader keyword arguments if you wish)
+loader = VideoCollectionDataLoader(ds, batch_size=8, num_workers=4)
+
+# Now you can iterate over all frames from all videos in a single iterator. Behind the
+# scenes, these frames are fetched in parallel (each worker handles one video at a time)
+for batch in loader:
+	frames = batch["frames"]  # torch.Tensor: B x C x H x W
+	video_paths = batch["video_paths"]  # list of Path or str, depending on input
+	frame_indices = batch["frame_indices"]  # list of int
+```
+
+When loading from video files (as_image_dirs=False), the dataset uses `torchcodec`'s `VideoDecoder` to decode frames and `get_video_metadata` to build per-video frame counts; you may want to enable caching if you index many large files.
+
+## Testing
+
+The test suite uses pytest. Run it from the repository root:
+
+```bash
+pytest tests
+```
+
+There are a few tests that write small MP4 files using imageio/ffmpeg; ensure `ffmpeg` is available in the environment where tests run.
+
+## Notes & troubleshooting
+
+- FFmpeg macroblock constraints: some ffmpeg builds require frame dimensions to be divisible by 16. If you see a warning about `macro_block_size=16` and unexpected resizing, choose frame sizes divisible by 16 in production pipelines.
+- If you plan to decode many large videos, enabling metadata caching (the package writes a `.metadata.json` next to each video when `get_video_metadata` is called) will speed up repeated indexing.
+- The PyTorch loader expects the dataset passed to `VideoCollectionDataLoader` to be an instance of `VideoCollectionDataset` and enforces the built-in collate function.

parallel_video_io-0.1.3/pyproject.toml

@@ -0,0 +1,31 @@
+[project]
+name = "parallel-video-io"
+version = "0.1.3"
+description = "Tools for reading and writing videos, and loading them efficiently with PyTorch."
+authors = [
+    {name = "Sibo Wang-Chen", email = "sibo.wang@epfl.ch"}
+]
+readme = "README.md"
+license = {text = "MIT"}
+requires-python = ">=3.10,<=3.13"
+dependencies = [
+    "torch>=2.8,<3",
+    "imageio>=2.37.0,<3",
+    "imageio-ffmpeg==0.6.0",
+    "torchcodec==0.7.0",
+    "joblib>=1.5.1,<2",
+    "tqdm>=4.67,<5",
+    "numpy>=2.0,<3",
+    "pytest>=8.4.2,<9",
+]
+
+[project.urls]
+repository = "https://github.com/sibocw/parallel-video-io"
+
+[tool.poetry]
+packages = [{include = "pvio", from = "src"}]
+
+
+[build-system]
+requires = ["poetry-core>=2.0.0,<3.0.0"]
+build-backend = "poetry.core.masonry.api"

parallel_video_io-0.1.3/src/pvio/__init__.py

@@ -0,0 +1,7 @@
+from .video_io import (
+    read_frames_from_video,
+    write_frames_to_video,
+    check_num_frames,
+    get_video_metadata,
+)
+from .torch import VideoCollectionDataset, VideoCollectionDataLoader

parallel_video_io-0.1.3/src/pvio/torch.py

@@ -0,0 +1,220 @@
+import torch
+import logging
+import re
+import imageio.v2 as imageio
+from typing import Callable
+from torchcodec.decoders import VideoDecoder
+from torch.utils.data import IterableDataset, DataLoader, get_worker_info
+from pathlib import Path
+from tqdm import tqdm
+from joblib import Parallel, delayed
+
+from .video_io import get_video_metadata
+from .util import balance_load_lpt
+
+
+class VideoCollectionDataset(IterableDataset):
+    """Yields individual frames from Spotlight behavior recordings, either
+    from videos or from image sequences."""
+
+    def __init__(
+        self,
+        paths: list[Path | str],
+        as_image_dirs: bool = False,
+        frame_sorting: None | str = None,
+        transform: Callable | None = None,
+    ):
+        r"""
+        Args:
+            paths (list[Path]): List of to video paths, or directories
+                containing frames as individual images.
+            as_image_dirs (bool): If True, treat each path as a directory
+                containing individual frames. Otherwise, treat it as a
+                video file.
+            frame_sorting (str | None): When `as_image_dirs` is True, this
+                argument specifies how images within each directory should
+                be sorted. If None, files are sorted by name. If given as a
+                string, it is used as a regex pattern to extract frame
+                numbers from filenames (e.g. r"frame\D*(\d+)(?!\d)").
+                When `as_image_dirs` is False, this argument is ignored.
+            transform (Callable | None): A function that is to be applied
+                to each frame after loading. Note that the following
+                operations are already applied to each frame:
+                (i) conversion from numpy array to torch tensor,
+                (ii) conversion from HWC to CHW format, and
+                (iii) conversion from uint8 in [0, 255] to float in [0, 1].
+                The transform function, if provided, is applied after these
+                operations.
+        """
+        self.video_paths = [Path(p) for p in paths]
+        self.worker_assignments = None
+        self.as_image_dirs = as_image_dirs
+        self.frame_sorting = frame_sorting
+        self.n_frames_lookup = None  # Populated by assign_workers()
+        self.transform = transform
+
+        # Check if the paths are all valid
+        for p in self.video_paths:
+            if self.as_image_dirs:
+                if not p.is_dir():
+                    raise ValueError(
+                        f"One of the specified paths {p} is not a valid directory. "
+                        "Directories containing individual frame images are expected."
+                    )
+            else:
+                if not p.is_file():
+                    raise ValueError(
+                        f"One of the specified paths {p} is not a valid file. "
+                        "Video files are expected."
+                    )
+
+        # Sort images if we're loading from directories of images
+        self.frame_sortings = {}
+        regex = re.compile(frame_sorting) if frame_sorting else None
+        if as_image_dirs:
+            # Iterate over the canonical Path objects (self.video_paths) so we
+            # consistently store Path keys and avoid relying on caller types
+            for path in self.video_paths:
+                all_files = [f for f in path.iterdir() if f.is_file()]
+                if regex is None:
+                    sorting_func = lambda f: f.name
+                else:
+                    sorting_func = lambda f: self._extract_frame_number(f.name, regex)
+                # Store a new sorted list (list.sort() returns None)
+                self.frame_sortings[path] = sorted(all_files, key=sorting_func)
+
+    def assign_workers(
+        self, n_frame_loading_workers: int, n_metadata_indexing_workers: int = -1
+    ):
+        # Check how many frame loading workers we're actually using (e.g. -1 actually
+        # means all available cores, so we need to figure out how many that is)
+        n_frame_loading_workers_effective = Parallel(
+            n_jobs=n_frame_loading_workers
+        )._effective_n_jobs()
+        logging.info(
+            f"Caller specified {n_frame_loading_workers} workers for frame loading. "
+            f"This is effectively {n_frame_loading_workers_effective} workers."
+        )
+
+        # Figure out how many frames there are in each video. This allows us to split
+        # the workload more evenly among workers by the number of frames.
+        if self.as_image_dirs:
+            self.n_frames_lookup = {
+                path: len(frames) for path, frames in self.frame_sortings.items()
+            }
+        else:
+            # Count frames in videos. This requires partially decoding the video files
+            # and it can be quite slow, so we do it in parallel and use caches.
+            logging.info(
+                f"Loading metadata for {len(self.video_paths)} videos. "
+                "This may take a while if no cached metadata is available."
+            )
+            metas = Parallel(n_jobs=n_metadata_indexing_workers)(
+                delayed(get_video_metadata)(path)
+                for path in tqdm(self.video_paths, desc="Indexing videos", disable=None)
+            )
+            self.n_frames_lookup = {
+                path: meta["n_frames"] for path, meta in zip(self.video_paths, metas)
+            }
+
+        # Split videos evenly among the available workers
+        self.worker_assignments = balance_load_lpt(
+            self.n_frames_lookup, max(1, n_frame_loading_workers_effective)
+        )
+
+    def __iter__(self):
+        # Get worker info for distributed loading
+        worker_info = get_worker_info()
+        if worker_info is None:
+            # Single process
+            video_subset = self.video_paths
+        else:
+            # Split videos among workers
+            video_subset = self.worker_assignments[worker_info.id]
+
+        # Each worker sequentially decodes its assigned videos
+        for video_path in video_subset:
+            if self.as_image_dirs:
+                # Read individual images
+                frame_files = self.frame_sortings[video_path]
+                for frame_idx, frame_file in enumerate(frame_files):
+                    frame = imageio.imread(frame_file)
+                    frame = torch.from_numpy(frame)
+                    if frame.ndim == 2:
+                        frame = frame.unsqueeze(-1)  # add channel dim
+                    frame = frame.permute(2, 0, 1)  # HWC to CHW
+                    frame = frame.float() / 255.0  # to float in [0, 1]
+                    if self.transform:
+                        frame = self.transform(frame)
+                    yield {
+                        "frame": frame,
+                        "video_path": video_path,
+                        "frame_idx": frame_idx,
+                    }
+            else:
+                # Use torchcodec to decode videos
+                decoder = VideoDecoder(video_path)
+                for frame_idx in range(len(decoder)):
+                    frame = decoder[frame_idx]  # returns tensor in CHW
+                    frame = frame.float() / 255.0  # to float in [0, 1]
+                    if self.transform:
+                        frame = self.transform(frame)
+                    yield {
+                        "frame": frame,
+                        "video_path": video_path,
+                        "frame_idx": frame_idx,
+                    }
+
+    def __len__(self):
+        if self.n_frames_lookup is None:
+            raise ValueError(
+                "VideoCollectionDataset length is unknown until workers are assigned. "
+                "Call `assign_workers()` before using `len()`."
+            )
+        return sum(self.n_frames_lookup.values())
+
+    @staticmethod
+    def _extract_frame_number(filename: str, regex_pattern: str) -> int:
+        matches = re.findall(regex_pattern, filename)
+        if len(matches) != 1:
+            raise ValueError(
+                f"{len(matches)} matches found in filename {filename} "
+                f"using regex pattern {regex_pattern}. Only one match is expected."
+            )
+        try:
+            return int(matches[0])
+        except ValueError as e:
+            raise ValueError(
+                f"Failed to parse '{matches[0]}' as int. This substring is extracted "
+                f"from filename {filename} using regex pattern {regex_pattern}."
+            ) from e
+
+
+class VideoCollectionDataLoader(DataLoader):
+    def __init__(self, dataset: VideoCollectionDataset, **kwargs):
+        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)
+
+        self.dataset.assign_workers(n_frame_loading_workers=self.num_workers)
+
+    @staticmethod
+    def _collate(batch):
+        """Receives a list of frame dicts, returns a batched dict"""
+        return {
+            "frames": torch.stack([item["frame"] for item in batch]),
+            "video_paths": [item["video_path"] for item in batch],
+            "frame_indices": [item["frame_idx"] for item in batch],
+        }

parallel_video_io-0.1.3/src/pvio/util.py

@@ -0,0 +1,34 @@
+import numpy as np
+from typing import Hashable
+
+
+def balance_load_lpt(
+    tasks: dict[Hashable, int], n_workers: int
+) -> list[list[Hashable]]:
+    """The Longest Processing Time (LPT) algorithm for load balancing: sort
+    tasks by decreasing duration and assigns each task to the worker with
+    the currently smallest total assigned load.
+
+    Args:
+        tasks (dict[Hashable, int]): A dict mapping task identifiers (can
+            be any hashable type) to their estimated durations.
+        n_workers (int): Number of workers to distribute tasks across.
+
+    Returns:
+        assignments (list[list[Hashable]]): A list of lists, where
+            assignments[i] contains the IDs of tasks assigned to worker i.
+    """
+    # Sort tasks by descending duration
+    sorted_tasks = sorted(tasks.items(), key=lambda x: x[1], reverse=True)
+
+    # Initialize workers and their current loads
+    worker_loads = np.zeros(n_workers)
+    assignments = [[] for _ in range(n_workers)]
+
+    # Assign each task to the currently least-loaded worker
+    for task_id, duration in sorted_tasks:
+        i = int(np.argmin(worker_loads))
+        assignments[i].append(task_id)
+        worker_loads[i] += duration
+
+    return assignments

parallel_video_io-0.1.3/src/pvio/video_io.py

@@ -0,0 +1,159 @@
+import numpy as np
+import json
+import logging
+import imageio.v2 as imageio
+from pathlib import Path
+
+
+def read_frames_from_video(
+    video_path: Path | str, frame_indices: list[int] | None = None
+) -> tuple[list[np.ndarray], float]:
+    """Read specific frames from a video file.
+
+    Args:
+        video_path (Path | str): Path to the video file.
+        frame_indices (list[int] | None): List of frame indices to read.
+            If None, read all frames.
+
+    Raises:
+        ValueError: If the video file cannot be read.
+        IndexError: If the frame indices are invalid.
+
+    Returns:
+        frames (list[np.ndarray]): List of frames as numpy arrays.
+        fps (float): FPS of the video.
+    """
+    frames = []
+    with imageio.get_reader(video_path) as reader:
+        if frame_indices is None:
+            frame_indices = list(range(reader.count_frames()))
+        for idx in frame_indices:
+            frames.append(reader.get_data(idx))
+        fps = reader.get_meta_data().get("fps", None)
+    return frames, fps
+
+
+_default_ffmpeg_params_for_video_writing = [
+    "-crf",
+    "15",  # Lower CRF = higher quality (15 is very high quality)
+    "-preset",
+    "slow",  # Slower preset = better compression efficiency
+    "-profile:v",
+    "high",  # Use high profile for better compression
+    "-level",
+    "4.0",  # H.264 level
+]
+
+
+def write_frames_to_video(
+    video_path: Path | str,
+    frames: list[np.ndarray],
+    fps: float,
+    codec: str = "libx264",
+    ffmpeg_params: list[str] = _default_ffmpeg_params_for_video_writing,
+    log_interval: int | None = None,
+    log_level: int = logging.INFO,
+):
+    """Write a sequence of frames to a video file.
+
+    Args:
+        video_path (Path | str): Path to save the video file.
+        frames (list[np.ndarray]): List of frames as numpy arrays (in
+            [height, width, channels] format).
+        fps (float): Frames per second for the output video.
+        codec (str): Codec to use. Default: 'libx264'.
+        ffmpeg_params (list[str]): Additional ffmpeg parameters.
+            Default is a set of parameters for high-quality H.264 encoding.
+            (see _default_ffmpeg_params_for_video_writing).
+        log_interval (int | None): If set, log progress every
+            `log_interval` frames at the specified log level.
+        log_level (int): Logging level for progress. Default: logging.INFO.
+    """
+    # Check frame size consistency
+    if len(frames) == 0:
+        raise ValueError("No frames provided to write_frames_to_video")
+    frame_size = frames[0].shape[:2]
+    for frame in frames:
+        if frame.shape[:2] != frame_size:
+            raise ValueError(
+                "All frames must have the same dimensions. The 0th frame has size "
+                f"{frame_size}, but at least one frame has size {frame.shape[:2]}."
+            )
+
+    # Use imageio to write video with ffmpeg backend
+    with imageio.get_writer(
+        str(video_path),
+        "ffmpeg",
+        fps=fps,
+        codec=codec,
+        quality=None,  # Use CRF (in ffmpeg_params) instead of quality
+        ffmpeg_params=ffmpeg_params,
+    ) as video_writer:
+        for i, frame in enumerate(frames):
+            video_writer.append_data(frame)
+
+            if log_interval is not None and i % log_interval == 0:
+                logging.log(log_level, f"Written frame {i + 1}/{len(frames)}")
+
+
+def check_num_frames(video_path: Path | str) -> int:
+    """Check number of frames in a video file."""
+    try:
+        with imageio.get_reader(video_path) as reader:
+            num_frames = reader.count_frames()
+    except Exception as e:
+        raise RuntimeError(f"Failed to open video file: {video_path}") from e
+    return num_frames
+
+
+def get_video_metadata(
+    video_path: Path | str,
+    cache_metadata: bool = True,
+    use_cached_metadata: bool = True,
+    metadata_suffix: str = ".metadata.json",
+):
+    """Get number of frames, frame size, and FPS of a video file.
+
+    Args:
+        video_path (Path | str): Path to the video file.
+        cache_metadata (bool): Whether to cache the metadata to a JSON
+            file. Default is True.
+        use_cached_metadata (bool): Whether to use cached metadata if
+            available. Default is True.
+        metadata_suffix (str): Suffix to use for the metadata cache file.
+            Default is ".metadata.json".
+
+    Returns:
+        dict: A dictionary containing the video metadata.
+    """
+    metadata = {}
+
+    video_path = Path(video_path)
+    cache_path = video_path.with_suffix(metadata_suffix)
+    if use_cached_metadata and cache_path.is_file():
+        try:
+            with open(cache_path, "r") as f:
+                metadata = json.load(f)
+            n_frames = metadata["n_frames"]
+            frame_size = tuple(metadata["frame_size"])
+            fps = metadata["fps"]
+        except Exception as e:
+            print(f"Corrupted metadata cache file {cache_path}")
+            raise e
+    else:
+        n_frames = check_num_frames(video_path)
+        sample_frames, fps = read_frames_from_video(
+            video_path, frame_indices=[0]
+        )
+        frame_size = sample_frames[0].shape[:2]
+
+        if cache_metadata:
+            metadata = {
+                "n_frames": n_frames,
+                "frame_size": list(frame_size),
+                "fps": fps,
+            }
+            with open(cache_path, "w") as f:
+                json.dump(metadata, f, indent=2)
+
+    return {"n_frames": n_frames, "frame_size": frame_size, "fps": fps}