copaint/gradio_ui.py

"""Gradio based ui for Copaint pdf generator.
"""
import gradio as gr
import shutil
import tempfile

from gradio_pdf import PDF
from importlib.resources import files
from pathlib import Path

from copaint.copaint import image_to_pdf 
import torchvision
import torch

fromPIltoTensor = torchvision.transforms.ToTensor()
fromTensortoPIL = torchvision.transforms.ToPILImage()


def add_grid_to_image(image, h_cells, w_cells):
    # image is a torch tensor of shape (3, h, w)
    image = image.convert("RGB")
    image = fromPIltoTensor(image)
    grid_color = torch.tensor([16,15,46]).unsqueeze(1).unsqueeze(1) / 255.0
    h,w = image.shape[1:]
    thickness = max(min(1, int(min(h,w)/100)), 1)
    print("thickness, h, w", thickness, h, w)
    for i in range(h_cells+1):
        idx_i = int(i*h/h_cells)
        image[:, idx_i-thickness:idx_i+thickness, :] = grid_color
        
    for j in range(w_cells+1):
        idx_j = int(j*w/w_cells)
        image[:, :, idx_j-thickness:idx_j+thickness] = grid_color
    image = fromTensortoPIL(image)

    return image


def canvas_ratio(image):
    w,h = image.size
    aspect_ratio = w/h
    if aspect_ratio > 1:
        aspect_ratio = 1/aspect_ratio
    
    # find nearest aspect ratio in the list of predefined ones 
    predefined_aspect_ratios = [2/3, 1/2, 1/1, 5/6, 4/5, 5/7]
    predefined_aspect_ratios_str = ["2/3", "1/2", "1/1", "5/6", "4/5", "5/7"]
    min_diff = float('inf')
    closest_ratio_idx = None
    for idx, ratio in enumerate(predefined_aspect_ratios):
        diff = abs(aspect_ratio - ratio)
        if diff < min_diff:
            min_diff = diff
            closest_ratio_idx = idx
    closest_ratio_str = predefined_aspect_ratios_str[closest_ratio_idx]
    
    if min_diff > 0.1:
        return None
    else:
        return f"Best canvas ratio: {closest_ratio_str}"


def add_grid_and_display_ratio(image, h_cells, w_cells):
    if image is None:
        return None, gr.update(visible=False)    
    return add_grid_to_image(image, h_cells, w_cells), gr.update(visible=True, value=canvas_ratio(image))


def process_copaint(
    input_image, 
    h_cells=None, 
    w_cells=None, 
    a4=False,
    high_res=False,
    cell_size_in_cm=None,
    min_cell_size_in_cm=2,
    copaint_name="",
    copaint_logo=None,
):
    """Process the input and generate CoPaint PDF"""
    # Create temporary directories for processing
    temp_input_dir = tempfile.mkdtemp()
    temp_output_dir = tempfile.mkdtemp()
    
    try:
        # Save uploaded images to temp directory
        input_path = Path(temp_input_dir) / "input_image.png"
        input_image.save(input_path)
        
        logo_path = None
        if copaint_logo is not None:
            logo_path = Path(temp_input_dir) / "logo.png"
            copaint_logo.save(logo_path)
        else:
            # Use default logo path from the package
            logo_path = files("copaint.static") / "logo_copaint.png"

        if copaint_name == "" or copaint_name is None:
            from copaint.cli import default_identifier
            copaint_name = default_identifier()

        if a4 == "A4":
            a4 = True
        else:
            a4 = False
        
        # Generate the PDF
        pdf_path = image_to_pdf(
            input_image=str(input_path), 
            logo_image=str(logo_path), 
            outputfolder=temp_output_dir, 
            h_cells=h_cells, 
            w_cells=w_cells, 
            unique_identifier=copaint_name,
            cell_size_in_cm=cell_size_in_cm,
            a4=a4,
            high_res=high_res,
            min_cell_size_in_cm=min_cell_size_in_cm
        )
        
        return pdf_path, None  # Return path and no error
    except Exception as e:
        # Return error message
        return None, f"Error generating PDF: {str(e)}"
    finally:
        # Clean up temporary input directory
        shutil.rmtree(temp_input_dir)


