utils/helpers.py

# 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)