import gradio as gr
import numpy as np
import zipfile
from io import BytesIO
from PIL import Image

def split_image_grid(image, grid_size):
    # Convert the image to a NumPy array
    img = np.array(image)
    width, height = img.shape[1], img.shape[0]
    grid_width, grid_height = grid_size

    # Calculate the size of each grid cell
    cell_width = width // grid_width
    cell_height = height // grid_height

    # Split the image into individual frames
    frames = []
    for i in range(grid_height):
        for j in range(grid_width):
            left = j * cell_width
            upper = i * cell_height
            right = left + cell_width
            lower = upper + cell_height
            frame = img[upper:lower, left:right]
            frames.append(frame)

    return frames

def zip_images(images):
    # Create a BytesIO object to hold the zip file
    zip_buffer = BytesIO()
    with zipfile.ZipFile(zip_buffer, 'w') as zipf:
        for idx, img in enumerate(images):
            # Save each image to the zip file
            img_buffer = BytesIO()
            img = Image.fromarray(img)  # Convert NumPy array to PIL Image
            img.save(img_buffer, format='PNG')
            img_buffer.seek(0)
            zipf.writestr(f'image_{idx}.png', img_buffer.getvalue())

    zip_buffer.seek(0)
    return zip_buffer

def process_image(image, grid_size):
    # Split the image into a grid of frames
    frames = split_image_grid(image, grid_size)
    # Zip the frames into a single zip file
    zip_file = zip_images(frames)
    return zip_file

with gr.Blocks() as demo:
    with gr.Row():
        image_input = gr.Image(label="Input Image")
        grid_size_input = gr.Slider(1, 10, value=2, step=1, label="Grid Size")
    zip_output = gr.File(label="Output Zip File")
    process_button = gr.Button("Process Image")

    process_button.click(process_image, inputs=[image_input, grid_size_input], outputs=zip_output)

demo.launch(show_error=True)