Add png output support for opengl

Author: ryan

manim/_config/utils.py

@@ -1108,8 +1108,7 @@ def dry_run(self):
             self.write_to_movie is False
             and self.write_all is False
             and self.save_last_frame is False
-            and self.save_pngs is False
-            and self.save_as_gif is False
+            and not self.format
         )
 
     @dry_run.setter
@@ -1118,8 +1117,7 @@ def dry_run(self, val: bool) -> None:
             self.write_to_movie = False
             self.write_all = False
             self.save_last_frame = False
-            self.save_pngs = False
-            self.save_as_gif = False
+            self.format = None
         else:
             raise ValueError(
                 "It is unclear what it means to set dry_run to "

manim/renderer/opengl_renderer.py

@@ -242,7 +242,7 @@ def init_scene(self, scene):
         )
         self.scene = scene
         if not hasattr(self, "window"):
-            if config["preview"]:
+            if config["preview"] and not config["format"] == "png":
                 from .opengl_renderer_window import Window
 
                 self.window = Window(self)
@@ -392,8 +392,7 @@ def render(self, scene, frame_offset, moving_mobjects):
         if self.skip_animations:
             return
 
-        if config["write_to_movie"]:
-            self.file_writer.write_frame(self)
+        self.file_writer.write_frame(self)
 
         if self.window is not None:
             self.window.swap_buffers()
@@ -417,7 +416,11 @@ def update_frame(self, scene):
         self.animation_elapsed_time = time.time() - self.animation_start_time
 
     def scene_finished(self, scene):
-        if config["save_last_frame"]:
+        # When num_plays is 0, no images have been output, so output a single
+        # image in this case
+        if config["save_last_frame"] or (
+            config["format"] == "png" and self.num_plays == 0
+        ):
             self.update_frame(scene)
             self.file_writer.save_final_image(self.get_image())
         self.file_writer.finish()
@@ -435,10 +438,11 @@ def get_image(self) -> Image.Image:
         PIL.Image
             The PIL image of the array.
         """
+        raw_buffer_data = self.get_raw_frame_buffer_object_data()
         image = Image.frombytes(
             "RGBA",
             self.get_pixel_shape(),
-            self.context.fbo.read(self.get_pixel_shape(), components=4),
+            raw_buffer_data,
             "raw",
             "RGBA",
             0,

manim/scene/scene_file_writer.py

@@ -284,25 +284,37 @@ def write_frame(self, frame_or_renderer):
             Pixel array of the frame.
         """
         if config.renderer == "opengl":
-            renderer = frame_or_renderer
-            self.writing_process.stdin.write(
-                renderer.get_raw_frame_buffer_object_data()
-            )
+            self.write_opengl_frame(frame_or_renderer)
         else:
             frame = frame_or_renderer
             if write_to_movie():
                 self.writing_process.stdin.write(frame.tobytes())
             if is_png_format() and not config["dry_run"]:
