hintauc 0.1.0__py3-none-any.whl

hintauc/__init__.py

@@ -0,0 +1,66 @@
+"""hintauc — Deterministic region-based hint generation (DHT) and Hint-AUC
+evaluation for line-art colorization.
+
+Library form of the method proposed in
+"Hint-AUC: Deterministic Region-based Hint Generation for Line Art
+Colorization Evaluation" (Madono, Mingcheng, Simo-Serra).
+
+Quick start
+-----------
+>>> import hintauc
+>>> hints = hintauc.generate_hints("illustration.png")       # 1. hints
+>>> color, mask = hints.at_ratio(0.10, hint_type="scribble") # top-10% regions
+>>> hints.save("out/illustration")                           # canonical files
+
+>>> ev = hintauc.Evaluator(metrics=("mse", "psnr", "ssim", "lpips"))
+>>> scores = ev("colorized.png", "ground_truth.png")         # 2. evaluation
+>>> auc = hintauc.hint_auc({0.0: 0.44, 0.01: 0.33, 0.03: 0.27,
+...                         0.05: 0.24, 0.10: 0.20, 0.25: 0.15,
+...                         0.50: 0.12, 1.00: 0.10})
+"""
+
+from .longest_path import geodesic_longest_path
+from .hints import (
+    DEFAULT_HINT_SIZE,
+    FELZENSZWALB_PARAMS,
+    HintResult,
+    generate_hints,
+    region_ids,
+    segment_regions,
+)
+from .metrics import (
+    DEFAULT_METRICS,
+    LOWER_IS_BETTER,
+    Evaluator,
+    evaluate_dirs,
+    evaluate_pair,
+)
+from .auc import (
+    DEFAULT_ALPHAS,
+    evaluate_hint_curve,
+    hint_auc,
+    hint_auc_table,
+    trapz,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "DEFAULT_ALPHAS",
+    "DEFAULT_HINT_SIZE",
+    "DEFAULT_METRICS",
+    "FELZENSZWALB_PARAMS",
+    "LOWER_IS_BETTER",
+    "Evaluator",
+    "HintResult",
+    "evaluate_dirs",
+    "evaluate_hint_curve",
+    "evaluate_pair",
+    "generate_hints",
+    "hint_auc",
+    "hint_auc_table",
+    "region_ids",
+    "segment_regions",
+    "trapz",
+    "__version__",
+]

hintauc/auc.py

@@ -0,0 +1,73 @@
+"""Hint-AUC: integrate per-ratio metric scores over the hint-ratio grid.
+
+Port of ``evaluation/calc_hint_auc.py``: plain trapezoidal integration
+of metric(alpha) over alpha in [0, 1] (the alpha range has length 1, so the
+integral equals the range-normalized value).
+"""
+
+from __future__ import annotations
+
+from typing import Dict, Mapping, Optional, Sequence
+
+#: the paper's front-loaded hint-ratio grid
+DEFAULT_ALPHAS = (0.00, 0.01, 0.03, 0.05, 0.10, 0.25, 0.50, 1.00)
+
+
+def trapz(xs: Sequence[float], ys: Sequence[float]) -> float:
+    """Trapezoidal rule (port of calc_hint_auc.trapz)."""
+    if len(xs) != len(ys) or len(xs) < 2:
+        raise ValueError("need >= 2 (x, y) points with matching lengths")
+    area = 0.0
+    for i in range(1, len(xs)):
+        area += (xs[i] - xs[i - 1]) * (ys[i] + ys[i - 1]) / 2.0
+    return area
+
+
+def hint_auc(scores_by_alpha: Mapping[float, float]) -> float:
+    """Hint-AUC of one metric from ``{alpha: score}``.
+
+    >>> hint_auc({0.0: 0.44, 0.01: 0.33, ..., 1.0: 0.10})
+    """
+    xs = sorted(scores_by_alpha)
+    ys = [scores_by_alpha[a] for a in xs]
+    return trapz(xs, ys)
+
+
+def hint_auc_table(
+    metrics_by_alpha: Mapping[float, Mapping[str, float]],
+    metrics: Optional[Sequence[str]] = None,
+) -> Dict[str, float]:
+    """Hint-AUC per metric from ``{alpha: {metric: score}}``."""
+    alphas = sorted(metrics_by_alpha)
+    if metrics is None:
+        metrics = sorted({m for a in alphas for m in metrics_by_alpha[a]})
+    out = {}
+    for m in metrics:
+        out[m] = hint_auc({a: metrics_by_alpha[a][m] for a in alphas})
+    return out
+
+
+def evaluate_hint_curve(
+    preds_by_alpha: Mapping[float, str],
+    gt_dir: str,
+    evaluator=None,
+    metrics: Sequence[str] = ("mse", "psnr", "ssim"),
+    pairing: str = "sorted",
+    limit: int = 0,
+) -> Dict[str, object]:
+    """End-to-end Hint-AUC over per-ratio prediction directories.
+
+    ``preds_by_alpha`` maps each hint ratio alpha to a directory of the
+    colorizations produced with hints at that ratio (paper grid:
+    :data:`DEFAULT_ALPHAS`).  Returns per-alpha mean metrics and the
+    Hint-AUC of every metric.
+    """
+    from .metrics import Evaluator, evaluate_dirs
+
+    ev = evaluator or Evaluator(metrics=metrics)
+    per_alpha: Dict[float, Dict[str, float]] = {}
+    for alpha in sorted(preds_by_alpha):
+        per_alpha[alpha] = evaluate_dirs(
+            preds_by_alpha[alpha], gt_dir, evaluator=ev,
+            pairing=pairing, limit=limit)
+    return {"per_alpha": per_alpha, "hint_auc": hint_auc_table(per_alpha)}

