# %%writefile encode_image.py
import base64
from typing import Union

import cv2
import numpy as np
from PIL import Image

Image.MAX_IMAGE_PIXELS = None  # Removes the limit, use with caution


def encode_image(image: Union[str, np.ndarray]) -> str:
    """
    Encodes an image as a base64 string.

    Args:
        image (Union[str, np.ndarray]): Path to the image file or a NumPy array representing the image.

    Returns:
        str: Base64-encoded image string.
    """
    if isinstance(image, str):  # If the input is a file path
        with open(image, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode("utf-8")
    elif isinstance(image, np.ndarray):  # If the input is a NumPy array
        _, buffer = cv2.imencode(".jpg", image)  # Encode image as JPEG
        return base64.b64encode(buffer).decode("utf-8")
    else:
        raise TypeError("Input must be a file path (str) or a NumPy array.")