import gradio as gr
from PIL import Image
import rembg
import io

# Function to handle image uploads, optional background removal, and collage generation
def process_images(images, remove_bg):
    # Create a blank transparent A4 canvas
    final_collage = Image.new("RGBA", (2480, 3508), (255, 255, 255, 0))  # A4 size canvas in pixels at 300 DPI

    # Positioning each image onto the A4 canvas
    x_offset = 0
    y_offset = 0
    for image in images:
        img = Image.open(image)

        # Optionally remove the background
        if remove_bg:
            img = Image.open(io.BytesIO(rembg.remove(image.read())))  # Use rembg to remove background
            
        # Resize image to fit within A4 dimensions
        img.thumbnail((500, 500))  # Limit the size of each sticker to fit the canvas
        
        # Paste the image onto the canvas
        final_collage.paste(img, (x_offset, y_offset), img if img.mode == 'RGBA' else None)
        
        # Update position for next image
        y_offset += img.height
        if y_offset + img.height > final_collage.height:
            y_offset = 0
            x_offset += img.width

        # If we exceed the canvas size, we stop adding more images
        if x_offset + img.width > final_collage.width:
            break

    return final_collage

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Sticker Collage Creator (A4)")

    # Multiple file upload component
    images = gr.Files(label="Upload your images", type="file", file_count="multiple", file_types=["image"])
    
    # Checkbox for background removal
    remove_bg = gr.Checkbox(label="Remove background?", value=True)
    
    # Output image
    output = gr.Image(label="Collage", type="pil")

    # Button to trigger the function
    submit = gr.Button("Create Collage")
    submit.click(process_images, inputs=[images, remove_bg], outputs=output)

demo.launch()