Save to memory instead of file (#90)

* save to memory

* typo

* more typos

* doesnt need it anymore

.github/workflows/python.yaml

@@ -30,4 +30,4 @@ jobs:
 
       - name: Run additional tests
         working-directory: src/python
-        run: pytest tests -s
+        run: pytest tests

src/python/pose_format/pose_visualizer.py

@@ -2,7 +2,8 @@
 import logging
 import math
 from functools import lru_cache
-from typing import Iterable, Tuple
+from typing import Iterable, Tuple, Union
+from io import BytesIO
 
 import numpy as np
 import numpy.ma as ma
@@ -209,14 +210,14 @@ def save_frame(self, f_name: str, frame: np.ndarray):
         """
         self.cv2.imwrite(f_name, frame)
 
-    def _save_image(self, f_name: str, frames: Iterable[np.ndarray], format: str = "GIF", transparency: bool = False):
+    def _save_image(self, f_name: Union[str, None], frames: Iterable[np.ndarray], format: str = "GIF", transparency: bool = False) -> Union[None, bytes]:
         """
         Save pose frames as Image (GIF or PNG).
 
         Parameters
         ----------
-        f_name : str
-            filename to save Image to.
+        f_name : Union[str, None]
+        	Filename to save Image to. If None, image will be saved to memory and returned as bytes.
         frames : Iterable[np.ndarray]
             Series of pose frames to be included in Image.
         format : str
@@ -226,7 +227,8 @@ def _save_image(self, f_name: str, frames: Iterable[np.ndarray], format: str = "
 
         Returns
         -------
-        None
+        Union[None, bytes]
+        	If f_name is None, returns the image data as bytes. Otherwise, returns None.
 
         Raises
         ------
@@ -244,59 +246,70 @@ def _save_image(self, f_name: str, frames: Iterable[np.ndarray], format: str = "
             cv_code = self.cv2.COLOR_BGR2RGB
 
         images = [Image.fromarray(self.cv2.cvtColor(frame, cv_code)) for frame in frames]
-        images[0].save(f_name,
-                       format=format,
-                       append_images=images[1:],
-                       save_all=True,
-                       duration=1000 / self.pose.body.fps,
-                       loop=0,
-                       disposal=2 if transparency else 0)
-
-    def save_gif(self, f_name: str, frames: Iterable[np.ndarray]):
+
+        def save_to(obj: Union[str, None]):
+            images[0].save(obj,
+							format=format,
+							append_images=images[1:],
+							save_all=True,
+							duration=1000 / self.pose.body.fps,
+							loop=0,
+							disposal=2 if transparency else 0)
+
+        if f_name:
+            save_to(f_name)
+        else:
+            with BytesIO() as mem:
+                save_to(mem)
+                return mem.getvalue()
+
+    def save_gif(self, f_name: Union[str, None], frames: Iterable[np.ndarray]) -> Union[None, bytes]:
         """
         Save pose frames as GIF.
 
         Parameters
         ----------
-        f_name : str
-            filename to save GIF to.
+        f_name : Union[str, None]
+       		Filename to save PNG to. If None, image will be saved to memory and returned as bytes.
         frames : Iterable[np.ndarray]
             Series of pose frames to be included in GIF.
 
         Returns
         -------
-        None
+        Union[None, bytes]
+        	If f_name is None, returns the PNG image data as bytes. Otherwise, returns None.
 
         Raises
         ------
         ImportError 
             If Pillow is not installed.
         """
-        self._save_image(f_name, frames, "GIF", False)
+        return self._save_image(f_name, frames, "GIF", False)
 
-    def save_png(self, f_name: str, frames: Iterable[np.ndarray], transparency: bool = True):
+    def save_png(self, f_name: Union[str, None], frames: Iterable[np.ndarray], transparency: bool = True) -> Union[None, bytes]:
         """
         Save pose frames as PNG.
 
         Parameters
         ----------
-        f_name : str
-            filename to save PNG to.
+        f_name : Union[str, None]
+        	Filename to save PNG to. If None, image will be saved to memory and returned as bytes.
         frames : Iterable[np.ndarray]
             Series of pose frames to be included in PNG.
         transparency : bool
             transparency decides opacity of background color.
 
         Returns
         -------
-        None
+        Union[None, bytes]
+        	If f_name is None, returns the PNG image data as bytes. Otherwise, returns None.
 
         Raises
         ------
         ImportError 
             If Pillow is not installed.
         """        
-        self._save_image(f_name, frames, "PNG", transparency)
+        return self._save_image(f_name, frames, "PNG", transparency)
 
     def save_video(self, f_name: str, frames: Iterable[np.ndarray], custom_ffmpeg=None):
         """

src/python/tests/visualization_test.py

@@ -23,8 +23,7 @@ def test_save_gif(self):
         with tempfile.NamedTemporaryFile(suffix='.gif', delete=False) as temp_gif:
             v.save_gif(temp_gif.name, v.draw())
             self.assertTrue(os.path.exists(temp_gif.name))
-            self.assertGreater(os.path.getsize(
-                temp_gif.name), 0)
+            self.assertGreater(os.path.getsize(temp_gif.name), 0)
 
     def test_save_png(self):
         """
@@ -35,11 +34,10 @@ def test_save_png(self):
 
         v = PoseVisualizer(pose)
 
-        with tempfile.TemporaryDirectory() as temp_dir:
-            temp_png = os.path.join(temp_dir, 'example.png')
-            v.save_png(temp_png, v.draw(transparency=True))
-            self.assertTrue(os.path.exists(temp_png))
-            self.assertGreater(os.path.getsize(temp_png), 0)
+        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_png:
+            v.save_png(temp_png.name, v.draw(transparency=True))
+            self.assertTrue(os.path.exists(temp_png.name))
+            self.assertGreater(os.path.getsize(temp_png.name), 0)
 
     def test_save_mp4(self):
         """
@@ -54,3 +52,14 @@ def test_save_mp4(self):
             v.save_video(temp_mp4.name, v.draw())
             self.assertTrue(os.path.exists(temp_mp4.name))
             self.assertGreater(os.path.getsize(temp_mp4.name), 0)
+
+    def test_save_to_memory(self):
+        """
+        Test saving pose visualization as bytes.
+        """
+        with open("tests/data/mediapipe_long_hand_normalized.pose", "rb") as f:
+            pose = Pose.read(f.read())
+
+        v = PoseVisualizer(pose)
+        file_bytes = v.save_png(None, v.draw())
+        self.assertGreater(len(file_bytes), 0)