Create utils/helpers.py

utils/helpers.py

@@ -0,0 +1,49 @@
+# utils/helpers.py
+import base64
+from io import BytesIO
+from PIL import Image
+import imageio
+import numpy as np
+from typing import List
+
+def encode_image_base64(image: Image.Image, format="PNG") -> str:
+    """Encodes a PIL Image to a base64 string."""
+    buffered = BytesIO()
+    image.save(buffered, format=format)
+    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    return img_str
+
+def decode_base64_image(image_base64: str) -> Image.Image:
+    """Decodes a base64 string to a PIL Image."""
+    img_bytes = base64.b64decode(image_base64)
+    img = Image.open(BytesIO(img_bytes))
+    return img.convert("RGB") # Ensure RGB format
+
+def encode_video_base64(video_frames: List[np.ndarray], fps: int, format="MP4") -> str:
+    """Encodes a list of numpy array frames (H, W, C) to a base64 video string."""
+    if not video_frames:
+        raise ValueError("Cannot encode empty video frames list.")
+
+    buffered = BytesIO()
+    # Ensure frames are uint8
+    frames_uint8 = [(frame * 255).astype(np.uint8) if frame.dtype != np.uint8 else frame for frame in video_frames]
+
+    if format.upper() == "MP4":
+        # imageio requires ffmpeg to be installed in the Space environment
+        # Add 'imageio[ffmpeg]' to requirements.txt if needed, or ensure ffmpeg is available
+        try:
+            imageio.mimwrite(buffered, frames_uint8, format='mp4', fps=fps, output_params=['-vf', 'pad=ceil(iw/2)*2:ceil(ih/2)*2']) # Ensure dimensions are even for some codecs
+            # quality=8 is a reasonable default, adjust if needed. Macro_block size needed for h264
+        except Exception as e:
+             print(f"WARNING: MP4 encoding failed ({e}). Check ffmpeg availability. Falling back to GIF.")
+             # Fallback to GIF if MP4 fails (e.g., ffmpeg not found)
+             imageio.mimwrite(buffered, frames_uint8, format='gif', duration=1000/fps) # duration is in ms
+             format = "GIF"
+
+    elif format.upper() == "GIF":
+        imageio.mimwrite(buffered, frames_uint8, format='gif', duration=1000/fps) # duration is in ms
+    else:
+        raise ValueError(f"Unsupported video format: {format}")
+
+    video_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    return video_str, format.upper() # Return actual format used (might differ on fallback)