diff --git a/README.md b/README.md index 4eadb94..322d807 100644 --- a/README.md +++ b/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. diff --git a/dmf/__init__.py b/dmf/__init__.py index e3d4d99..78e1cd2 100644 --- a/dmf/__init__.py +++ b/dmf/__init__.py @@ -4,7 +4,7 @@ from .__version__ import __version__ -subpackages = ["alerts", "io", "env"] +subpackages = ["alerts", "io", "env", "video"] __getattr__, __dir__, __all__ = lazy.attach(__name__, subpackages) @@ -12,6 +12,7 @@ from . import alerts from . import io from . import env + from . import video __all__ = ["__version__", "alerts", "io", "env"] diff --git a/dmf/__version__.py b/dmf/__version__.py index c7bf681..4f59c08 100644 --- a/dmf/__version__.py +++ b/dmf/__version__.py @@ -1,2 +1,2 @@ -__version__ = "0.1.1" \ No newline at end of file +__version__ = "0.1.2" \ No newline at end of file diff --git a/dmf/io/__init__.py b/dmf/io/__init__.py index 038492a..a0327f3 100644 --- a/dmf/io/__init__.py +++ b/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"] \ No newline at end of file +__all__ = ["compress", "decompress", "load", "save"] \ No newline at end of file diff --git a/dmf/io/save.py b/dmf/io/save.py index 80feb27..0eae233 100644 --- a/dmf/io/save.py +++ b/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) \ No newline at end of file + from ..video.video_writer import write_video + write_video(file_path=file_path, frames=data, **kwargs) \ No newline at end of file diff --git a/dmf/video/__init__.py b/dmf/video/__init__.py new file mode 100644 index 0000000..902a534 --- /dev/null +++ b/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"] diff --git a/dmf/video/video_reader.py b/dmf/video/video_reader.py new file mode 100644 index 0000000..75c093f --- /dev/null +++ b/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() diff --git a/dmf/io/video.py b/dmf/video/video_writer.py similarity index 86% rename from dmf/io/video.py rename to dmf/video/video_writer.py index f1d71ad..056b203 100644 --- a/dmf/io/video.py +++ b/dmf/video/video_writer.py @@ -13,7 +13,8 @@ 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: @@ -21,6 +22,8 @@ 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. diff --git a/docs/index.rst b/docs/index.rst index c7c9654..e2f7b55 100644 --- a/docs/index.rst +++ b/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. diff --git a/docs/modules/index.rst b/docs/modules/index.rst index c94ec7a..b827b84 100644 --- a/docs/modules/index.rst +++ b/docs/modules/index.rst @@ -9,5 +9,6 @@ Below are the different modules available within the package: :maxdepth: 1 alerts - io env + io + video diff --git a/docs/modules/io.rst b/docs/modules/io.rst index 19e0764..f0a165d 100644 --- a/docs/modules/io.rst +++ b/docs/modules/io.rst @@ -172,32 +172,4 @@ Examples from dmf.io import decompress decompress("my_folder.zip") - -Other Utilities ---------------- - -In addition to saving, loading, and compression, the IO module includes utilities such as `VideoWriter`, which can be used to create videos from image frames. - -.. autosummary:: - :toctree: autosummary - - dmf.io.VideoWriter - - -Examples -~~~~~~~~ - -**Creating a Video from Image Frames**: - -.. code-block:: python - - import cv2 - from dmf.io.video import VideoWriter - - # Initialize the VideoWriter - with VideoWriter("output.mp4", fps=30) as writer: - for i in range 100): - frame = cv2.imread(f"frame_{i}.png") - writer.add_frame(frame) - -This will create a video file `output.mp4` from a sequence of image frames. + \ No newline at end of file diff --git a/docs/modules/video.rst b/docs/modules/video.rst new file mode 100644 index 0000000..10194d6 --- /dev/null +++ b/docs/modules/video.rst @@ -0,0 +1,89 @@ +Video +===== + +The `video` module in DMF Utils provides utilities for reading and writing video files. It offers functions to easily handle video frames, allowing you to write videos from frames and read videos into various formats such as NumPy arrays or PIL images. + +This module is included in the base package: + +.. code-block:: bash + + pip install dmf-utils[video] + +Overview +-------- + +The `video` module allows you to: + +- Write videos from frames using `VideoWriter` or the `write_video` function. +- Read videos into frames as either NumPy arrays or PIL images using `VideoReader` or the `read_video` function. +- Handle video processing tasks easily with support for common video formats. + +Video Functions and Classes +--------------------------- + +Functions and classes included in this module: + +.. autosummary:: + :toctree: autosummary + + dmf.video.write_video + dmf.video.read_video + dmf.video.VideoWriter + dmf.video.VideoReader + +Examples +-------- + +**Writing a Video**: + +.. code-block:: python + + from dmf.video import write_video + + # Example: Writing a video from NumPy arrays + import numpy as np + frames = [np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8) for _ in range(10)] + write_video("output.mp4", frames, fps=30) + + # Example: Writing a video from image paths + frames = ["frame1.png", "frame2.png", "frame3.png"] + write_video("output.mp4", frames, fps=30) + +**Reading a Video**: + +.. code-block:: python + + from dmf.video import read_video + + # Example: Reading a video as a NumPy array + frames = read_video("input.mp4", output_type="numpy") + print(frames.shape) # (num_frames, height, width, 3) + + # Example: Reading a video as a list of PIL images + frames = read_video("input.mp4", output_type="pil") + print(len(frames)) # Number of frames in the video + +Classes +------- + +**VideoWriter**: + +.. code-block:: python + + from dmf.video import VideoWriter + + frames = ["frame1.png", "frame2.png", "frame3.png"] + with VideoWriter("output.mp4", fps=30) as writer: + for frame in frames: + writer.add_frame(frame) + +**VideoReader**: + +.. code-block:: python + + from dmf.video import VideoReader + + with VideoReader("input.mp4", output_type="pil") as reader: + for frame in reader: + # Process each frame + pass diff --git a/pyproject.toml b/pyproject.toml index c3886ca..c4368a2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,21 +64,27 @@ alerts = [ "slack_sdk", "requests", ] +video = [ + "numpy", + "opencv-python", + "pillow", +] all = [ "sphinx", "pydata-sphinx-theme", "slack_sdk", "requests", + "numpy", + "opencv-python", + "pillow", ] extra = [ "pandas", "h5py", - "pillow", "torch", "pyyaml", "scipy", "librosa", - "opencv-python", "matplotlib", "py7zr", ]