copaint/gradio_ui.py

"""Gradio based ui for Copaint pdf generator.
"""
import gradio as gr
import shutil
import tempfile
import logging
import os
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

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

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

from PIL import Image

def remove_transparency(image: Image.Image):
    # Convert transparency to white background
    if image.mode in ("RGBA", "LA"):
        background = Image.new("RGB", image.size, (255, 255, 255))  # white background
        background.paste(image, mask=image.split()[-1])  # paste using alpha channel as mask
        image = background
    else:
        image = image.convert("RGB")  # just to be safe

    return image  # or continue processing

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)
    
    # Calculate average image brightness
    avg_brightness = image.mean()
    # Use white grid for dark images, dark blue for bright ones
    grid_color = torch.tensor([255,255,255]).unsqueeze(1).unsqueeze(1) / 255.0 if avg_brightness < 0.3 else 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)
    logger.debug(f"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 get_canvas_ratio_message(h_cells, w_cells):
    # Calculate aspect ratio
    ratio = h_cells / w_cells if h_cells > w_cells else w_cells / h_cells
    closest_ratio = None
    closest_ratio_str = None
    min_diff = float('inf')
    
    # Common aspect ratios
    ratios = {
        "1:1": 1,
        "5:6": 5/6,
        "3:4": 3/4,
        "2:3": 2/3
    }
    
    # Find closest ratio
    for ratio_str, r in ratios.items():
        diff = abs(ratio - r)
        if diff < min_diff:
            min_diff = diff
            closest_ratio = r
            closest_ratio_str = ratio_str
    
    example_str = ""
    
    if closest_ratio_str == "1:1":
        if h_cells == 2 and h_cells == 2:
            example_str = ' (for example, a 6" by 6" canvas)'
        if h_cells == 3 and h_cells == 3:
            example_str = ' (for example, an 8" by 8" canvas)'
    elif closest_ratio_str == "5:6":
        example_str = ' (for example, a 10" by 12" canvas)'
    elif closest_ratio_str == "3:4":
        if (h_cells == 3 and w_cells == 4) or (h_cells == 4 and w_cells == 3):
            example_str = ' (for example, a 6" by 8" canvas)'
        if (h_cells == 4 and w_cells == 6) or (h_cells == 6 and w_cells == 4):
            example_str = ' (for example, an 18" by 24", or a 12" by 16" canvas)'
    elif closest_ratio_str == "2:3" and ((h_cells == 6 and w_cells == 9) or (h_cells == 9 and w_cells == 6)):
        example_str = ' (for example, a 24" by 36" canvas)'

    return_str = f"You have chosen a <b>{h_cells}x{w_cells} Grid ({h_cells*w_cells} squares)</b>.\n\n"
    return_str += f"Preparing your canvas: <b>choose a canvas with a {closest_ratio_str} ratio</b> to respect your design's size{example_str}."
    return f"<div style='font-size: 1.2em; line-height: 1.5;'>{return_str}</div>"

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=get_canvas_ratio_message(h_cells, w_cells))


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
    # Use /dev/shm if available for better performance
    if os.path.exists('/dev/shm'):
        temp_input_dir = tempfile.mkdtemp(dir='/dev/shm')
        temp_output_dir = tempfile.mkdtemp(dir='/dev/shm')
    else:
        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", preprocess=handle_upload)
            image_input.upload(
                fn=remove_transparency,
                inputs=image_input,
                outputs=image_input
            )
            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=None, show_label=False,  show_fullscreen_button=False, interactive=False, show_download_button=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. In most situations, you don't need this.")

                                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")#, show_label=False, show_fullscreen_button=False, interactive=False, show_download_button=False)

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

        # 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