src/app.py

#!/usr/bin/env python3

import logging
import shutil
import tempfile
from functools import partial
from pathlib import Path

import gradio as gr

from png23mf import (
    flip_image_horizontally,
    generate_3mf_file,
    generate_qrcode,
    overlay_images_core,
)

logger = logging.getLogger(__name__)

temp_dir = Path(tempfile.gettempdir())


def calculate_zoom_value(
    *,
    base_img,
    overlay_img,
    current_zoom,
    resolution,
):
    """Calculate zoom to fit overlay on base image."""
    logger.debug(f"Entering calculate_zoom_value with arguments: {locals()}")
    if base_img is None or overlay_img is None:
        result = current_zoom
        logger.debug(f"Exiting calculate_zoom_value with result: {result!r}")
        return result
    base = base_img.convert(mode="RGBA")
    w_base, h_base = base.size
    if max(w_base, h_base) > resolution:
        scale_factor = resolution / max(w_base, h_base)
        w_base = int(w_base * scale_factor)
        h_base = int(h_base * scale_factor)
    base_longer = max(w_base, h_base)
    w_overlay, h_overlay = overlay_img.size
    overlay_longer = max(w_overlay, h_overlay)
    if overlay_longer > base_longer:
        zoom = round(base_longer / overlay_longer, 2)
        zoom = max(0.01, min(10.0, zoom))
        result = zoom
        logger.debug(f"Exiting calculate_zoom_value with result: {result!r}")
        return result
    else:
        result = current_zoom
        logger.debug(f"Exiting calculate_zoom_value with result: {result!r}")
        return result


def update_width_from_height(
    width,
    height,
    use_aspect_ratio,
    base_img,
    overlay_img,
):
    """Calculate new width based on height and aspect ratio."""
    logger.debug(
        f"Entering update_width_from_height with arguments: {locals()}"
    )
    if use_aspect_ratio:
        if base_img is not None:
            w_base, h_base = base_img.size
        elif overlay_img is not None:
            w_base, h_base = overlay_img.size
        else:
            result = width
            logger.debug(
                f"Exiting update_width_from_height with result: {result!r}"
            )
            return result
        new_width = round(height * w_base / h_base, 2)
        if abs(new_width - width) > 1e-6:
            result = new_width
            logger.debug(
                f"Exiting update_width_from_height with result: {result!r}"
            )
            return result
    result = width
    logger.debug(f"Exiting update_width_from_height with result: {result!r}")

    return result


def update_height_from_width(
    width,
    height,
    use_aspect_ratio,
    base_img,
    overlay_img,
):
    """Calculate new height based on width and aspect ratio."""
    logger.debug(
        f"Entering update_height_from_width with arguments: {locals()}"
    )
    if use_aspect_ratio:
        if base_img is not None:
            w_base, h_base = base_img.size
        elif overlay_img is not None:
            w_base, h_base = overlay_img.size
        else:
            result = height
            logger.debug(
                f"Exiting update_height_from_width with result: {result!r}"
            )
            return result
        new_height = round(width * h_base / w_base, 2)
        if abs(new_height - height) > 1e-6:
            result = new_height
            logger.debug(
                f"Exiting update_height_from_width with result: {result!r}"
            )
            return result
    result = height
    logger.debug(f"Exiting update_height_from_width with result: {result!r}")

    return result


def overlay_front_images(
    base_img,
    front_img,
    x,
    y,
    rot,
    zoom,
    indexed_colors,
    use_common_colors,
    morph_size,
    resolution,
):
    """Overlay an image onto a base image and extract masks."""
    base_rgb, masks, mask_colors = overlay_images_core(
        base_img=base_img,
        overlay_img=front_img,
        x=x,
        y=y,
        rot=rot,
        zoom=zoom,
        indexed_colors=indexed_colors,
        use_common_colors=use_common_colors,
        morph_size=morph_size,
        resolution=resolution,
    )
    if base_rgb is None:
        return None, None, None, None, gr.update(interactive=False)
    return base_rgb, masks, masks, mask_colors, gr.update(interactive=True)


