utils.py

import os
import base64
import json
from PIL import Image
from typing import List, Dict
import cv2  # CHANGE: Added for video processing


def print_png_size_in_mb(file_path):
    if not file_path.lower().endswith('.png'):
        print("File is not a PNG image.")
        return
    try:
        with Image.open(file_path) as img:
            img.verify()
        size_bytes = os.path.getsize(file_path)
        size_mb = size_bytes / (1024 * 1024)
        print(f"Image size: {size_mb:.2f} MB")
    except Exception as e:
        print(f"Error reading image: {e}")


def get_image_size_from_bytes(image_bytes: bytes):
    size_bytes = len(image_bytes)
    size_mb = size_bytes/(1024*1024)
    return size_mb


def image_file_to_base64(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")

# ...existing code...


def encode_image_bytes(image_bytes: bytes) -> str:
    """
    CHANGE: Added this function to encode image bytes to a base64 string.
    Args:
        image_bytes (bytes): Image data in bytes.
    Returns:
        str: Base64-encoded string of the image.
    """
    return base64.b64encode(image_bytes).decode("utf-8")


def get_frames(video_path: str, frame_step: int = 1) -> list:
    """
    Args:
        video_path (str): Path to the local video file.
        frame_step (int): Step between frames to extract (default 1 = every frame).
    Returns:
        list: List of dicts with 'idx' (frame index) and 'image' (image data as bytes).
    """
    frames = []
    cap = cv2.VideoCapture(video_path)
    idx = 0
    frame_idx = 0
    if not cap.isOpened():
        print(f"Error: Cannot open video file {video_path}")
        return frames
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        if frame_idx % frame_step == 0:
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            success, buffer = cv2.imencode('.png', frame_rgb)
            if success:
                frames.append({'idx': idx, 'image': buffer.tobytes()})
                idx += 1
        frame_idx += 1
    cap.release()
    return frames


def observe_run(states: List[Dict]):
    with open("/output/states.json", "w") as f:
        json.dump(states, f, indent=2)