visualize_layout.py

"""Visualize PP-DocLayoutV3 results served over HTTP.

Sends each image to a running server (serve_pp_doclayout_v3.py), draws the
returned boxes/polygons on top of it — a small numbered badge per box marks
its reading order, and a legend strip maps each label's color to its name —
and writes both the annotated image and the raw JSON response to disk (JSON
carries the per-element score; the image would get unreadable if it tried to
print full "order:label score" banners on every box — dense pages can have
30+ small, tightly packed regions where that text is bigger than the box
itself and blots out neighbours). Talks to the server only over HTTP — no
onnxruntime/torch import needed here.

Usage:
    python serve_pp_doclayout_v3.py --onnx pp_doclayoutv3.onnx &
    python visualize_layout.py --images "examples/inputs/*.jpg" --out-dir examples/outputs
    python visualize_layout.py --images page.png --url http://localhost:8000/v1/layout --threshold 0.4
"""

from __future__ import annotations

import argparse
import colorsys
import glob
import json
import math
from pathlib import Path
from typing import Any

import cv2
import numpy as np
import requests

GOLDEN_RATIO_CONJUGATE = 0.618033988749895
FONT = cv2.FONT_HERSHEY_SIMPLEX
BADGE_FONT_SCALE = 0.35
BADGE_RADIUS_MIN = 9
LEGEND_ROW_H = 22
LEGEND_COL_W = 140
LEGEND_PAD = 8


def label_color(label_id: int) -> tuple[int, int, int]:
    """Deterministic, well-separated BGR color per label id (golden-angle hue spacing)."""
    hue = (label_id * GOLDEN_RATIO_CONJUGATE) % 1.0
    r, g, b = colorsys.hsv_to_rgb(hue, 0.65, 0.95)
    return int(b * 255), int(g * 255), int(r * 255)


def draw_badge(img: np.ndarray, order: int, x: int, y: int, color: tuple[int, int, int]) -> None:
    """Small filled circle with the reading-order number, clamped so it stays on-canvas."""
    h, w = img.shape[:2]
    text = str(order)
    (tw, th), _ = cv2.getTextSize(text, FONT, BADGE_FONT_SCALE, 1)
    radius = max(BADGE_RADIUS_MIN, tw // 2 + 3)
    cx, cy = min(max(x, radius), w - radius - 1), min(max(y, radius), h - radius - 1)
    cv2.circle(img, (cx, cy), radius, color, -1, cv2.LINE_AA)
    cv2.putText(img, text, (cx - tw // 2, cy + th // 2), FONT, BADGE_FONT_SCALE, (255, 255, 255), 1, cv2.LINE_AA)


def build_legend(elements: list[dict], width: int) -> np.ndarray | None:
    """White strip mapping each label present to its color, for appending below the image."""
    entries = sorted({(el["label_id"], el["label"]) for el in elements})
    if not entries:
        return None
    cols = max(1, width // LEGEND_COL_W)
    rows = math.ceil(len(entries) / cols)
    legend = np.full((LEGEND_PAD * 2 + rows * LEGEND_ROW_H, width, 3), 255, dtype=np.uint8)
    for i, (label_id, label) in enumerate(entries):
        row, col = divmod(i, cols)
        x, y = LEGEND_PAD + col * LEGEND_COL_W, LEGEND_PAD + row * LEGEND_ROW_H
        cv2.rectangle(legend, (x, y + 4), (x + 14, y + 18), label_color(label_id), -1)
        cv2.putText(legend, label, (x + 20, y + 16), FONT, 0.42, (20, 20, 20), 1, cv2.LINE_AA)
    return legend


def draw(image_path: Path, elements: list[dict], out_path: Path, show_polygons: bool) -> None:
    img = cv2.imread(str(image_path))
    if img is None:
        raise FileNotFoundError(f"could not read image: {image_path}")

    for el in elements:
        color = label_color(el["label_id"])
        if show_polygons and el.get("polygon"):
            poly = np.asarray(el["polygon"], dtype=np.int32).reshape(-1, 1, 2)
            cv2.polylines(img, [poly], True, color, 2, cv2.LINE_AA)
        else:
            x0, y0, x1, y1 = map(int, el["box"])
            cv2.rectangle(img, (x0, y0), (x1, y1), color, 2)

        x0, y0 = int(el["box"][0]), int(el["box"][1])
        draw_badge(img, el["order"], x0, y0, color)

    legend = build_legend(elements, img.shape[1])
    if legend is not None:
        img = np.vstack([img, legend])

    out_path.parent.mkdir(parents=True, exist_ok=True)
    cv2.imwrite(str(out_path), img)


def call_server(url: str, image_path: Path, threshold: float, polygons: bool, timeout: float) -> dict[str, Any]:
    with open(image_path, "rb") as f:
        files = {"file": (image_path.name, f, "application/octet-stream")}
        params = {"threshold": threshold, "polygons": str(polygons).lower()}
        resp = requests.post(url, files=files, params=params, timeout=timeout)
    resp.raise_for_status()
    return resp.json()


def resolve_images(patterns: list[str]) -> list[Path]:
    paths: list[Path] = []
    for pattern in patterns:
        matched = sorted(glob.glob(pattern))
        if matched:
            paths.extend(Path(m) for m in matched)
        else:
            paths.append(Path(pattern))
    return paths


def main() -> int:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--images", nargs="+", required=True, help="image paths and/or glob patterns")
    p.add_argument("--url", default="http://localhost:8000/v1/layout")
    p.add_argument("--out-dir", type=Path, default=Path("examples/outputs"))
    p.add_argument("--threshold", type=float, default=0.5)
    p.add_argument("--no-polygons", dest="polygons", action="store_false")
    p.add_argument("--timeout", type=float, default=60.0)
    args = p.parse_args()

    paths = resolve_images(args.images)
    if not paths:
        raise SystemExit(f"no images matched: {args.images}")

    args.out_dir.mkdir(parents=True, exist_ok=True)
    summary = []
    for path in paths:
        print(f"-> {path.name}")
        result = call_server(args.url, path, args.threshold, args.polygons, args.timeout)
        elements = result["results"][0]["elements"]

        json_path = args.out_dir / f"{path.stem}.json"
        json_path.write_text(json.dumps(result, indent=2))

        vis_path = args.out_dir / f"{path.stem}_annotated.jpg"
        draw(path, elements, vis_path, show_polygons=args.polygons)

        print(f"   {len(elements)} elements, {result['latency_ms']} ms -> {vis_path.name}, {json_path.name}")
        summary.append(
            {"image": path.name, "annotated": vis_path.name, "json": json_path.name, "count": len(elements)}
        )

    (args.out_dir / "summary.json").write_text(json.dumps(summary, indent=2))
    print(f"\n{len(summary)} image(s) -> {args.out_dir}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())