Merge pull request #26 from ahuang11/fix_max_frames

Fix max frames

Author: ahuang11

.github/workflows/build.yml

@@ -1,6 +1,6 @@
 name: Build
 
-on: [push, pull_request]
+on: [pull_request]
 
 jobs:
   test:

docs/supported_formats.md

@@ -2,7 +2,7 @@
 
 StreamJoy supports a variety of input types!
 
-## 📋 List of Images or URLs
+## 📋 List of Images, GIFs, Videos, or URLs
 
 ```python
 from streamjoy import stream
@@ -15,7 +15,7 @@ stream([URL_FMT.format(day=day) for day in range(1, 31)], uri="2024_jan_sea_ice.
 <source src="https://github.com/ahuang11/streamjoy/assets/15331990/7c933cd4-aa15-461a-af79-f508d9d76aa5" type="video/mp4">
 </video>
 
-## 📁 Directory of Images or URLs
+## 📁 Directory of Images, GIFs, Videos, or URLs
 
 ```python
 from streamjoy import stream
@@ -29,6 +29,20 @@ stream(URL_DIR, uri="air_temperature.mp4", pattern="air.sig995.194*.nc")
 <source src="https://github.com/ahuang11/streamjoy/assets/15331990/93cb0c1b-46d3-48e6-be2c-e3b1487f9117" type="video/mp4">
 </video>
 
