app.py

'''
this script is used to split an image into tiles and download them as a zip file.

The script defines two functions: split_image and create_zip_file. The
split_image function takes an input image and tile size as arguments and splits
the image into smaller tiles. The create_zip_file function takes a list of image
tiles and a prefix as arguments and creates a zip file containing all the tiles.

The process_image function combines the two functions to split the input image
into tiles and create a zip file of the tiles. The main function launches a Gradio
app that allows users to upload an image, specify the tile size, view the resulting
tiles, and download all tiles in a zip archive.

To run the script, simply execute it in a Python environment. The Gradio app will
open in a new tab in your web browser, allowing you to interact with the image
splitting functionality.

How to use the script:
1. Run the script in a Python environment.
```python
python split_and_zip.py
```
2. Open the provided local URL in a web browser.
3. Upload an image file.
4. Adjust the tile size using the slider.
5. Click the "Process Image" button to split the image into tiles.
6. View the resulting tiles in the gallery.
7. Click the "Download Tiles" button to download all tiles as a zip file.



'''
import os
import cv2
import numpy as np
import gradio as gr
import tempfile
import zipfile
from typing import List, Optional, Tuple


def split_image(image: np.ndarray, tile_size: int) -> List[np.ndarray]:
    """
    Split a large image into smaller tiles.

    Parameters
    ----------
    image : np.ndarray
        Input image as a NumPy array.
    tile_size : int
        Size of each square tile in pixels.

    Returns
    -------
    List[np.ndarray]
        A list of image tiles as NumPy arrays.

    Examples
    --------
    >>> image = cv2.imread('path/to/image.jpg')
    >>> tiles = split_image(image, tile_size=500)
    >>> len(tiles)
    16
    """
    if image is None:
        return []

    h, w = image.shape[:2]
    tiles = []

    for y in range(0, h, tile_size):
        for x in range(0, w, tile_size):
            end_y = min(y + tile_size, h)
            end_x = min(x + tile_size, w)
            tile = image[y:end_y, x:end_x]
            if tile.shape[0] > 0 and tile.shape[1] > 0:
                tiles.append(tile)
    return tiles


def create_zip_file(tiles: List[np.ndarray], prefix: str = "tile") -> str:
    """
    Create a zip file containing all image tiles.

    Parameters
    ----------
    tiles : List[np.ndarray]
        List of image tiles as NumPy arrays.
    prefix : str, optional
        Prefix for each tile filename, by default "tile".

    Returns
    -------
    str
        Path to the created zip file.

    Examples
    --------
    >>> zip_path = create_zip_file(tiles, prefix='sample')
    >>> os.path.exists(zip_path)
    True
    """
    temp_dir = tempfile.mkdtemp()
    zip_path = os.path.join(temp_dir, "tiles.zip")

    with zipfile.ZipFile(zip_path, 'w') as zf:
        for i, tile in enumerate(tiles):
            tile_path = os.path.join(temp_dir, f"{prefix}_{i}.png")
            cv2.imwrite(tile_path, cv2.cvtColor(tile, cv2.COLOR_RGB2BGR))
            zf.write(tile_path, f"{prefix}_{i}.png")

    return zip_path


def process_image(image: np.ndarray, tile_size: int) -> Tuple[List[np.ndarray], str]:
    """
    Split the input image into tiles and create a zip file of the tiles.

    Parameters
    ----------
    image : np.ndarray
        Input image as a NumPy array.
    tile_size : int
        Size of each square tile in pixels.

    Returns
    -------
    Tuple[List[np.ndarray], str]
        A tuple containing the list of image tiles and the path to the zip file.

    Examples
    --------
    >>> tiles, zip_path = process_image(image, tile_size=500)
    >>> len(tiles)
    16
    >>> os.path.exists(zip_path)
    True
    """
    tiles = split_image(image, tile_size)
    zip_path = create_zip_file(tiles) if tiles else ""
    return tiles, zip_path


def main():
    """
    Launch the Gradio app for splitting images into tiles and downloading them as a zip file.

    The app allows users to upload an image, specify the tile size, view the resulting tiles,
    and download all tiles in a zip archive.

    Examples
    --------
    Run the script and open the provided local URL in a web browser.
    """
    with gr.Blocks() as interface:
        with gr.Row():
            input_image = gr.Image(type="numpy", label="Input Image")
            tile_size = gr.Slider(
                minimum=100, maximum=1000, step=100, value=500, label="Tile Size"
            )

        with gr.Row():
            submit_btn = gr.Button("Process Image")

        with gr.Row():
            gallery = gr.Gallery(label="Tiles", columns=3)
            download_btn = gr.File(label="Download Tiles")

        submit_btn.click(
            fn=process_image,
            inputs=[input_image, tile_size],
            outputs=[gallery, download_btn],
        )

    interface.launch()


if __name__ == '__main__':
    main()