Image processing functions for the Rainfall Rescue imagesΒΆ

# Functions to cut up images (for tokenization) and to stitch them together again

from PIL import Image
import numpy as np


# Cut an image into square blocks of a specified size
#  With padding where necessary.
# def cut_image(image, block_size):
#     width, height = image.size
#     blocks = []

#     # Calculate the number of blocks in each dimension
#     num_blocks_x = (width + block_size - 1) // block_size
#     num_blocks_y = (height + block_size - 1) // block_size

#     # Create blocks with padding if necessary
#     for i in range(num_blocks_y):
#         for j in range(num_blocks_x):
#             left = j * block_size
#             upper = i * block_size
#             right = min(left + block_size, width)
#             lower = min(upper + block_size, height)

#             # Create a new image for the block
#             cropped = image.crop((left, upper, right, lower))
#             # If the block is smaller than the specified size, pad it
#             if cropped.width < block_size or cropped.height < block_size:
#                 block = Image.new(
#                     image.mode, (block_size, block_size), color="white"
#                 )  # White padding
#                 block.paste(cropped, (0, 0))
#             else:
#                 block = cropped

#             blocks.append(block)

#     return blocks


def cut_image(image, block_size, overlap=0.0):
    """
    Cut an image into square blocks of a specified size with adjustable overlap.

    Args:
        image (PIL.Image): The input image.
        block_size (int): The size of each block (in pixels).
        overlap (float): Fraction of block size to overlap between blocks (0.0 to 1.0).

    Returns:
        list: List of image blocks.
    """
    width, height = image.size
    blocks = []

    # Calculate the step size based on overlap
    step_size = int(block_size * (1 - overlap))

    # Create blocks with padding if necessary
    for upper in range(0, height, step_size):
        for left in range(0, width, step_size):
            right = min(left + block_size, width)
            lower = min(upper + block_size, height)

            # Create a new image for the block
            cropped = image.crop((left, upper, right, lower))
            # If the block is smaller than the specified size, pad it
            if cropped.width < block_size or cropped.height < block_size:
                block = Image.new(
                    image.mode, (block_size, block_size), color="white"
                )  # White padding
                block.paste(cropped, (0, 0))
            else:
                block = cropped

            blocks.append(block)

    return blocks