app.py

# annotate_concat_demo.py
# pip install -U gradio pillow

import os
import time
from pathlib import Path
from typing import List, Tuple, Optional

import gradio as gr
from PIL import Image, ImageOps

# Your existing implementations are assumed available:
from unibox.utils.image_utils import (
    concatenate_images_horizontally,
    add_annotation,
)

# ------------------------- helpers -------------------------

def _norm_path(f) -> Optional[str]:
    if f is None:
        return None
    if isinstance(f, (str, Path)):
        return str(f)
    if hasattr(f, "name"):
        return str(getattr(f, "name"))
    if isinstance(f, dict):
        return str(f.get("name") or f.get("path") or "")
    return str(f)

def _load_images(files) -> List[Tuple[Image.Image, str]]:
    out: List[Tuple[Image.Image, str]] = []
    for f in (files or []):
        p = _norm_path(f)
        if not p:
            continue
        im = Image.open(p)
        # auto-orient, ensure RGB; supports PNG/JPEG/WebP/GIF/BMP/TIFF…
        im = ImageOps.exif_transpose(im).convert("RGB")
        out.append((im, os.path.basename(p)))
    return out

def _parse_descriptions(text: str, n: int):
    lines = (text or "").splitlines()
    if len(lines) > n:
        return None, f"Too many description lines ({len(lines)}) for {n} image(s). Provide ≤ one per image."
    lines = lines + [""] * (max(0, n - len(lines)))  # pad with blanks
    return lines[:n], None

def _build_stats(files, desc_text: str) -> dict:
    pairs = _load_images(files)
    n = len(pairs)
    lines, err = _parse_descriptions(desc_text, n) if n > 0 else ((desc_text or "").splitlines(), None)
    mapping = {}
    for i, (_, fname) in enumerate(pairs):
        mapping[fname] = (lines[i] if isinstance(lines, list) and i < len(lines) else "")
    return {
        "num_images": n,
        "num_descriptions": len((desc_text or "").splitlines()),
        "mapping": mapping,
        **({"error": err} if err else {}),
    }

# --------------------- core actions ------------------------

def concatenate_with_annotations(
    files,
    desc_text: str,
    max_height: int,
    position: str,
    alignment: str,
    size_adj: str,
):
    logs = []
    out_img = None
    out_file = None

    pairs = _load_images(files)
    if not pairs:
        logs.append("ERROR: Please upload at least one image.")
        return out_img, out_file, "\n".join(logs), _build_stats(files, desc_text)

    lines, err = _parse_descriptions(desc_text, len(pairs))
    if err:
        logs.append(f"ERROR: {err}")
        return out_img, out_file, "\n".join(logs), _build_stats(files, desc_text)

    # For left/right, alignment must be center (matches add_annotation behavior)
    if position in ("left", "right"):
        alignment = "center"

    annotated = []
    for (im, fname), line in zip(pairs, lines):
        if line.strip():
            im2 = add_annotation(
                pil_image=im,
                annotation=line,
                position=position,
                alignment=alignment,
                size=size_adj,
            )
            annotated.append(im2)
            logs.append(f"Annotated: {fname}")
        else:
            annotated.append(im)
            logs.append(f"Skipped (no description): {fname}")

    started = time.time()
    merged = concatenate_images_horizontally(annotated, max_height=max_height)
    if merged is None:
        logs.append("ERROR: Concatenation produced no result.")
        return None, None, "\n".join(logs), _build_stats(files, desc_text)

    # Save JPEG with required name
    out_dir = Path("outputs")
    out_dir.mkdir(parents=True, exist_ok=True)
    stamp = time.strftime("%Y%m%d_%H%M%S")
    out_name = f"concatenate_{stamp}.jpg"
    out_path = out_dir / out_name
    merged.save(str(out_path), format="JPEG", quality=95, optimize=True)

    w, h = merged.size
    size_bytes = out_path.stat().st_size
    latency = time.time() - started
    logs.append(f"Output: {out_name} — {w}×{h}px — {size_bytes} bytes — {latency:.3f}s")

    return merged, str(out_path), "\n".join(logs), _build_stats(files, desc_text)

def check_stats_only(files, desc_text: str, *_):
    stats = _build_stats(files, desc_text)
    log = f"Images: {stats.get('num_images', 0)}; Description lines: {stats.get('num_descriptions', 0)}"
    if "error" in stats:
        log += f"\nERROR: {stats['error']}"
    return None, None, log, stats

# ----------------------- UI wiring -------------------------

theme = gr.themes.Monochrome(primary_hue="slate", radius_size="sm")

with gr.Blocks(
    title="Annotated Concatenation — Demo",
    theme=theme,
    analytics_enabled=False,
) as demo:

    gr.Markdown("# Annotate & Concatenate Images")
    gr.Markdown(
        "Upload images (PNG/JPEG/WebP…), add one description per line (blank = skip), "
        "and concatenate horizontally. The output JPEG is named `concatenate_{timestamp}.jpg`."
    )

    with gr.Row(variant="panel"):
        with gr.Column(scale=2):
            files_in = gr.Files(
                label="Image files",
                # Explicit list ensures WebP acceptance across Gradio builds
                file_types=[
                    ".png", ".jpg", ".jpeg", ".webp", ".gif",
                    ".bmp", ".tif", ".tiff", ".jfif"
                ],
                type="filepath",
                interactive=True,
            )
            desc_in = gr.Textbox(
                label="Descriptions (one per line; blank lines allowed to skip)",
                placeholder="e.g.\nLeft image label\n\nRight image label",
                lines=8,
            )
            max_h = gr.Number(
                label="Max height (px) for concatenated image",
                value=1024,
                precision=0,
                minimum=64,
                interactive=True,
            )

            # Folded by default
            with gr.Accordion("Annotation settings", open=False):
                pos = gr.Dropdown(
                    label="Position",
                    choices=["top", "bottom", "left", "right"],
                    value="bottom",
                )
                align = gr.Radio(
                    label="Alignment (applies to top/bottom)",
                    choices=["left", "center", "right"],
                    value="center",
                )
                size_adj = gr.Radio(
                    label="Text size",
                    choices=["default", "larger", "smaller", "smallest", "largest"],
                    value="default",
                )

                def _toggle_align(p):
                    return gr.update(value="center", interactive=False) if p in ("left", "right") else gr.update(interactive=True)

                pos.change(_toggle_align, inputs=[pos], outputs=[align])

            with gr.Row():
                concat_btn = gr.Button("Concatenate image", variant="primary")
                stats_btn = gr.Button("Check stats")

        with gr.Column(scale=3):
            out_img = gr.Image(
                label="Concatenated image (preview)",
                interactive=False,
                format="jpeg",
                show_download_button=False,
            )
            download_file = gr.File(
                label="Download JPEG (named as saved)",
                interactive=False,
                height=72,  # compact
            )

            with gr.Accordion("Logs", open=False):
                logs_out = gr.Textbox(
                    label="Info / Errors",
                    lines=10,
                    interactive=False,
                )

            with gr.Accordion("Stats", open=False):
                stats_out = gr.JSON(label="Counts and current filename→description mapping")

    concat_btn.click(
        concatenate_with_annotations,
        inputs=[files_in, desc_in, max_h, pos, align, size_adj],
        outputs=[out_img, download_file, logs_out, stats_out],
        api_name="concatenate",
    )

    stats_btn.click(
        check_stats_only,
        inputs=[files_in, desc_in, max_h, pos, align, size_adj],
        outputs=[out_img, download_file, logs_out, stats_out],
        api_name="check_stats",
    )

if __name__ == "__main__":
    demo.queue(max_size=8).launch()