def overlay_back_images(
    base_img,
    back_img,
    x,
    y,
    rot,
    zoom,
    indexed_colors,
    use_common_colors,
    morph_size,
    resolution,
):
    """Overlay an image onto a base image and extract masks."""
    flipped_base_img = (
        flip_image_horizontally(image=base_img)
        if base_img is not None
        else None
    )
    base_rgb, masks, mask_colors = overlay_images_core(
        base_img=flipped_base_img,
        overlay_img=back_img,
        x=x,
        y=y,
        rot=rot,
        zoom=zoom,
        indexed_colors=indexed_colors,
        use_common_colors=use_common_colors,
        morph_size=morph_size,
        resolution=resolution,
    )
    if base_rgb is None:
        return None, None, None, None, gr.update(interactive=False)
    return base_rgb, masks, masks, mask_colors, gr.update(interactive=True)


def update_slider_ranges(base_img, overlay_img, resolution):
    if base_img is None:
        if overlay_img is not None:
            w, h = overlay_img.size
        else:
            w, h = resolution, resolution
    else:
        w, h = base_img.size
    max_dim = resolution
    if max(w, h) > max_dim:
        scale_factor = max_dim / max(w, h)
        w = int(w * scale_factor)
        h = int(h * scale_factor)
    default_width = round(100, 2)
    default_height = round(default_width * h / w, 2)
    if overlay_img is not None:
        zoom_val = calculate_zoom_value(
            base_img=base_img,
            overlay_img=overlay_img,
            current_zoom=1.0,
            resolution=resolution,
        )
    else:
        zoom_val = 1.0
    return (
        gr.update(minimum=-w, maximum=w, step=1, value=0),
        gr.update(minimum=-h, maximum=h, step=1, value=0),
        gr.update(maximum=180, step=1, value=0),
        gr.update(maximum=10.0, step=0.01, value=zoom_val),
        gr.update(value=4),
        gr.update(value=default_width),
        gr.update(value=default_height),
    )


def get_user_dir(req: gr.Request) -> str:
    session_hash = (
        str(req.session_hash) if req.session_hash is not None else "unknown"
    )
    user_dir = temp_dir / session_hash
    user_dir.mkdir(parents=True, exist_ok=True)
    return str(user_dir)


def delete_user_dir(req: gr.Request) -> None:
    session_hash = (
        str(req.session_hash) if req.session_hash is not None else "unknown"
    )
    user_dir = temp_dir / session_hash
    if user_dir.exists():
        shutil.rmtree(user_dir)


def switch_to_model_tab():
    return gr.update(selected="model")


generate_3mf_with_frames = partial(generate_3mf_file, animate_frames=18)


def generate_qr_code_from_text(text: str):
    """Generate a QR code image from input text."""
    logger.debug(
        f"Entering generate_qr_code_from_text with arguments: {locals()}"
    )
    image = generate_qrcode(fill_color="random", text=text)
    result = image
    logger.debug(f"Exiting generate_qrcode with result: {result!r}")

    return result


