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 load_with_transparency(image_file):
    # Open manually to preserve RGBA
    image = image_file #Image.open(image_file.name).convert("RGBA")
    
    # Example: return image size and check mode
    print( f"Image mode: {image.mode}, size: {image.size}")

    logger.info(f"TRANSPARENCY MODE: {image.mode}")
    # 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
        logger.info(f"TRANSPARENCY REMOVED WITH WHITE BACKGROUND")
    else:
        logger.info(f"TRANSPARENCY REMOVED")
        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(image, h_cells, w_cells, language="en"):

    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 = [1/3, 1/2, 2/3, 1/1, 5/6, 4/5, 5/7]
    predefined_aspect_ratios_str = ["1:3", "1:2", "2:3", "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.2:
        return None
    else:

        example_str = ""
        if language == "en":
            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."
            elif closest_ratio_str == "1:2":
                example_str = ", for example, a 12” by 24” canvas."
        elif language == "fr":
            if closest_ratio_str == "1:1":
                if h_cells == 2 and h_cells == 2:
                    example_str = ", par exemple, une toile de 15cmx15cm."
                if h_cells == 3 and h_cells == 3:
                    example_str = ", par exemple, une toile de 15cmx15cm ou 20cmx20cm."
            elif closest_ratio_str == "5:6":
                if h_cells == 6 and w_cells == 9:
                    example_str = ", par exemple, une toile de 60cmx90cm."
                if h_cells == 9 and w_cells == 6:
                    example_str = ", par exemple, une toile de 90cmx60cm."
            elif closest_ratio_str == "4:5":
                if h_cells == 4 and w_cells == 6:
                    example_str = ", par exemple, une toile de 15cmx20cm."
                if h_cells == 6 and w_cells == 4:
                    example_str = ", par exemple, une toile de 45cmx60cm ou 30cmx40cm."
            elif closest_ratio_str == "2:3" and ((h_cells == 6 and w_cells == 9) or (h_cells == 9 and w_cells == 6)):
                example_str = ", par exemple, une toile de 60cmx90cm."
            elif closest_ratio_str == "1:2":
                example_str = ", par exemple, une toile de 30cmx60cm."
    
    if language == "en":
        return_str = f"You have selected 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} size ratio</b> to match your design's ({w} x {h} pixels){example_str}"
    else:
        return_str = f"Vous avez choisi un découpage  de {h_cells} par {w_cells} ({h_cells*w_cells} tuiles).\n\n"
        return_str += f"Choix de la toile: prenez une toile dont le rapport largeur sur longueur est de <b>{closest_ratio_str}</b> pour respecter la taille de votre design ({w}x{h} pixels){example_str}"
    print(f"return_str: {return_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, language="en"):
    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(image, h_cells, w_cells, language))


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,
    language="en"
):
    """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(language=language)

        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(language="en"):
    # Create Gradio Interface
    with gr.Blocks(title="Copaint Generator", theme='NoCrypt/miku') as demo:

        if language == "en":
            gr.Markdown("# 🤖 Copaint Generator")
            gr.Markdown("Upload an image,  set grid parameters, and we generate your Copaint PDF  🖨️📄✂️ for your next collaborative painting activity. 🎨🖌️")
        elif language == "fr":
            gr.Markdown("# 🤖 Générateur de Copaint")
            gr.Markdown("Téléchargez une image, définissez les paramètres de découpage et on vous génère votre PDF Copaint 🖨️📄✂️ pour votre prochaine activité de peinture collaborative. 🎨🖌️")

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

            # Upload Design Template
            with gr.Column(scale=2):
                if language == "en":
                    label_input_image = "Upload Your Design"
                elif language == "fr":
                    label_input_image = "Déposer votre image"
                input_image = gr.Image(type="pil", image_mode="RGBA", label=label_input_image)
                
            input_image.upload(
                fn=load_with_transparency,
                inputs=input_image,
                outputs=input_image
            )
            with gr.Column(scale=1):
                # Grid
                if language == "en":
                    label_grid_layout = "Grid Layout"
                elif language == "fr":
                    label_grid_layout = "Découpage"
                with gr.Tab(label_grid_layout):
                    if language == "en":
                        gr.Markdown("<div style='text-align: left; 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)
                        language_gr_state = gr.State("en")
                    elif language == "fr":
                        gr.Markdown("<div style='text-align: center; font-weight: bold;'>Découpage (nombre de tuiles)</div>")
                        w_cells = gr.Number(label="↔ (largeur)", value=4, precision=0)
                        h_cells = gr.Number(label=" par ↕ (hauteur)", value=6, precision=0)
                        language_gr_state = gr.State("fr")
                    

                    examples_labels_en = [
                        "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)"
                    ]
                    examples_labels_fr = [
                        "Copaint Mariage 6 par 9 tuiles (54 tuiles)", 
                        "Copaint Classique 4 par 6 tuiles (24 tuiles)", 
                        "Copaint Mini 3 par 3 tuiles (9 tuiles)", 
                        "Copaint Mini 3 par 4 tuiles (12 tuiles)", 
                        "Copaint Mini 2 par 2 tuiles (4 tuiles)"
                    ]
                    if language == "en":
                        example_labels = examples_labels_en
                        label_example = "Examples"
                    elif language == "fr":
                        example_labels = examples_labels_fr
                        label_example = "Exemples"
                    gr.Examples(
                        examples=[
                            [6, 9],
                            [4, 6],
                            [3, 3],
                            [3, 4],
                            [2, 2]
                        ],
                        example_labels=example_labels,
                        inputs=[w_cells, h_cells],
                        label=label_example
                    )

                    # Grid + Design preview
                    if language == "en":
                        gr.Markdown("<div style='text-align: center; font-weight: bold;'>Preview</div>")
                    elif language == "fr":
                        gr.Markdown("<div style='text-align: center; font-weight: bold;'>Aperçu</div>")
                    
                    output_image = gr.Image(label=None, show_label=False,  show_fullscreen_button=False, interactive=False, show_download_button=False)

                    # canvas ratio message
                    if language == "en":
                        canvas_msg_label = "Canvas Ratio"
                    elif language == "fr":
                        canvas_msg_label = "Aspect Ratio"
                    canvas_msg =  gr.Markdown(label=canvas_msg_label, visible=False)

                # PDF options
                if language == "en":
                    label_pdf_options = "PDF Printing Options"
                elif language == "fr":
                    label_pdf_options = "Options d'impression"
                with gr.Tab(label_pdf_options):
                    if language == "en":
                        use_a4 = gr.Dropdown(choices=["US letter", "A4"], label="Paper Format", value="US letter")
                    elif language == "fr":
                        use_a4 = gr.Dropdown(choices=["A4", "US letter"], label="Format de papier", value="A4")

                    if language == "en":
                        label_advanced_settings = "Advanced settings (optional)"
                    elif language == "fr":
                        label_advanced_settings = "Options avancées (optionnel)"
                    with gr.Accordion(label_advanced_settings, open=False):
                        with gr.Row():
                            with gr.Column(scale=1):
                                if language == "en":
                                    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>")
                                elif language == "fr":
                                    high_res = gr.Checkbox(label="Mode haute résolution (plus de 20 secondes de traitement)")
                                    cell_size = gr.Number(label="Taille des tuiles, en cm (optionnel)",
                                                        value="",
                                                        info="Si elle n'est pas spécifiée, la hauteur des tuiles est automatiquement calculée afin que votre impression tienne sur une seule page, recto-verso.")
                                    copaint_name = gr.Textbox(label="Ajouter un nom personnalisé à votre design (optionnel)",
                                                            value="",
                                                            max_length=10,
                                                            info="Vous pouvez ajouter un nom personnalisé: : il apparaîtra au dos de chaque tuile, en haut à gauche.")
                                    copaint_logo = gr.Image(type="pil",
                                                            label="Ajouter un logo personnalisé (optionnel)")
                                    gr.Markdown(
                                        "<div style='font-size: 0.85em;'>"
                                        "Vous pouvez ajouter un logo personnalisé: il apparaîtra au dos de chaque tuile, en bas à droite."
                                        "</div>")

        # --- outputs ---
        with gr.Row():
            # PDF
            with gr.Column(scale=1):
                if language == "en":
                    submit_btn = gr.Button("Generate Copaint PDF", variant="primary")
                elif language == "fr":
                    submit_btn = gr.Button("Générer mon PDF Copaint", variant="primary")

        with gr.Row():
            with gr.Column(scale=1):
                if language == "en":
                    output_file = gr.File(label="Download PDF", visible=False, interactive=False)
                elif language == "fr":
                    output_file = gr.File(label="Télécharger le PDF", visible=False, interactive=False)
            with gr.Column(scale=1):
                if language == "en":
                    output_error_msg = gr.Textbox(label="Error Message", visible=False)
                elif language == "fr":
                    output_error_msg = gr.Textbox(label="Message d'erreur", visible=False)
        
        with gr.Row():
            with gr.Column(scale=1):
                if language == "en":
                    output_pdf = PDF(label="PDF Preview")#, show_label=False, show_fullscreen_button=False, interactive=False, show_download_button=False)
                elif language == "fr":
                    output_pdf = PDF(label="Aperçu du PDF")#, 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, language_gr_state],
                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,
                language=language
            )
                
            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