src/helpers.py

import numpy as np
from PIL import Image

def resize_image(image_array, max_size=1048576):
    """
    Resizes a NumPy image array to a maximum total size of 1048576 pixels while preserving the aspect ratio.

    Args:
        image_array (numpy.ndarray): The input image array.
        max_size (int): The maximum total size of the image in pixels (default is 1048576).

    Returns:
        numpy.ndarray: The resized image array.
    """
    # Get the current size of the image
    height, width, _ = image_array.shape

    # Calculate the total size of the image
    total_size = height * width

    # If the image is already smaller than the max size, return the original image
    if total_size <= max_size:
        return image_array

    # Calculate the aspect ratio
    aspect_ratio = width / height

    # Calculate the new dimensions while preserving the aspect ratio
    if aspect_ratio > 1:
        new_width = int(np.sqrt(max_size * aspect_ratio))
        new_height = int(new_width / aspect_ratio)
    else:
        new_height = int(np.sqrt(max_size / aspect_ratio))
        new_width = int(new_height * aspect_ratio)

    # Resize the image using Pillow
    image = Image.fromarray(image_array)
    image = image.resize((new_width, new_height), resample=Image.BICUBIC)

    # Convert the resized image back to a NumPy array
    resized_image = np.array(image)

    return resized_image