hintauc/cli.py

@@ -0,0 +1,84 @@
+"""Command-line interface: ``hintauc generate`` / ``hintauc eval``."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+
+
+def main(argv=None):
+    p = argparse.ArgumentParser(
+        prog="hintauc",
+        description="Deterministic hint generation (DHT) and Hint-AUC evaluation")
+    sub = p.add_subparsers(dest="cmd", required=True)
+
+    g = sub.add_parser("generate", help="generate deterministic hints for an image")
+    g.add_argument("image", help="input image (ground-truth color image)")
+    g.add_argument("-o", "--out_stem", default=None,
+                   help="output stem (default: <image without extension>)")
+    g.add_argument("--size", type=int, default=64, help="hint resolution (default 64)")
+    g.add_argument("--segmenter", default="felzenszwalb",
+                   choices=["felzenszwalb", "slic", "quickshift"])
+    g.add_argument("--ratio", type=float, default=None,
+                   help="also write hints masked to this hint ratio (0..1)")
+    g.add_argument("--hint_type", default="scribble", choices=["scribble", "dot"])
+    g.add_argument("--path_method", default="filfinder", choices=["filfinder", "geodesic"],
+                   help="longest-path extraction: 'filfinder' (paper) or 'geodesic' (dependency-free, deterministic)")
+    g.add_argument("--tie_break", default="default", choices=["default", "stable"],
+                   help="order of equal-area regions when selecting by --ratio: 'default' (paper: NumPy argsort) or 'stable' (ascending label; version-independent)")
+    g.add_argument("--dot_method", default="medoid", choices=["medoid", "mean", "nearest_mean"],
+                   help="dot placement: 'mean' (paper: truncated mean of the path, may leave the region) or 'medoid' (in-region path pixel nearest to the mean)")
+    g.add_argument("-v", "--verbose", action="store_true")
+
+    e = sub.add_parser("eval", help="evaluate colorization(s) against ground truth")
+    e.add_argument("pred", help="predicted image or directory")
+    e.add_argument("gt", help="ground-truth image or directory")
+    e.add_argument("--metrics", nargs="+",
+                   default=["mse", "psnr", "ssim"],
+                   help="subset of: mse psnr ssim lpips openclip dino dreamsim")
+    e.add_argument("--pairing", default="sorted", choices=["sorted", "name"])
+    e.add_argument("--resize", type=int, default=256)
+    e.add_argument("--device", default=None)
+    e.add_argument("--limit", type=int, default=0)
+
+    args = p.parse_args(argv)
+
+    if args.cmd == "generate":
+        from .hints import generate_hints
+        import cv2
+
+        res = generate_hints(args.image, size=args.size,
+                             segmenter=args.segmenter, verbose=args.verbose,
+                             path_method=args.path_method, dot_method=args.dot_method)
+        stem = args.out_stem or os.path.splitext(args.image)[0]
+        paths = res.save(stem)
+        if args.ratio is not None:
+            color, mask = res.at_ratio(args.ratio, hint_type=args.hint_type, tie_break=args.tie_break)
+            pct = int(round(args.ratio * 100))
+            cv2.imwrite(f"{stem}_{args.hint_type}_col{res.size}_r{pct}.png", color)
+            cv2.imwrite(f"{stem}_{args.hint_type}_mask{res.size}_r{pct}.png", mask)
+            paths[f"ratio_{pct}"] = f"{stem}_{args.hint_type}_*{res.size}_r{pct}.png"
+        info = {"n_regions": res.n_regions(), "failed_regions": res.failed_regions,
+                "path_method": res.path_method, "dot_method": res.dot_method, "outputs": paths}
+        print(json.dumps(info, indent=1))
+        return 0
+
+    if args.cmd == "eval":
+        from .metrics import Evaluator, evaluate_dirs
+
+        ev = Evaluator(metrics=args.metrics, device=args.device, resize=args.resize)
+        if os.path.isdir(args.pred):
+            scores = evaluate_dirs(args.pred, args.gt, evaluator=ev,
+                                   pairing=args.pairing, limit=args.limit)
+        else:
+            scores = ev(args.pred, args.gt)
+        print(json.dumps(scores, indent=1))
+        return 0
+
+    return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())