def build_gradio_ui():
    # Create Gradio Interface
    with gr.Blocks(title="CoPaint Generator", theme='NoCrypt/miku') as demo:

        gr.Markdown("# 🤖 CoPaint Generator")
        gr.Markdown("Upload an image with your painting design and set grid parameters to generate a CoPaint PDF template 🖨️📄✂️ for your next collaborative painting activities. 🎨🖌️")

        # --- inputs ---
        with gr.Row(equal_height=True):

            # Upload Design Template
            with gr.Column(scale=2):
                input_image = gr.Image(type="pil", label="Upload Your Design")

            with gr.Column(scale=1):
                # Grid
                with gr.Tab("Grid Layout"):
                    gr.Markdown("<div style='text-align: center; font-weight: bold;'>Squares' Grid</div>")
                    w_cells = gr.Number(label="↔ (width)", value=4, precision=0)
                    h_cells = gr.Number(label=" by ↕ (heigth)", value=6, precision=0)

                    gr.Examples(
                        examples=[
                            [6, 9],
                            [4, 6],
                            [3, 3],
                            [3, 4],
                            [2, 2]
                        ],
                        example_labels=[
                            "Copaint Wedding 6x9 Grid (54 squares)", 
                            "Copaint Classic 4x6 Grid (24 squares)", 
                            "Copaint Mini 3x3 Grid (9 squares)", 
                            "Copaint Mini 3x4 Grid (12 squares)", 
                            "Copaint Mini 2x2 Grid (4 squares)"],
                        inputs=[w_cells, h_cells],
                    )

                    # Grid + Design preview
                    gr.Markdown("<div style='text-align: center; font-weight: bold;'>Preview</div>")
                    output_image = gr.Image(label="Squares' Grid Preview", interactive=False)

                    # canvas ratio message
                    canvas_msg =  gr.Markdown(label="Canvas Ratio", visible=False)

                # PDF options
                with gr.Tab("PDF Printing Options"):
                    use_a4 = gr.Dropdown(choices=["US letter", "A4"], label="Paper Format", value="US letter")

                    with gr.Accordion("Advanced settings (optional)", open=False):
                        with gr.Row():
                            with gr.Column(scale=1):
                                high_res = gr.Checkbox(label="High Resolution Mode (>20sec long processing)")

                                cell_size = gr.Number(label="Square Size, in cm (optional)",
                                                      value="",
                                                      info="If none is provided, the design size automatically adjusts to fit on a single page.")

                                copaint_name = gr.Textbox(label="Add a Custom Design Name (optional)",
                                                          value="",
                                                          max_length=10,
                                                          info="You can add a custom design name: it will appear on the back of each square, in the top left corner.")
                                
                                copaint_logo = gr.Image(type="pil",
                                                        label="Add a Custom Logo (optional)")
                                gr.Markdown(
                                    "<div style='font-size: 0.85em;'>"
                                    "You can add a custom logo: it will appear on the back of each square, in the bottom right corner."
                                    "</div>")

        # --- outputs ---
        with gr.Row():
            # PDF
            with gr.Column(scale=1):
                submit_btn = gr.Button("Generate Copaint PDF", variant="primary")

        with gr.Row():
            with gr.Column(scale=1):
                output_file = gr.File(label="Download PDF", visible=False, interactive=False)
            with gr.Column(scale=1):
                output_error_msg = gr.Textbox(label="Error Message", visible=False)
        
        with gr.Row():
            with gr.Column(scale=1):
                output_pdf = PDF(label="PDF Preview")

        # Update output_image: trigger update when any input changes
        input_image.change(
                fn=add_grid_and_display_ratio,
                inputs=[input_image, h_cells, w_cells],
                outputs=[output_image, canvas_msg]
            )
        for component in [h_cells, w_cells]:
            component.change(
                fn=add_grid_to_image,
                inputs=[input_image, h_cells, w_cells],
                outputs=output_image
            )

        # Submit function: generate pdf
        def on_submit(input_image, h_cells, w_cells, use_a4, high_res, cell_size, copaint_name, copaint_logo):
            if input_image is None:
                return None, None, gr.update(visible=True, value="Please upload an image first 👀")
            
            if cell_size is None or cell_size == "" or cell_size == 0:
                cell_size = None
            
            pdf_path, error = process_copaint(
                input_image=input_image,
                h_cells=int(h_cells),
                w_cells=int(w_cells),
                a4=use_a4,
                high_res=high_res,
                cell_size_in_cm=cell_size if cell_size else None,
                min_cell_size_in_cm=float(2),
                copaint_name=copaint_name,
                copaint_logo=copaint_logo
            )
                
            if error:
                # Show error message
                return None, None, gr.update(visible=True, value=error)
            else:
                # Show successful PDF
                return pdf_path, gr.update(visible=True, value=pdf_path, interactive=False), gr.update(visible=False)
        
        submit_btn.click(
            on_submit,
            inputs=[input_image, h_cells, w_cells, use_a4, high_res, cell_size, copaint_name, copaint_logo],
            outputs=[output_pdf, output_file, output_error_msg]
        )
        
    return demo