Movie storage

.github/workflows/tests_all.yml

@@ -12,11 +12,19 @@ jobs:
 
     steps:
     - uses: actions/checkout@v3
+
     - name: Set up Python
       uses: actions/setup-python@v4
       with:
         python-version: '3.12'
 
+    - uses: FedericoCarboni/setup-ffmpeg@v3
+      # install ffmpeg as special requirement
+      id: setup-ffmpeg
+      with:
+        ffmpeg-version: release
+        github-token: ${{ github.server_url == 'https://github.com' && github.token || '' }}
+
     - name: Install dependencies
       # install all requirements
       run: |

docs/source/_static/requirements_optional.csv

@@ -1,4 +1,5 @@
 Package,Minimal version,Usage
+ffmpeg-python,0.2,Reading and writing videos
 h5py,2.10,Storing data in the hierarchical file format
 ipywidgets,8,Jupyter notebook support
 mpi4py,3,Parallel processing using MPI

pde/storage/__init__.py

@@ -7,9 +7,11 @@
    ~memory.get_memory_storage
    ~memory.MemoryStorage
    ~file.FileStorage
+   ~movie.MovieStorage
 
 .. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de> 
 """
 
 from .file import FileStorage
 from .memory import MemoryStorage, get_memory_storage
+from .movie import MovieStorage

pde/storage/_ffmpeg.py

@@ -0,0 +1,79 @@
+"""
+Functions for interacting with FFmpeg
+
+.. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de>
+"""
+
+# import subprocess as sp
+from dataclasses import dataclass
+
+import numpy as np
+from numpy.typing import DTypeLike
+
+# def _run_ffmpeg(args: list[str]):
+#     return sp.check_output(["ffmpeg"] + args)
+#
+#
+# def codecs() -> list[str]:
+#     """list: all supported ffmpeg codecs"""
+#     res = _run_ffmpeg(["-codecs"])
+#
+#
+# def get_pixel_formats(encoder=None):
+#     if encoder is None:
+#         res = _run_ffmpeg(["-pix_fmts"])
+#     else:
+#         res = _run_ffmpeg(["-h", f"encoder={encoder}"])
+
+
+@dataclass
+class FFmpegFormat:
+    """defines a FFmpeg format used for storing field data in a video"""
+
+    pix_fmt_file: str
+    codec: str
+    channels: int
+    value_max: int
+    dtype: DTypeLike
+
+    @property
+    def pix_fmt_data(self) -> str:
+        return {1: "gray", 3: "rgb24", 4: "rgba"}[self.channels]
+
+    def data_to_frame(self, normalized_data: np.ndarray) -> np.ndarray:
+        return (normalized_data * self.value_max).astype(self.dtype)
+
+    def data_from_frame(self, frame_data: np.ndarray):
+        return frame_data.astype(float) / self.value_max
+
+
+formats = {
+    "gray": FFmpegFormat(
+        pix_fmt_file="gray",
+        codec="libx264",
+        channels=1,
+        value_max=255,
+        dtype=np.uint8,
+    ),
+    "rgb24": FFmpegFormat(
+        pix_fmt_file="rgb24",
+        codec="libx264rgb",
+        channels=3,
+        value_max=255,
+        dtype=np.uint8,
+    ),
+    "bgr24": FFmpegFormat(
+        pix_fmt_file="bgr24",
+        codec="libx264rgb",
+        channels=3,
+        value_max=255,
+        dtype=np.uint8,
+    ),
+    "rgb32": FFmpegFormat(
+        pix_fmt_file="rgb32",
+        codec="libx264rgb",
+        channels=4,
+        value_max=255,
+        dtype=np.uint8,
+    ),
+}

pde/storage/file.py

@@ -25,7 +25,7 @@ class FileStorage(StorageBase):
 
     def __init__(
         self,
-        filename: str,
+        filename: str | Path,
         *,
         info: InfoDict | None = None,
         write_mode: WriteModeType = "truncate_once",