-                target_dir, extension = os.path.splitext(self.image_file_path)
-                if config["zero_pad"]:
-                    Image.fromarray(frame).save(
-                        f"{target_dir}{str(self.frame_count).zfill(config['zero_pad'])}{extension}"
-                    )
-                else:
-                    Image.fromarray(frame).save(
-                        f"{target_dir}{self.frame_count}{extension}"
-                    )
-                self.frame_count += 1
+                self.output_image_from_array(frame)
+
+    def write_opengl_frame(self, renderer):
+        if write_to_movie():
+            self.writing_process.stdin.write(
+                renderer.get_raw_frame_buffer_object_data()
+            )
+        elif is_png_format() and not config["dry_run"]:
+            target_dir, extension = os.path.splitext(self.image_file_path)
+            self.output_image(
+                renderer.get_image(), target_dir, extension, config["zero_pad"]
+            )
+
+    def output_image_from_array(self, frame_data):
+        target_dir, extension = os.path.splitext(self.image_file_path)
+        self.output_image(
+            Image.fromarray(frame_data), target_dir, extension, config["zero_pad"]
+        )
+
+    def output_image(self, image: Image.Image, target_dir, ext, zero_pad: bool):
+        if zero_pad:
+            image.save(f"{target_dir}{str(self.frame_count).zfill(zero_pad)}{ext}")
+        else:
+            image.save(f"{target_dir}{self.frame_count}{ext}")
+        self.frame_count += 1
 
     def save_final_image(self, image):
         """

tests/test_config.py

@@ -141,19 +141,15 @@ def test_temporary_dry_run():
 
 def test_dry_run_with_png_format():
     """Test that there are no exceptions when running a png without output"""
-    with tempconfig(
-        {"write_to_movie": False, "disable_caching": True, "format": "png"}
-    ):
+    with tempconfig({"write_to_movie": False, "disable_caching": True}):
         assert config["dry_run"] is True
         scene = MyScene()
         scene.render()
 
 
 def test_dry_run_with_png_format_skipped_animations():
     """Test that there are no exceptions when running a png without output and skipped animations"""
-    with tempconfig(
-        {"write_to_movie": False, "disable_caching": True, "format": "png"}
-    ):
+    with tempconfig({"write_to_movie": False, "disable_caching": True}):
         assert config["dry_run"] is True
         scene = MyScene(skip_animations=True)
         scene.render()

tests/test_scene_rendering/test_cli_flags.py

@@ -121,6 +121,32 @@ def test_s_flag(tmp_path, manim_cfg_file, simple_scenes_path):
     assert not is_empty, "running manim with -s flag did not render an image"
 
 
+@pytest.mark.slow
+def test_s_flag_opengl_renderer(tmp_path, manim_cfg_file, simple_scenes_path):
+    scene_name = "SquareToCircle"
+    command = [
+        sys.executable,
+        "-m",
+        "manim",
+        "-ql",
+        "-s",
+        "--renderer",
+        "opengl",
+        "--media_dir",
+        str(tmp_path),
+        simple_scenes_path,
+        scene_name,
+    ]
+    out, err, exit_code = capture(command)
+    assert exit_code == 0, err
+
+    exists = (tmp_path / "videos").exists()
+    assert not exists, "running manim with -s flag rendered a video"
+
+    is_empty = not any((tmp_path / "images" / "simple_scenes").iterdir())
+    assert not is_empty, "running manim with -s flag did not render an image"
+
+
 @pytest.mark.slow
 def test_r_flag(tmp_path, manim_cfg_file, simple_scenes_path):
     scene_name = "SquareToCircle"
@@ -365,6 +391,33 @@ def test_images_are_created_when_png_format_set(
     assert expected_png_path.exists(), "png file not found at " + str(expected_png_path)
 
 
+@pytest.mark.slow
+def test_images_are_created_when_png_format_set_for_opengl(
+    tmp_path, manim_cfg_file, simple_scenes_path
+):
+    """Test images are created in media directory when --format png is set for opengl"""
+    scene_name = "SquareToCircle"
+    command = [
+        sys.executable,
+        "-m",
+        "manim",
+        "-ql",
+        "--renderer",
+        "opengl",
+        "--media_dir",
+        str(tmp_path),
+        "--format",
+        "png",
+        simple_scenes_path,
+        scene_name,
+    ]
+    out, err, exit_code = capture(command)
+    assert exit_code == 0, err
+
+    expected_png_path = tmp_path / "images" / "simple_scenes" / "SquareToCircle0000.png"
+    assert expected_png_path.exists(), "png file not found at " + str(expected_png_path)
+
+
 @pytest.mark.slow
 def test_images_are_zero_padded_when_zero_pad_set(
     tmp_path, manim_cfg_file, simple_scenes_path
@@ -397,6 +450,40 @@ def test_images_are_zero_padded_when_zero_pad_set(
     assert expected_png_path.exists(), "png file not found at " + str(expected_png_path)
 
 
+@pytest.mark.slow
+def test_images_are_zero_padded_when_zero_pad_set_for_opengl(
+    tmp_path, manim_cfg_file, simple_scenes_path
+):
+    """Test images are zero padded when --format png and --zero_pad n are set with the opengl renderer"""
+    scene_name = "SquareToCircle"
+    command = [
+        sys.executable,
+        "-m",
+        "manim",
+        "-ql",
+        "--renderer",
+        "opengl",
+        "--media_dir",
+        str(tmp_path),
+        "--format",
+        "png",
+        "--zero_pad",
+        "3",
+        simple_scenes_path,
+        scene_name,
+    ]
+    out, err, exit_code = capture(command)
+    assert exit_code == 0, err
+
+    unexpected_png_path = tmp_path / "images" / "simple_scenes" / "SquareToCircle0.png"
+    assert not unexpected_png_path.exists(), "non zero padded png file found at " + str(
+        unexpected_png_path
+    )
+
+    expected_png_path = tmp_path / "images" / "simple_scenes" / "SquareToCircle000.png"
+    assert expected_png_path.exists(), "png file not found at " + str(expected_png_path)
+
+
 @pytest.mark.slow
 def test_webm_format_output(tmp_path, manim_cfg_file, simple_scenes_path):
     """Test only webm created when --format webm is set"""