utils/image_utils.py

# Standard library imports
import io
import base64
import urllib.request

# Third-party imports
from PIL import Image
import numpy as np


def load_image(image_path):
    """
    Loads an image from a URL, base64 string, or file.

    Args:
        image_path (str): The path to the image. It can be a URL, a base64 string, or a file path.

    Returns:
        PIL.Image.Image: The loaded image.
    """
    try:
        if image_path.startswith("http://") or image_path.startswith("https://"):
            # Debug url
            print("Debug URL:", image_path)
            # Load image from URL
            with urllib.request.urlopen(image_path) as response:
                image = Image.open(io.BytesIO(response.read()))
        elif image_path.startswith("data:image"):
            # Load image from base64 string
            image_data = base64.b64decode(image_path.split(",")[1])
            image = Image.open(io.BytesIO(image_data))
        else:
            # Load image from file
            image = Image.open(image_path)
        return image
    except Exception as e:
        print(f"Error loading image: {e}")
        return None


def preprocess_image(image):
    """
    Preprocesses the image for the models.

    Args:
        image (PIL.Image.Image): The image to preprocess.

    Returns:
        numpy.ndarray: The preprocessed image as a NumPy array.
    """
    # Ensure image is a PIL Image before converting
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)

    image = image.convert("RGB")
    image = np.array(image)
    return image


def get_image_from_input(input_type, uploaded_image, image_url, base64_string):
    """
    Centralized function to get an image from various input types.

    Args:
        input_type (str): The selected input method ("Upload File", "Enter URL", "Enter Base64").
        uploaded_image (PIL.Image.Image): The uploaded image (if input_type is "Upload File").
        image_url (str): The image URL (if input_type is "Enter URL").
        base64_string (str): The image base64 string (if input_type is "Enter Base64").

    Returns:
        PIL.Image.Image: The loaded image, or None if an error occurred or no valid input was provided.
    """
    image = None
    input_value = None

    if input_type == "Upload File" and uploaded_image is not None:
        image = uploaded_image  # This is a PIL Image from gr.Image(type="pil")
        print("Using uploaded image (PIL)")  # Debug print

    elif input_type == "Enter URL" and image_url and image_url.strip():
        input_value = image_url
        print(f"Using URL: {input_value}")  # Debug print

    elif input_type == "Enter Base64" and base64_string and base64_string.strip():
        input_value = base64_string
        print("Using Base64 string")  # Debug print

    else:
        print("No valid input provided based on selected type.")
        return None  # No valid input

    # If input_value is set (URL or Base64), use the local load_image
    if input_value:
        image = load_image(input_value)
        if image is None:
            print("Error: Could not load image from provided input.")
            return None  # load_image failed

    # Now 'image' should be a PIL Image or None
    if image is None:
        print("Image is None after loading/selection.")
        return None

    return image