with gr.Blocks() as demo:
    gr.Markdown(value="# PNG to scad/3mf")

    with gr.Row():
        with gr.Column():
            with gr.Tabs() as left_tabs:
                with gr.Tab("Front Image"):
                    with gr.Row():
                        front_overlay_img = gr.Image(
                            label="Front Image",
                            type="pil",
                            image_mode="RGBA",
                        )
                    with gr.Row():
                        qr_text_front = gr.Textbox(
                            label="Generate QR-code",
                            placeholder="Enter text",
                            submit_btn="▣ QRCode",
                        )
                    with gr.Column():
                        x_slider = gr.Slider(
                            label="X (px)",
                            minimum=-500,
                            maximum=500,
                            step=1,
                            value=0,
                            interactive=False,
                            info="Overlay image X coordinate relative to base",
                        )
                        y_slider = gr.Slider(
                            label="Y (px)",
                            minimum=-500,
                            maximum=500,
                            step=1,
                            value=0,
                            interactive=False,
                            info="Overlay image Y coordinate relative to base",
                        )
                        rot_slider = gr.Slider(
                            label="Rotation (°)",
                            minimum=-180,
                            maximum=180,
                            step=1,
                            value=0,
                            interactive=False,
                            info="Rotate overlay image",
                        )
                        zoom_slider = gr.Slider(
                            label="Zoom",
                            minimum=0.01,
                            maximum=10.0,
                            step=0.01,
                            value=1.0,
                            interactive=False,
                            info="Make overlay image bigger or smaller",
                        )
                        indexed_slider = gr.Slider(
                            label="Indexed Colors",
                            minimum=2,
                            maximum=255,
                            step=1,
                            value=4,
                            info="Overlay image colors will be quantized",
                        )
                        use_common_colors_checkbox = gr.Checkbox(
                            label="Use most common colors for quantization",
                            value=True,
                            interactive=True,
                            info="Best for the logo's.",
                        )
                        morphology_slider = gr.Slider(
                            label="Morphology Size",
                            minimum=1,
                            maximum=20,
                            step=1,
                            value=1,
                            info=(
                                "Size of morphological opening/closing "
                                "to remove small features"
                            ),
                        )
                with gr.Tab("Back Image"):
                    with gr.Row():
                        back_overlay_img = gr.Image(
                            label="Back Image",
                            type="pil",
                            image_mode="RGBA",
                        )
                    with gr.Row():
                        qr_text_back = gr.Textbox(
                            label="Generate QR-code",
                            placeholder="Enter text",
                            submit_btn="▣ QRCode",
                        )
                    with gr.Column():
                        x_slider_back = gr.Slider(
                            label="X (px)",
                            minimum=-500,
                            maximum=500,
                            step=1,
                            value=0,
                            interactive=False,
                            info="Overlay image X coordinate relative to base",
                        )
                        y_slider_back = gr.Slider(
                            label="Y (px)",
                            minimum=-500,
                            maximum=500,
                            step=1,
                            value=0,
                            interactive=False,
                            info="Overlay image Y coordinate relative to base",
                        )
                        rot_slider_back = gr.Slider(
                            label="Rotation (°)",
                            minimum=-180,
                            maximum=180,
                            step=1,
                            value=0,
                            interactive=False,
                            info="Rotate overlay image",
                        )
                        zoom_slider_back = gr.Slider(
                            label="Zoom",
                            minimum=0.01,
                            maximum=10.0,
                            step=0.01,
                            value=1.0,
                            interactive=False,
                            info="Make overlay image bigger or smaller",
                        )
                        indexed_slider_back = gr.Slider(
                            label="Indexed Colors",
                            minimum=2,
                            maximum=255,
                            step=1,
                            value=4,
                            info="Overlay image colors will be quantized",
                        )
                        use_common_colors_checkbox_back = gr.Checkbox(
                            label="Use most common colors for quantization",
                            value=True,
                            interactive=True,
                            info="Best for the logo's.",
                        )
                        morphology_slider_back = gr.Slider(
                            label="Morphology Size",
                            minimum=1,
                            maximum=20,
                            step=1,
                            value=1,
                            info=(
                                "Size of morphological opening/closing "
                                "to remove small features"
                            ),
                        )
                with gr.Tab("Base Image"):
                    gr.Markdown(
                        value="**Tip:** The base image must be "
                        "black and white. Black is treated as transparent"
                        "and White will be extruded."
                    )
                    with gr.Row():
                        base_img = gr.Image(
                            label="Base Image",
                            type="pil",
                            image_mode="RGB",
                        )
                    gr.Markdown(value="**Examples**")
                    base_examples_dir = (
                        Path(__file__).parent.parent / "examples" / "base"
                    )
                    gr.Examples(
                        examples=str(base_examples_dir),
                        inputs=[base_img],
                        label="Base",
                    )
                    front_examples_dir = (
                        Path(__file__).parent.parent / "examples" / "front"
                    )
                    gr.Examples(
                        examples=str(front_examples_dir),
                        inputs=[front_overlay_img],
                        label="Front",
                    )
                    back_examples_dir = (
                        Path(__file__).parent.parent / "examples" / "back"
                    )
                    gr.Examples(
                        examples=str(back_examples_dir),
                        inputs=[back_overlay_img],
                        label="Back",
                    )
                with gr.Tab("Settings"):
                    with gr.Column():
                        resolution_dropdown = gr.Dropdown(
                            label="Resolution (px)",
                            choices=[512, 1024, 2048],
                            value=512,
                            interactive=True,
                        )
                        model_width = gr.Number(
                            label="Width (mm)",
                            minimum=0,
                            step=0.01,
                            value=100,
                            interactive=True,
                        )
                        model_height = gr.Number(
                            label="Height (mm)",
                            minimum=0,
                            step=0.01,
                            value=100,
                            interactive=True,
                        )
                        use_aspect_ratio = gr.Checkbox(
                            label="Use image aspect ratio (base/front)",
                            value=True,
                            interactive=True,
                        )
                        base_thickness = gr.Number(
                            label="Base Thickness (mm)",
                            minimum=0,
                            step=0.1,
                            value=1.0,
                        )
                        overlay_thickness = gr.Number(
                            label="Overlay Thickness (mm)",
                            minimum=0,
                            step=0.1,
                            value=0.5,
                        )

        with gr.Column():
            with gr.Tabs() as right_tabs:
                with gr.Tab("Preview", id="preview"):
                    front_out = gr.Image(label="Front Preview")
                    back_out = gr.Image(label="Back Preview")
                with gr.Tab("Masks", id="masks"):
                    mask_gallery_front = gr.Gallery(label="Front Masks")
                    mask_gallery_back = gr.Gallery(label="Back Masks")
                with gr.Tab("Model", id="model"):
                    animate_gif = gr.Image(
                        label="Animation GIF",
                        type="filepath",
                        interactive=False,
                    )
                    scad_file = gr.File(label="scad/3mf archive")
            generate_button = gr.Button(
                "Generate Model",
                variant="primary",
                interactive=False,
            )

    masks_state_front = gr.State()
    mask_colors_state_front = gr.State()
    masks_state_back = gr.State()
    mask_colors_state_back = gr.State()
    user_dir = gr.Text(visible=False)

    model_height.input(
        fn=update_width_from_height,
        inputs=[
            model_width,
            model_height,
            use_aspect_ratio,
            base_img,
            front_overlay_img,
        ],
        outputs=[model_width],
    )
    model_width.input(
        fn=update_height_from_width,
        inputs=[
            model_width,
            model_height,
            use_aspect_ratio,
            base_img,
            front_overlay_img,
        ],
        outputs=[model_height],
    )

    qr_text_front.submit(
        fn=generate_qr_code_from_text,
        inputs=[qr_text_front],
        outputs=[front_overlay_img],
    )
    qr_text_back.submit(
        fn=generate_qr_code_from_text,
        inputs=[qr_text_back],
        outputs=[back_overlay_img],
    )

    front_overlay_img.change(
        fn=overlay_front_images,
        inputs=[
            base_img,
            front_overlay_img,
            x_slider,
            y_slider,
            rot_slider,
            zoom_slider,
            indexed_slider,
            use_common_colors_checkbox,
            morphology_slider,
            resolution_dropdown,
        ],
        outputs=[
            front_out,
            mask_gallery_front,
            masks_state_front,
            mask_colors_state_front,
            generate_button,
        ],
    )

    back_overlay_img.change(
        fn=overlay_back_images,
        inputs=[
            base_img,
            back_overlay_img,
            x_slider_back,
            y_slider_back,
            rot_slider_back,
            zoom_slider_back,
            indexed_slider_back,
            use_common_colors_checkbox_back,
            morphology_slider_back,
            resolution_dropdown,
        ],
        outputs=[
            back_out,
            mask_gallery_back,
            masks_state_back,
            mask_colors_state_back,
            generate_button,
        ],
    )

    for comp in [
        x_slider,
        y_slider,
        rot_slider,
        zoom_slider,
        indexed_slider,
        use_common_colors_checkbox,
        morphology_slider,
        resolution_dropdown,
    ]:
        comp.change(
            fn=overlay_front_images,
            inputs=[
                base_img,
                front_overlay_img,
                x_slider,
                y_slider,
                rot_slider,
                zoom_slider,
                indexed_slider,
                use_common_colors_checkbox,
                morphology_slider,
                resolution_dropdown,
            ],
            outputs=[
                front_out,
                mask_gallery_front,
                masks_state_front,
                mask_colors_state_front,
                generate_button,
            ],
        )

    for comp in [
        x_slider_back,
        y_slider_back,
        rot_slider_back,
        zoom_slider_back,
        indexed_slider_back,
        use_common_colors_checkbox_back,
        morphology_slider_back,
        resolution_dropdown,
    ]:
        comp.change(
            fn=overlay_back_images,
            inputs=[
                base_img,
                back_overlay_img,
                x_slider_back,
                y_slider_back,
                rot_slider_back,
                zoom_slider_back,
                indexed_slider_back,
                use_common_colors_checkbox_back,
                morphology_slider_back,
                resolution_dropdown,
            ],
            outputs=[
                back_out,
                mask_gallery_back,
                masks_state_back,
                mask_colors_state_back,
                generate_button,
            ],
        )

    base_img.change(
        fn=overlay_front_images,
        inputs=[
            base_img,
            front_overlay_img,
            x_slider,
            y_slider,
            rot_slider,
            zoom_slider,
            indexed_slider,
            use_common_colors_checkbox,
            morphology_slider,
            resolution_dropdown,
        ],
        outputs=[
            front_out,
            mask_gallery_front,
            masks_state_front,
            mask_colors_state_front,
            generate_button,
        ],
    )

    base_img.change(
        fn=overlay_back_images,
        inputs=[
            base_img,
            back_overlay_img,
            x_slider_back,
            y_slider_back,
            rot_slider_back,
            zoom_slider_back,
            indexed_slider_back,
            use_common_colors_checkbox_back,
            morphology_slider_back,
            resolution_dropdown,
        ],
        outputs=[
            back_out,
            mask_gallery_back,
            masks_state_back,
            mask_colors_state_back,
            generate_button,
        ],
    )

    base_img.change(
        fn=update_slider_ranges,
        inputs=[base_img, front_overlay_img, resolution_dropdown],
        outputs=[
            x_slider,
            y_slider,
            rot_slider,
            zoom_slider,
            indexed_slider,
            model_width,
            model_height,
        ],
    )

    base_img.change(
        fn=update_slider_ranges,
        inputs=[base_img, back_overlay_img, resolution_dropdown],
        outputs=[
            x_slider_back,
            y_slider_back,
            rot_slider_back,
            zoom_slider_back,
            indexed_slider_back,
            model_width,
            model_height,
        ],
    )

    front_overlay_img.change(
        fn=update_height_from_width,
        inputs=[
            model_width,
            model_height,
            use_aspect_ratio,
            base_img,
            front_overlay_img,
        ],
        outputs=[model_height],
    )

    back_overlay_img.change(
        fn=update_height_from_width,
        inputs=[
            model_width,
            model_height,
            use_aspect_ratio,
            base_img,
            back_overlay_img,
        ],
        outputs=[model_height],
    )

    front_overlay_img.change(
        fn=lambda img, base, current_zoom, res: (
            gr.update(interactive=bool(img)),
            gr.update(interactive=bool(img)),
            gr.update(interactive=bool(img)),
            gr.update(
                interactive=bool(img),
                value=calculate_zoom_value(
                    base_img=base,
                    overlay_img=img,
                    current_zoom=current_zoom,
                    resolution=res,
                )
                if img
                else current_zoom,
            ),
        ),
        inputs=[front_overlay_img, base_img, zoom_slider, resolution_dropdown],
        outputs=[x_slider, y_slider, rot_slider, zoom_slider],
    )

    back_overlay_img.change(
        fn=lambda img, base, current_zoom, res: (
            gr.update(interactive=bool(img)),
            gr.update(interactive=bool(img)),
            gr.update(interactive=bool(img)),
            gr.update(
                interactive=bool(img),
                value=calculate_zoom_value(
                    base_img=base,
                    overlay_img=img,
                    current_zoom=current_zoom,
                    resolution=res,
                )
                if img
                else current_zoom,
            ),
        ),
        inputs=[
            back_overlay_img,
            base_img,
            zoom_slider_back,
            resolution_dropdown,
        ],
        outputs=[
            x_slider_back,
            y_slider_back,
            rot_slider_back,
            zoom_slider_back,
        ],
    )

    resolution_dropdown.change(
        fn=update_slider_ranges,
        inputs=[base_img, front_overlay_img, resolution_dropdown],
        outputs=[
            x_slider,
            y_slider,
            rot_slider,
            zoom_slider,
            indexed_slider,
            model_width,
            model_height,
        ],
    )
    resolution_dropdown.change(
        fn=update_slider_ranges,
        inputs=[base_img, back_overlay_img, resolution_dropdown],
        outputs=[
            x_slider_back,
            y_slider_back,
            rot_slider_back,
            zoom_slider_back,
            indexed_slider_back,
            model_width,
            model_height,
        ],
    )

    demo.load(get_user_dir, outputs=user_dir)
    demo.unload(delete_user_dir)

    generate_button.click(
        fn=switch_to_model_tab,
        inputs=[],
        outputs=right_tabs,
    ).then(
        fn=generate_3mf_with_frames,
        inputs=[
            base_img,
            model_width,
            model_height,
            base_thickness,
            overlay_thickness,
            resolution_dropdown,
            masks_state_front,
            mask_colors_state_front,
            masks_state_back,
            mask_colors_state_back,
            user_dir,
        ],
        outputs=[scad_file, animate_gif],
    )

if __name__ == "__main__":
    demo.launch(share=False, server_name="0.0.0.0")