import gradio as gr
import numpy as np
import zipfile
from PIL import Image
import tempfile

def split_image(image):
    im = Image.fromarray(np.uint8(image))

    # Obtiene las dimensiones de la imagen
    width, height = im.size

    with zipfile.ZipFile("images.zip", 'w') as zip:
        # Divide la imagen en 9 imágenes más pequeñas
        for i in range(3):
            for j in range(3):
                box = (j*width/3, i*height/3, (j+1)*width/3, (i+1)*height/3)
                region = im.crop(box)
                region = region.convert("RGB")
                # region = region.resize((width//3, height//3), Image.ANTIALIAS)
    
                with tempfile.NamedTemporaryFile(suffix='.png') as temp:
                    region.save(temp.name)
                    zip.write(temp.name)
                    temp.flush()
                    temp.close()
    return "images.zip"

gr.Interface(fn=split_image,
             inputs="image",
             outputs="file",
             title="Convert collage to images",
             examples = ["example.jpg"]  
            ).launch();