+## 🧮 Numpy Array
+
+```python
+from streamjoy import stream
+import imageio.v3 as iio
+
+array = iio.imread("imageio:newtonscradle.gif")  # is a 4D numpy array
+stream(array, max_frames=-1).write("newtonscradle.mp4")
+```
+
+<video controls="true" allowfullscreen="true">
+<source src="https://github.com/ahuang11/streamjoy/assets/15331990/7687e951-654c-4719-b50a-4aabc0ddf2e4" type="video/mp4">
+</video>
+
 ## 🐼 Pandas DataFrame or Series
 
 ```python

streamjoy/_utils.py

@@ -327,14 +327,6 @@ def imread_with_pause(
 def subset_resources_renderer_iterables(
     resources: Any, renderer_iterables: list[Any], max_frames: int
 ):
-    if len(resources) > max_frames and max_frames != -1:
-        color = config["logging_warning_color"]
-        reset = config["logging_reset_color"]
-        logging.warning(
-            f"There are a total of {len(resources)} frames, "
-            f"but streaming only {color}{max_frames}{reset}. "
-            f"Set max_frames to -1 to stream all frames."
-        )
     resources = resources[: max_frames or max_frames]
     renderer_iterables = [
         iterable[: len(resources)] for iterable in renderer_iterables or []

streamjoy/serializers.py

@@ -48,21 +48,50 @@ def _select_obj_handler(resources: Any) -> MediaStream:
     if isinstance(resources, (Path, str)):
         return serialize_paths
 
+    resources_type = type(resources)
+    module = getattr(resources_type, "__module__").split(".", maxsplit=1)[0]
+    type_ = resources_type.__name__
     for class_or_package_name, function_name in obj_handlers.items():
-        module = getattr(resources, "__module__", "").split(".", maxsplit=1)[0]
-        type_ = type(resources).__name__
         if (
             f"{module}.{type_}" == class_or_package_name
             or module == class_or_package_name
         ):
             return globals()[function_name]
 
     raise ValueError(
-        f"Could not find a method to handle {type(resources)}; "
+        f"Could not find a method to handle {resources_type}; "
         f"supported classes/packages are {list(obj_handlers.keys())}."
     )
 
 
+def serialize_numpy(
+    stream_cls,
+    resources: np.ndarray,
+    renderer: Callable | None = None,
+    renderer_iterables: list[Any] | None = None,
+    renderer_kwargs: dict | None = None,
+    **kwargs,
+) -> Serialized:
+    """
+    Serialize numpy arrays for streaming or rendering.
+
+    Args:
+        stream_cls: The class reference used for logging and utility functions.
+        resources: The numpy array to be serialized.
+        renderer: The rendering function to use on the array.
+        renderer_iterables: Additional iterable arguments to pass to the renderer.
+        renderer_kwargs: Additional keyword arguments to pass to the renderer.
+        **kwargs: Additional keyword arguments, including 'dim' and 'var' for xarray selection.
+
+    Returns:
+        A tuple containing the serialized resources, renderer, renderer_iterables, renderer_kwargs, and any additional keyword arguments.
+    """
+    resources = [resource for resource in resources]
+    renderer_kwargs = renderer_kwargs or {}
+    renderer_kwargs.update(_utils.pop_from_cls(stream_cls, kwargs))
+    return Serialized(resources, renderer, renderer_iterables, renderer_kwargs, kwargs)
+
+
 def serialize_xarray(
     stream_cls,
     resources: xr.Dataset | xr.DataArray,
@@ -230,8 +259,8 @@ def serialize_polars(
     groupby = kwargs.get("groupby")
 
     if groupby:
-        group_sizes = resources.groupby(groupby).agg(pl.count()).sort(by="count")
-        total_frames = group_sizes.select(pl.col("count").max()).to_numpy()[0, 0]
+        group_sizes = resources.groupby(groupby).agg(pl.len())
+        total_frames = group_sizes.select(pl.col("len").max()).to_numpy()[0, 0]
     else:
         total_frames = len(resources)
 
@@ -276,6 +305,12 @@ def serialize_polars(
         if "ylabel" not in renderer_kwargs:
             renderer_kwargs["ylabel"] = renderer_kwargs["y"].title().replace("_", " ")
 
+    if kwargs.get("processes"):
+        logging.warning(
+            "Polars (HoloViews) rendering does not support processes; "
+            "setting processes=False."
+        )
+    kwargs["processes"] = False
     return Serialized(
         resources_expanded, renderer, renderer_iterables, renderer_kwargs, kwargs
     )

streamjoy/settings.py

@@ -48,7 +48,7 @@
     "pandas.Series": "serialize_pandas",
     "holoviews": "serialize_holoviews",
     "polars.DataFrame": "serialize_polars",
-    "polars.Series": "serialize_polars",
+    "numpy.ndarray": "serialize_numpy",
 }
 
 file_handlers = {

streamjoy/streams.py

@@ -196,6 +196,25 @@ def __init__(self, resources: list[Any] | None = None, **params) -> None:
 
         super().__init__(**params)
 
+    @classmethod
+    def from_numpy(
+        cls,
+        array: np.ndarray,
+        renderer: Callable | None = None,
+        renderer_iterables: list[Any] | None = None,
+        renderer_kwargs: dict | None = None,
+        **kwargs,
+    ) -> MediaStream:
+        serialized = serialize_appropriately(
+            cls,
+            resources=array,
+            renderer=renderer,
+            renderer_iterables=renderer_iterables,
+            renderer_kwargs=renderer_kwargs,
+            **kwargs,
+        )
+        return cls(**serialized.param.values(), **serialized.kwargs)
+
     @classmethod
     def from_xarray(
         cls,
@@ -538,6 +557,13 @@ def write(
         if resources is None:
             resources = self.resources
             renderer_iterables = self.renderer_iterables
+            max_frames = _utils.get_max_frames(
+                len(resources), kwargs.get("max_frames", self.max_frames)
+            )
+            kwargs["max_frames"] = max_frames
+            resources, renderer_iterables = _utils.subset_resources_renderer_iterables(
+                resources, renderer_iterables, max_frames
+            )
         else:
             serialized = serialize_appropriately(
                 self, resources, renderer, renderer_kwargs, **kwargs

tests/conftest.py

@@ -1,6 +1,7 @@
 from pathlib import Path
 
 import hvplot.xarray  # noqa: F401
+import imageio.v3 as iio
 import pandas as pd
 import polars as pl
 import pytest
@@ -16,6 +17,11 @@
 PARQUET_PATH = DATA_DIR / "gapminder.parquet"
 
 
+@pytest.fixture
+def array():
+    return iio.imread("imageio:newtonscradle.gif")
+
+
 @pytest.fixture
 def ds():
     return xr.open_zarr(ZARR_PATH)

tests/test_polars.py

@@ -0,0 +1,16 @@
+from io import BytesIO
+
+import streamjoy.polars  # noqa: F401
+
+
+class TestPolars:
+    def test_dataframe(self, pl_df):
+        stream = pl_df.streamjoy(groupby="Country")
+        assert stream.renderer_kwargs == {
+            "groupby": "Country",
+            "x": "Year",
+            "y": "fertility",
+            "xlabel": "Year",
+            "ylabel": "Fertility",
+        }
+        assert isinstance(stream.write(), BytesIO)

tests/test_serializers.py

@@ -0,0 +1,68 @@
+from streamjoy.models import Serialized
+from streamjoy.serializers import (
+    serialize_holoviews,
+    serialize_numpy,
+    serialize_pandas,
+    serialize_polars,
+    serialize_xarray,
+)
+
+
+class TestSerializeNumpy:
+    def test_serialize_numpy(self, array):
+        serialized = serialize_numpy(None, array)
+        assert isinstance(serialized, Serialized)
+        assert len(serialized.resources) == 36
+        assert isinstance(serialized.resources, list)
+        assert not serialized.renderer
+        assert serialized.renderer_iterables is None
+        assert isinstance(serialized.renderer_kwargs, dict)
+        assert isinstance(serialized.kwargs, dict)
+
+
+class TestSerializeXarray:
+    def test_serialize_xarray(self, ds):
+        serialized = serialize_xarray(None, ds)
+        assert isinstance(serialized, Serialized)
+        assert len(serialized.resources) == 3
+        assert isinstance(serialized.resources, list)
+        assert callable(serialized.renderer)
+        assert serialized.renderer_iterables is None
+        assert isinstance(serialized.renderer_kwargs, dict)
+        assert isinstance(serialized.kwargs, dict)
+
+
+class TestSerializePandas:
+    def test_serialize_pandas(self, df):
+        serialized = serialize_pandas(None, df)
+        assert isinstance(serialized, Serialized)
+        assert len(serialized.resources) == 3
+        assert isinstance(serialized.resources, list)
+        assert callable(serialized.renderer)
+        assert serialized.renderer_iterables is None
+        assert isinstance(serialized.renderer_kwargs, dict)
+        assert isinstance(serialized.kwargs, dict)
+
+
+class TestSerializePolars:
+    def test_serialize_polars(self, pl_df):
+        serialized = serialize_polars(None, pl_df)
+        assert isinstance(serialized, Serialized)
+        assert len(serialized.resources) == 3
+        assert isinstance(serialized.resources, list)
+        assert callable(serialized.renderer)
+        assert serialized.renderer_iterables is None
+        assert isinstance(serialized.renderer_kwargs, dict)
+        assert isinstance(serialized.kwargs, dict)
+
+
+class TestSerializeHoloviews:
+    def test_serialize_holoviews(self, hmap):
+        serialized = serialize_holoviews(None, hmap)
+        assert isinstance(serialized, Serialized)
+        assert len(serialized.resources) == 20
+        assert isinstance(serialized.resources, list)
+        assert callable(serialized.renderer)
+        assert serialized.renderer_iterables is None
+        assert isinstance(serialized.renderer_kwargs, dict)
+        assert isinstance(serialized.kwargs, dict)

tests/test_streams.py

@@ -14,6 +14,10 @@ def _assert_stream_and_props(self, sj, stream_cls):
         props.n_images == 3
         return props
 
+    def test_from_numpy(self, stream_cls, array):
+        sj = stream_cls.from_numpy(array)
+        self._assert_stream_and_props(sj, stream_cls)
+
     def test_from_pandas(self, stream_cls, df):
         sj = stream_cls.from_pandas(df)
         self._assert_stream_and_props(sj, stream_cls)
@@ -69,6 +73,12 @@ def test_holoviews_bokeh_backend(self, stream_cls, ds):
         props = self._assert_stream_and_props(sj, stream_cls)
         assert props.shape[1] == 300
 
+    def test_write_max_frames(self, stream_cls, df):
+        sj = stream_cls.from_pandas(df, max_frames=3)
+        buf = sj.write(max_frames=2)
+        props = improps(buf)
+        assert props.n_images == 2
+
 
 class TestGifStream(AbstractTestMediaStream):
     @pytest.fixture(scope="class")