[BUG FIX] Fix 'gs.tools.animate' dropping last frame by migrating from MoviePy to PyAV.

genesis/utils/tools.py

@@ -13,33 +13,114 @@ def animate(imgs, filename=None, fps=60):
     """
     Create a video from a list of images.
 
+    Images must be uint8 arrays of shape ``(H, W, 3)`` (RGB), ``(H, W, 4)`` (RGBA,
+    alpha stripped automatically), or ``(H, W)`` (grayscale). PIL Images are also
+    accepted. Float arrays are *not* supported; convert to uint8 before calling
+    (e.g. ``(img * 255).astype(np.uint8)``).
+
+    All frames must share the same height, width, and channel count as the first
+    frame; a ``ValueError`` is raised immediately if a mismatch is detected.
+
     Args:
-        imgs (list): List of input images.
-        filename (str, optional): Name of the output video file. If not provided, the name will be default to the name of the caller file, with a timestamp and '.mp4' extension.
+        imgs (list): List of input images (numpy arrays or PIL Images).
+        filename (str, optional): Output video path (.mp4). Defaults to
+            ``<caller_script>_<timestamp>.mp4`` in the current directory.
+        fps (int, optional): Frames per second. Defaults to 60.
     """
     assert isinstance(imgs, list)
     if len(imgs) == 0:
         gs.logger.warning("No image to save.")
         return
 
+    fps = max(1, int(round(fps)))  # PyAV requires an integer time-base denominator
+
     if filename is None:
         caller_file = inspect.stack()[-1].filename
         # caller file + timestamp + .mp4
         filename = os.path.splitext(os.path.basename(caller_file))[0] + f"_{time.strftime('%Y%m%d_%H%M%S')}.mp4"
-    os.makedirs(os.path.abspath(os.path.dirname(filename)), exist_ok=True)
+    os.makedirs(os.path.abspath(os.path.dirname(filename) or "."), exist_ok=True)
 
     gs.logger.info(f'Saving video to ~<"{filename}">~...')
-    from moviepy import ImageSequenceClip
-
-    imgs = ImageSequenceClip(imgs, fps=fps)
-    imgs.write_videofile(
-        filename,
-        fps=fps,
-        logger=None,
-        codec="libx264",
-        preset="ultrafast",
-        # ffmpeg_params=["-crf", "0"],
-    )
+
+    _av_ok = False
+    try:
+        import av
+
+        # libx264 must be compiled into this PyAV build; fall back to moviepy otherwise.
+        if "libx264" not in av.codecs_available:
+            raise ImportError("PyAV build does not include libx264")
+
+        first = imgs[0]
+        if not isinstance(first, np.ndarray):
+            first = np.array(first)
+        if not np.issubdtype(first.dtype, np.unsignedinteger) or first.dtype != np.uint8:
+            raise ValueError(
+                f"animate() requires uint8 images; got dtype={first.dtype}. "
+                "Convert float images with (img * 255).astype(np.uint8) before calling."
+            )
+        # Strip alpha channel if present; libx264/yuv420p only accepts RGB or grayscale.
+        if first.ndim == 3 and first.shape[2] == 4:
+            first = first[..., :3]
+        if first.ndim not in (2, 3) or (first.ndim == 3 and first.shape[2] != 3):
+            raise ValueError(
+                f"animate() requires images of shape (H, W), (H, W, 3), or (H, W, 4); "
+                f"got shape={first.shape}."
+            )
+        height, width = first.shape[:2]
+        is_color = first.ndim == 3 and first.shape[2] == 3
+        fmt = "rgb24" if is_color else "gray"
+
+        container = av.open(filename, mode="w")
+        try:
+            stream = container.add_stream("libx264", rate=fps)
+            stream.width = width
+            stream.height = height
+            stream.pix_fmt = "yuv420p"
+            stream.codec_context.options = {"preset": "ultrafast"}
+
+            for i, img in enumerate(imgs):
+                if not isinstance(img, np.ndarray):
+                    img = np.array(img)
+                if img.dtype != np.uint8:
+                    raise ValueError(
+                        f"animate() requires uint8 images; frame {i} has dtype={img.dtype}."
+                    )
+                # Strip alpha channel for consistency with `first`.
+                if img.ndim == 3 and img.shape[2] == 4:
+                    img = img[..., :3]
+                if img.shape[:2] != (height, width) or img.ndim != first.ndim:
+                    raise ValueError(
+                        f"Frame {i} shape {img.shape} does not match first frame shape "
+                        f"{first.shape} (after alpha strip). All frames must be identical in "
+                        "size and channel count."
+                    )
+                # from_ndarray handles stride/padding internally, avoiding the
+                # line_size // channels reshape bug for non-aligned widths.
+                frame = av.VideoFrame.from_ndarray(img, format=fmt)
+                for packet in stream.encode(frame):
+                    container.mux(packet)
+
+            for packet in stream.encode(None):
+                container.mux(packet)
+        finally:
+            container.close()
+
+        _av_ok = True
+
+    except Exception as exc:
+        if isinstance(exc, ValueError):
+            raise
+        gs.logger.warning(
+            f"PyAV unavailable or failed ({type(exc).__name__}: {exc}); falling back to moviepy. "
+            "Note: moviepy ≥ 2.x may drop the last frame. Install 'av' for reliable output."
+        )
+
+    if not _av_ok:
+        from moviepy import ImageSequenceClip
+
+        clip = ImageSequenceClip(imgs, fps=fps)
+        clip.write_videofile(filename, fps=fps, logger=None, codec="libx264", preset="ultrafast")
+
     gs.logger.info("Video saved.")
 
 

