"""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())