Add video module. VideoWriter moved to video module

README.md

@@ -35,8 +35,10 @@ See the [documentation](https://dmf-utils.readthedocs.io/) for more installation
 ## Modules
 
 * [Alerts](https://dmf-utils.readthedocs.io/en/latest/modules/alerts.html): Get notified when a function finishes running and send messages or files to Slack and Telegram.
-* [IO](https://dmf-utils.readthedocs.io/en/latest/modules/io.html): Load and save data from different formats, and manage compressed files.
 * [Env](https://dmf-utils.readthedocs.io/en/latest/modules/env.html): Manage environment variables.
+* [IO (Input/Output)](https://dmf-utils.readthedocs.io/en/latest/modules/io.html): Load and save data from different formats, and manage compressed files.
+* [Video](https://dmf-utils.readthedocs.io/en/latest/modules/video.html): Utilities to work with video files.
+
 
 See the [modules documentation](https://dmf-utils.readthedocs.io/en/latest/modules/index.html) for more information.
 

dmf/__init__.py

@@ -4,14 +4,15 @@
 
 from .__version__ import __version__
 
-subpackages = ["alerts", "io", "env"]
+subpackages = ["alerts", "io", "env", "video"]
 
 __getattr__, __dir__, __all__ = lazy.attach(__name__, subpackages)
 
 if TYPE_CHECKING:
     from . import alerts
     from . import io
     from . import env
+    from . import video
 
 __all__ = ["__version__", "alerts", "io", "env"]
 

dmf/__version__.py

@@ -1,2 +1,2 @@
 
-__version__ = "0.1.1"
+__version__ = "0.1.2"

dmf/io/__init__.py

@@ -7,7 +7,6 @@
     "decompress": ["decompress"],
     "load": ["load"],
     "save": ["save"],
-    "video": ["VideoWriter"],
 }
 
 __getattr__, __dir__, __all__ = lazy.attach(__name__, submod_attrs=submod_attrs)
@@ -17,7 +16,6 @@
     from .decompress import decompress
     from .load import load
     from .save import save
-    from .video import VideoWriter
 
 
-__all__ = ["compress", "decompress", "load", "save", "VideoWriter"]
+__all__ = ["compress", "decompress", "load", "save"]

dmf/io/save.py

@@ -311,6 +311,5 @@ def save_audio(data: Any, file_path: Path, **kwargs):
 @register_saver("video", ["mp4", "avi", "mov", "mkv"])
 def save_video(data: Any, file_path: Path, **kwargs):
     """Save data using the video saver."""
-    from .video import VideoWriter
-    writer = VideoWriter(file_path, **kwargs)
-    writer.generate_video(data)
+    from ..video.video_writer import write_video
+    write_video(file_path=file_path, frames=data, **kwargs)

dmf/video/__init__.py

@@ -0,0 +1,16 @@
+from typing import TYPE_CHECKING
+
+import lazy_loader as lazy
+
+submod_attrs = {
+    "video_writer": ["VideoWriter", "write_video"],
+    "video_reader": ["VideoReader", "read_video"],
+}
+
+__getattr__, __dir__, __all__ = lazy.attach(__name__, submod_attrs=submod_attrs)
+
+if TYPE_CHECKING:
+    from .video_writer import VideoWriter, write_video
+    from .video_reader import VideoReader, read_video
+
+__all__ = ["write_video", "read_video", "VideoWriter", "VideoReader"]

dmf/video/video_reader.py

@@ -0,0 +1,203 @@
+from pathlib import Path
+from typing import Union, List, Literal, Iterator
+
+try:
+    import cv2
+except ImportError:
+    raise ImportError(
+        "OpenCV is required for video reading. "
+        "Install it using `pip install opencv-python`."
+    )
+
+try:
+    import numpy as np
+except ImportError:
+    raise ImportError(
+        "NumPy is required for video reading. "
+        "Install it using `pip install numpy`."
+    )
+
+from PIL import Image
+
+__all__ = ["read_video", "VideoReader"]
+
+OutputType = Literal["numpy", "pil"]
+
+
+def read_video(
+    file_path: Union[str, Path], 
+    output_type: OutputType = "numpy"
+) -> Union[np.ndarray, List[Image.Image]]:
+    """
+    Read an entire video and return the frames as either NumPy arrays or PIL images.
+
+    Parameters
+    ----------
+    file_path : Union[str, Path]
+        The path to the input video file.
+    output_type : Literal["numpy", "pil"], default="numpy"
+        The desired output type for the frames. "numpy" returns frames as NumPy arrays,
+        while "pil" returns frames as PIL images.
+
+    Returns
+    -------
+    Union[np.ndarray, List[Image.Image]]
+        A NumPy array of shape (num_frames, height, width, 3) if output_type is "numpy",
+        or a list of PIL images if output_type is "pil".
+    """
+    with VideoReader(file_path, output_type=output_type) as reader:
+        return reader.read_video()
+
+
+class VideoReader:
+    """
+    A utility class to read videos and return frames as either NumPy arrays or PIL images.
+
+    Parameters
+    ----------
+    file_path : Union[str, Path]
+        The path to the input video file.
+    output_type : Literal["numpy", "pil"], default="numpy"
+        The desired output type for the frames. "numpy" returns frames as NumPy arrays,
+        while "pil" returns frames as PIL images.
+    
+    Examples
+    --------
+    Reading the entire video as NumPy arrays:
+
+    .. code-block:: python
+
+        reader = VideoReader("input.mp4", output_type="numpy")
+        frames = reader.read_video()
+
+    Iterating over video frames as PIL images:
+
+    .. code-block:: python
+
+        reader = VideoReader("input.mp4", output_type="pil")
+        for frame in reader:
+            frame.show()  # Display the frame using PIL's show method
+
+    Accessing a specific frame by index:
+
+    .. code-block:: python
+
+        reader = VideoReader("input.mp4", output_type="numpy")
+        frame = reader[10]  # Get the 11th frame (index starts at 0)
+    """
+
+    def __init__(self, file_path: Union[str, Path], output_type: OutputType = "numpy"):
+        self.file_path = Path(file_path)
+        self.output_type = output_type
+        self._cap = None
+        self._frame_count = 0
+        self._initialize_reader()
+
+    def _initialize_reader(self):
+        """Initialize the video reader."""
+        self._cap = cv2.VideoCapture(str(self.file_path))
+        if not self._cap.isOpened():
+            raise FileNotFoundError(f"Unable to open video file: {self.file_path}")
+        self._frame_count = int(self._cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+    def read_video(self) -> Union[np.ndarray, List[Image.Image]]:
+        """
+        Read the entire video and return all frames.
+
+        Returns
+        -------
+        Union[np.ndarray, List[Image.Image]]
+            A NumPy array of shape (num_frames, height, width, 3) if output_type is "numpy",
+            or a list of PIL images if output_type is "pil".
+        """
+        frames = []
+        while True:
+            ret, frame = self._cap.read()
+            if not ret:
+                break
+            frames.append(self._process_frame(frame))
+
+        self._cap.release()
+
+        if self.output_type == "numpy":
+            return np.array(frames)
+        return frames
+
+    def _process_frame(self, frame: np.ndarray) -> Union[np.ndarray, Image.Image]:
+        """Convert the frame to the desired output type."""
+        if self.output_type == "pil":
+            return Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+        return frame
+
+    def __getitem__(self, index: int) -> Union[np.ndarray, Image.Image]:
+        """
+        Retrieve a specific frame by index.
+
+        Parameters
+        ----------
+        index : int
+            The index of the frame to retrieve (0-based).
+
+        Returns
+        -------
+        Union[np.ndarray, Image.Image]
+            The frame at the specified index.
+        """
+        if index < 0 or index >= self._frame_count:
+            raise IndexError("Frame index out of range")
+
+        self._cap.set(cv2.CAP_PROP_POS_FRAMES, index)
+        ret, frame = self._cap.read()
+        if not ret:
+            raise ValueError(f"Failed to retrieve frame at index {index}")
+
+        return self._process_frame(frame)
+
+    def __iter__(self) -> Iterator[Union[np.ndarray, Image.Image]]:
+        """
+        Iterate over video frames one by one.
+
+        Yields
+        ------
+        Union[np.ndarray, Image.Image]
+            The next frame in the video as either a NumPy array or a PIL image.
+        """
+        self._initialize_reader()
+        return self
+
+    def __next__(self) -> Union[np.ndarray, Image.Image]:
+        """
+        Return the next frame in the video.
+
+        Returns
+        -------
+        Union[np.ndarray, Image.Image]
+            The next frame in the video as either a NumPy array or a PIL image.
+
+        Raises
+        ------
+        StopIteration
+            If the video has no more frames.
+        """
+        if not self._cap.isOpened():
+            raise StopIteration
+
+        ret, frame = self._cap.read()
+        if not ret:
+            self._cap.release()
+            raise StopIteration
+
+        return self._process_frame(frame)
+
+    def __len__(self) -> int:
+        """Return the total number of frames in the video."""
+        return self._frame_count
+
+    def __enter__(self):
+        """Context management enter method."""
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """Context management exit method."""
+        if self._cap:
+            self._cap.release()

dmf/io/video.py → dmf/video/video_writer.py

@@ -13,14 +13,17 @@
     import numpy as np
 except ImportError:
     raise ImportError(
-        "NumPy is required for video writing. Install it using `pip install numpy`."
+        "NumPy is required for video writing. "
+        "Install it using `pip install numpy`."
     )
 
 if TYPE_CHECKING:
     import numpy as np
     from PIL import Image
     import matplotlib.pyplot as plt
 
+__all__ = ["VideoWriter", "write_video"]
+
 FrameType = Union["np.ndarray", "Image.Image", "plt.Figure", Path, str]
 
 CODECS_MAPPING = {
@@ -31,6 +34,39 @@
 }
 
 
+def write_video(
+    file_path: Union[str, Path],
+    frames: Iterable[FrameType],
+    fps: int = 30,
+    codec: Optional[str] = None,
+) -> Path:
+    """
+    Write a video from a list of frames.
+
+    Parameters
+    ----------
+    frames : Iterable[FrameType]
+        An iterable of frames to add to the video.
+    file_path : Union[str, Path]
+        The path where the output video file will be saved.
+    fps : int, default=30
+        Frames per second (FPS) for the output video.
+    codec : Optional[str], default=None
+        The codec to use for video compression.
+        Use the file extension to infer the 
+        codec if not specified.
+
+    Returns
+    -------
+    Path
+        The path to the output video file.
+    """
+
+    writer = VideoWriter(file_path=file_path, codec=codec, fps=fps)
+
+    return writer.write_video(frames)
+
+
 class VideoWriter:
     """
     A utility class to create videos from a sequence of image frames.
@@ -92,7 +128,7 @@ def __init__(
         fps : int, default=30
             Frames per second (FPS) for the output video.
         """
-        
+
         self.file_path = Path(file_path)
         self.codec = codec or self._get_codec()
         self.fps = fps
@@ -125,7 +161,7 @@ def add_frame(self, frame: FrameType) -> int:
         self.n_frames += 1
         return self.n_frames
 
-    def generate_video(self, frames: Iterable[FrameType]) -> Path:
+    def write_video(self, frames: Iterable[FrameType]) -> Path:
         """
         Generate a video from a list of frames.
 

docs/index.rst

@@ -64,8 +64,9 @@ Modules
 DMF Utils is designed in a modular way, allowing you to install only the components needed for your specific project. The available modules include:
 
 - :doc:`modules/alerts`: Tools for sending notifications and alerts via Slack and Telegram.
-- :doc:`modules/io`: Input/output utilities for file handling and data management.
 - :doc:`modules/env`: Facities to manage environemt variables.
+- :doc:`modules/io`: Input/output utilities for file handling and data management.
+- :doc:`modules/video`: Utilities for reading and writing video files.
 
 
 For more detailed information about each module, see the :doc:`modules/index` section.

docs/modules/index.rst

@@ -9,5 +9,6 @@ Below are the different modules available within the package:
    :maxdepth: 1
 
    alerts
-   io
    env
+   io
+   video