tests/test_utils.py

@@ -1,4 +1,6 @@
 import math
+import os
+import tempfile
 from functools import partial
 from unittest.mock import patch
 
@@ -10,7 +12,7 @@
 
 import genesis as gs
 import genesis.utils.geom as gu
-from genesis.utils.tools import FPSTracker
+from genesis.utils.tools import FPSTracker, animate
 from genesis.utils.misc import tensor_to_array
 from genesis.utils import warnings as warnings_mod
 from genesis.utils.warnings import warn_once
@@ -331,6 +333,129 @@ def test_geom_tensor_identity(batch_shape):
         np.testing.assert_allclose(tensor_to_array(tc_args[0]), tensor_to_array(tc_args[-1]), atol=1e2 * gs.EPS)
 
 
+@pytest.mark.required
+def test_animate_frame_count_rgb():
+    """animate() must encode every frame; no frame should be dropped."""
+    av = pytest.importorskip("av", reason="PyAV not installed")
+    if "libx264" not in av.codecs_available:
+        pytest.skip("PyAV build does not include libx264")
+
+    n_frames = 5
+    imgs = [np.full((8, 8, 3), i * 50, dtype=np.uint8) for i in range(n_frames)]
+    with tempfile.TemporaryDirectory() as tmp:
+        path = os.path.join(tmp, "out.mp4")
+        animate(imgs, filename=path, fps=30)
+        container = av.open(path)
+        frames = list(container.decode(video=0))
+        container.close()
+    assert len(frames) == n_frames, f"Expected {n_frames} frames, got {len(frames)}"
+
+
+@pytest.mark.required
+def test_animate_frame_count_grayscale():
+    """animate() must encode every grayscale frame without dropping any."""
+    av = pytest.importorskip("av", reason="PyAV not installed")
+    if "libx264" not in av.codecs_available:
+        pytest.skip("PyAV build does not include libx264")
+
+    n_frames = 4
+    imgs = [np.full((10, 10), i * 60, dtype=np.uint8) for i in range(n_frames)]
+    with tempfile.TemporaryDirectory() as tmp:
+        path = os.path.join(tmp, "gray.mp4")
+        animate(imgs, filename=path, fps=24)
+        container = av.open(path)
+        frames = list(container.decode(video=0))
+        container.close()
+    assert len(frames) == n_frames, f"Expected {n_frames} frames, got {len(frames)}"
+
+
+@pytest.mark.required
+def test_animate_rejects_float_input():
+    """animate() must raise ValueError for float images, not silently truncate."""
+    av = pytest.importorskip("av", reason="PyAV not installed")
+    if "libx264" not in av.codecs_available:
+        pytest.skip("PyAV build does not include libx264")
+
+    imgs = [np.zeros((8, 8, 3), dtype=np.float32)]
+    with tempfile.TemporaryDirectory() as tmp:
+        with pytest.raises(ValueError, match="uint8"):
+            animate(imgs, filename=os.path.join(tmp, "out.mp4"))
+
+
+@pytest.mark.required
+def test_animate_rejects_mismatched_frame_shape():
+    """animate() must raise ValueError when a later frame has a different shape."""
+    av = pytest.importorskip("av", reason="PyAV not installed")
+    if "libx264" not in av.codecs_available:
+        pytest.skip("PyAV build does not include libx264")
+
+    imgs = [
+        np.zeros((8, 8, 3), dtype=np.uint8),
+        np.zeros((16, 8, 3), dtype=np.uint8),  # different height
+    ]
+    with tempfile.TemporaryDirectory() as tmp:
+        with pytest.raises(ValueError, match="shape"):
+            animate(imgs, filename=os.path.join(tmp, "out.mp4"))
+
+
+@pytest.mark.required
+def test_animate_rgba_strips_alpha():
+    """animate() must accept RGBA inputs and produce the correct frame count."""
+    av = pytest.importorskip("av", reason="PyAV not installed")
+    if "libx264" not in av.codecs_available:
+        pytest.skip("PyAV build does not include libx264")
+
+    n_frames = 3
+    imgs = [np.full((8, 8, 4), i * 80, dtype=np.uint8) for i in range(n_frames)]
+    with tempfile.TemporaryDirectory() as tmp:
+        path = os.path.join(tmp, "rgba.mp4")
+        animate(imgs, filename=path, fps=30)
+        container = av.open(path)
+        frames = list(container.decode(video=0))
+        container.close()
+    assert len(frames) == n_frames
+
+
+@pytest.mark.required
+def test_animate_av_error_falls_back_to_moviepy(monkeypatch):
+    """When PyAV raises a non-ValueError (e.g. FFmpegError), animate() falls back to moviepy."""
+    pytest.importorskip("av", reason="PyAV not installed")
+
+    import av as av_mod
+    from unittest.mock import MagicMock, patch
+
+    # Make av.open raise a generic RuntimeError to simulate an av encode failure.
+    with patch.object(av_mod, "open", side_effect=RuntimeError("simulated av failure")):
+        moviepy_called = []
+
+        class FakeClip:
+            def write_videofile(self, *a, **kw):
+                moviepy_called.append(True)
+
+        with patch("genesis.utils.tools.ImageSequenceClip", return_value=FakeClip(), create=True):
+            # Also ensure the import inside animate() works without calling real moviepy
+            with patch.dict("sys.modules", {"moviepy": MagicMock()}):
+                import sys
+
+                sys.modules["moviepy"].ImageSequenceClip = lambda imgs, fps: FakeClip()
+                with tempfile.TemporaryDirectory() as tmp:
+                    animate(
+                        [np.full((8, 8, 3), 128, dtype=np.uint8)],
+                        filename=os.path.join(tmp, "out.mp4"),
+                    )
+    assert moviepy_called, "moviepy fallback was not triggered after av failure"
+
+
+@pytest.mark.required
+def test_animate_value_error_propagates_through_av_block():
+    """ValueError from input validation must propagate even if av is available."""
+    pytest.importorskip("av", reason="PyAV not installed")
+    imgs = [np.zeros((8, 8, 3), dtype=np.float32)]  # wrong dtype
+    with tempfile.TemporaryDirectory() as tmp:
+        with pytest.raises(ValueError, match="uint8"):
+            animate(imgs, filename=os.path.join(tmp, "out.mp4"))
+
+
 def test_fps_tracker():
     n_envs = 23
     tracker = FPSTracker(alpha=0.0, minimum_interval_seconds=0.1, n_envs=n_envs)