trace-digitiser 0.1.0__py3-none-any.whl

trace_digitiser/__init__.py

@@ -0,0 +1,222 @@
+"""trace_digitiser — template-free digitisation of raster scientific line plots.
+
+Quick start::
+
+    from trace_digitiser import digitise
+
+    result = digitise(
+        "figure.jpg",
+        layout_mode="stacked",
+        expected_rows=2,
+        expected_cols=1,
+        output_dir="outputs",
+    )
+
+    print(result.trace_data.head())
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional
+
+import pandas as pd
+
+from .calibration import propagate_y_calibration_across_rows, robust_y_calibration
+from .diagnostics import draw_digitised_trace, draw_panel_overlay, draw_trace_mask, print_panel_summary
+from .digitise import digitise_trace_mask
+from .io import build_panel_metadata, load_image, save_outputs
+from .models import Calibration, DigitiserResult, Panel, Trace, XLabel
+from .ocr import detect_x_labels
+from .panel_detection import find_plot_panels
+from .summarise import summarise_by_detected_labels
+from .trace_detection import detect_trace_masks
+from .x_calibration import calibrate_x_axis
+
+__all__ = [
+    "digitise",
+    "Calibration",
+    "DigitiserResult",
+    "Panel",
+    "Trace",
+    "XLabel",
+]
+
+
+def digitise(
+    image_path: str | Path,
+    *,
+    layout_mode: str = "auto",
+    expected_rows: Optional[int] = None,
+    expected_cols: Optional[int] = None,
+    expected_panels: Optional[int] = None,
+    output_dir: Optional[str | Path] = None,
+    output_prefix: Optional[str] = None,
+    show_debug: bool = False,
+    save_diagnostics: bool = False,
+) -> DigitiserResult:
+    """End-to-end chart digitisation pipeline.
+
+    Parameters
+    ----------
+    image_path : str or Path
+        Path to the input raster image.
+    layout_mode : str
+        ``"auto"``, ``"single"``, ``"stacked"``, ``"horizontal"``, or
+        ``"grid"``.
+    expected_rows, expected_cols, expected_panels : int, optional
+        Layout hints that constrain panel selection.
+    output_dir : str or Path, optional
+        Directory for CSV outputs.  Defaults to current directory.
+    output_prefix : str, optional
+        Prefix for output filenames.  Defaults to the image stem.
+    show_debug : bool
+        If True, display inline diagnostic plots (for interactive use).
+    save_diagnostics : bool
+        If True, write diagnostic overlay PNGs to *output_dir*.
+
+    Returns
+    -------
+    DigitiserResult
+        Structured result with panels, traces, DataFrames, and paths.
+    """
+    image_path = Path(image_path)
+    rgb = load_image(image_path)
+
+    if output_prefix is None:
+        output_prefix = image_path.stem
+
+    diag_dir: Optional[Path] = None
+    if save_diagnostics:
+        diag_dir = Path(output_dir or ".") / "diagnostics"
+        diag_dir.mkdir(parents=True, exist_ok=True)
+
+    if show_debug:
+        print("Processing:", image_path)
+        print("Image size:", rgb.shape[1], "×", rgb.shape[0])
+
+    # ------------------------------------------------------------------
+    # 1. Detect panels
+    # ------------------------------------------------------------------
+    panels, h_lines, v_lines = find_plot_panels(
+        rgb,
+        layout_mode=layout_mode,
+        expected_rows=expected_rows,
+        expected_cols=expected_cols,
+        expected_panels=expected_panels,
+    )
+
+    if show_debug or save_diagnostics:
+        print_panel_summary(panels, h_lines, v_lines, layout_mode)
+        draw_panel_overlay(rgb, panels, h_lines, v_lines, output_dir=diag_dir, show=show_debug)
+
+    # ------------------------------------------------------------------
+    # 2. Y-axis calibration
+    # ------------------------------------------------------------------
+    for p in panels:
+        calib = robust_y_calibration(rgb, p, verbose=show_debug)
+        p["y_calibration"] = calib.to_dict()
+
+    panels = propagate_y_calibration_across_rows(panels, verbose=show_debug)
+
+    # ------------------------------------------------------------------
+    # 2b. X-axis calibration (numeric x ticks)
+    # ------------------------------------------------------------------
+    for p in panels:
+        x_cal = calibrate_x_axis(rgb, p, verbose=show_debug)
+        if x_cal is not None:
+            p["x_calibration"] = x_cal
+
+    # ------------------------------------------------------------------
+    # 3. Trace detection and digitisation
+    # ------------------------------------------------------------------
+    all_trace_frames: list[pd.DataFrame] = []
+    trace_debug: list[tuple[dict, dict]] = []
+
+    for p in panels:
+        masks = detect_trace_masks(rgb, p)
+        if show_debug:
+            print(f"Panel {p['panel_id']}: detected {len(masks)} coloured trace(s)")
+
+        for tr in masks:
+            if show_debug:
+                print(" ", {k: v for k, v in tr.items() if k != "mask"})
+            all_trace_frames.append(digitise_trace_mask(p, tr))
+            trace_debug.append((p, tr))
+
+    trace_data = pd.concat(all_trace_frames, ignore_index=True) if all_trace_frames else pd.DataFrame()
+
+    # ------------------------------------------------------------------
+    # 4. X-label OCR
+    # ------------------------------------------------------------------
+    for p in panels:
+        p["x_labels"] = detect_x_labels(rgb, p)
+        if show_debug:
+            print(f"Panel {p['panel_id']} x labels:")
+            for lab in p["x_labels"]:
+                print(f"  {lab['text']:>8s}  x={lab['x']:.1f}  conf={lab['conf']:.1f}")
+
+    # ------------------------------------------------------------------
+    # 5. Interval summaries
+    # ------------------------------------------------------------------
+    summary_by_label = summarise_by_detected_labels(trace_data, panels)
+
+    # ------------------------------------------------------------------
+    # 6. Diagnostics
+    # ------------------------------------------------------------------
+    if show_debug or save_diagnostics:
+        for p, tr in trace_debug:
+            draw_trace_mask(rgb, p, tr, output_dir=diag_dir, show=show_debug)
+
+        if not trace_data.empty:
+            for (panel_id, trace_id), _ in trace_data.groupby(["panel_id", "trace_id"]):
+                draw_digitised_trace(trace_data, panel_id, trace_id, output_dir=diag_dir, show=show_debug)
+
+    # ------------------------------------------------------------------
+    # 7. Save outputs
+    # ------------------------------------------------------------------
+    panel_metadata = build_panel_metadata(panels)
+    trace_csv, meta_csv, summary_csv = save_outputs(
+        trace_data, panel_metadata, summary_by_label, output_prefix, output_dir
+    )
+
+    if show_debug:
+        print("Wrote:")
+        print(" -", trace_csv)
+        print(" -", meta_csv)
+        if summary_csv:
+            print(" -", summary_csv)
+        else:
+            print(" - no label summary; fewer than two labels detected")
+
+    # ------------------------------------------------------------------
+    # 8. Build structured result
+    # ------------------------------------------------------------------
+    return DigitiserResult(
+        image_path=image_path,
+        rgb=rgb,
+        panels=[
+            Panel(
+                panel_id=p["panel_id"],
+                x0=p["x0"],
+                x1=p["x1"],
+                y_top=p["y_top"],
+                y_bottom=p["y_bottom"],
+                gridline_y=p["gridline_y"],
+                source=p["source"],
+                score=p["score"],
+                layout_mode=p.get("layout_mode", layout_mode),
+                calibration=Calibration(**p["y_calibration"]) if "y_calibration" in p else None,
+                x_labels=[
+                    XLabel(**lab) for lab in p.get("x_labels", [])
+                ],
+            )
+            for p in panels
+        ],
+        trace_data=trace_data,
+        summary_by_label=summary_by_label,
+        panel_metadata=panel_metadata,
+        trace_csv_path=trace_csv,
+        summary_csv_path=summary_csv,
+        metadata_csv_path=meta_csv,
+    )

trace_digitiser/calibration.py

@@ -0,0 +1,283 @@
+"""Y-axis calibration: OCR-based linear fitting and cross-panel propagation.
+
+The calibration model is::
+
+    y_value = a * y_pixel + b
+
+Because image y-coordinates increase downward, normal scientific axes
+will have a negative ``a``.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from .models import Calibration
+from .ocr import ocr_number_near_y
+
+
+# ===================================================================
+# Orientation guard
+# ===================================================================
+
+def orient_calibration_upward(a: float, b: float, panel: dict) -> tuple[float, float, bool]:
+    """Ensure the calibration maps the top of the panel to a higher value.
+
+    Returns ``(a, b, was_flipped)``.
+    """
+    top_y = panel["y_top"]
+    bottom_y = panel["y_bottom"]
+    top_val = a * top_y + b
+    bottom_val = a * bottom_y + b
+
+    if top_val >= bottom_val:
+        return float(a), float(b), False
+
+    new_top_val = max(top_val, bottom_val)
+    new_bottom_val = min(top_val, bottom_val)
+    new_a = (new_bottom_val - new_top_val) / max(1, bottom_y - top_y)
+    new_b = new_top_val - new_a * top_y
+    return float(new_a), float(new_b), True
+
+
+# ===================================================================
+# Robust y-axis calibration
+# ===================================================================
+
+def robust_y_calibration(
+    rgb: np.ndarray,
+    panel: dict,
+    verbose: bool = False,
+) -> Calibration:
+    """Fit a linear y-pixel → y-value mapping for one panel.
+
+    Steps:
+
+    1. OCR numbers near all detected gridlines.
+    2. Try candidate line fits between OCR pairs (linear and log).
+    3. Score fits by number of inlier gridlines.
+    4. Refit from inliers if possible.
+    5. Enforce upward orientation.
+    6. Fall back to normalised 0–1 scale if OCR is insufficient.
+
+    Log-scale detection: if the OCR'd values are all positive and
+    better explained by ``log10(y_value) = a * y_pixel + b``, the
+    calibration uses ``scale_type = "log_ocr_value"``.
+    """
+    y_lines = panel["gridline_y"]
+    ocr: dict[int, list[float]] = {int(y): ocr_number_near_y(rgb, panel, int(y)) for y in y_lines}
+
+    pairs: list[tuple[float, float]] = []
+    for y, vals in ocr.items():
+        for v in vals:
+            if abs(v) < 10_000:
+                pairs.append((float(y), float(v)))
+
+    best: dict | None = None
+    best_log: dict | None = None
+
+    if len(pairs) >= 2:
+        for i in range(len(pairs)):
+            y1, v1 = pairs[i]
+            for j in range(i + 1, len(pairs)):
+                y2, v2 = pairs[j]
+                if abs(y2 - y1) < 20 or abs(v2 - v1) < 1:
+                    continue
+
+                # --- Linear fit ---
+                a = (v2 - v1) / (y2 - y1)
+                b = v1 - a * y1
+
+                # Normal y-axis: values increase upward → a < 0.
+                if a < 0:
+                    inliers: list[tuple[float, float]] = []
+                    score = 0
+                    total_error = 0.0
+
+                    for yy, vals in ocr.items():
+                        if not vals:
+                            continue
+                        pred = a * yy + b
+                        errors = [abs(pred - vv) for vv in vals]
+                        k = int(np.argmin(errors))
+                        err = errors[k]
+                        tol = max(18.0, 0.08 * abs(v2 - v1))
+
+                        if err <= tol:
+                            score += 1
+                            total_error += err
+                            inliers.append((float(yy), float(vals[k])))
+
+                    if best is None or (score, -total_error) > (best["score"], -best["total_error"]):
+                        best = {
+                            "a": a, "b": b,
+                            "score": score, "total_error": total_error,
+                            "inliers": inliers,
+                        }
+
+                # --- Log fit: log10(value) = a_log * y_pixel + b_log ---
+                if v1 > 0 and v2 > 0:
+                    lv1, lv2 = np.log10(v1), np.log10(v2)
+                    if abs(lv2 - lv1) < 0.3:
+                        continue
+                    a_log = (lv2 - lv1) / (y2 - y1)
+                    b_log = lv1 - a_log * y1
+
+                    if a_log >= 0:
+                        continue
+
+                    log_inliers: list[tuple[float, float]] = []
+                    log_score = 0
+                    log_error = 0.0
+
+                    for yy, vals in ocr.items():
+                        if not vals:
+                            continue
+                        pred_log = a_log * yy + b_log
+                        pred_val = 10 ** pred_log
+                        errors = [abs(pred_val - vv) / max(1e-9, abs(vv)) for vv in vals if vv > 0]
+                        if not errors:
+                            continue
+                        k = int(np.argmin(errors))
+                        rel_err = errors[k]
+
+                        if rel_err <= 0.20:  # 20 % relative tolerance
+                            log_score += 1
+                            log_error += rel_err
+                            pos_vals = [vv for vv in vals if vv > 0]
+                            log_inliers.append((float(yy), float(pos_vals[k])))
+
+                    if log_score >= 2:
+                        if best_log is None or (log_score, -log_error) > (best_log["score"], -best_log["total_error"]):
+                            best_log = {
+                                "a": a_log, "b": b_log,
+                                "score": log_score, "total_error": log_error,
+                                "inliers": log_inliers,
+                            }
+
+    orientation_flipped = False
+    use_log = False
+
+    # Decide between linear and log.
+    if best_log is not None and best_log["score"] >= 2:
+        linear_score = best["score"] if best is not None else 0
+        if best_log["score"] > linear_score or (best_log["score"] == linear_score and best is None):
+            use_log = True
+
+    if use_log and best_log is not None:
+        inlier_arr = np.array(best_log["inliers"], dtype=float)
+        log_vals = np.log10(inlier_arr[:, 1])
+        if len(inlier_arr) >= 2:
+            a, b = np.polyfit(inlier_arr[:, 0], log_vals, deg=1)
+        else:
+            a, b = best_log["a"], best_log["b"]
+        scale_type = "log_ocr_value"
+        # For log scale, "upward" means higher values at top → a < 0.
+        if a > 0:
+            a, b = -a, -(b)  # flip
+            orientation_flipped = True
+        best_used = best_log
+    elif best is not None and best["score"] >= 2:
+        inlier_arr = np.array(best["inliers"], dtype=float)
+        if len(inlier_arr) >= 2:
+            a, b = np.polyfit(inlier_arr[:, 0], inlier_arr[:, 1], deg=1)
+        else:
+            a, b = best["a"], best["b"]
+
+        scale_type = "ocr_value"
+        a, b, orientation_flipped = orient_calibration_upward(a, b, panel)
+        best_used = best
+    else:
+        a = -1.0 / max(1, panel["y_bottom"] - panel["y_top"])
+        b = -a * panel["y_bottom"]
+        scale_type = "normalised_0_to_1"
+        best_used = {"score": 0, "inliers": []}
+        a, b, orientation_flipped = orient_calibration_upward(a, b, panel)
+
+    if verbose:
+        print(f"Panel {panel['panel_id']} y-axis OCR candidates:")
+        for y in y_lines:
+            print(f"  y={y}: {ocr[y]}")
+        if scale_type == "log_ocr_value":
+            print(f"  calibration: log10(value) = {a:.6g} * y_pixel + {b:.6g}")
+            print(f"  top value: {10 ** (a * panel['y_top'] + b):.3g}")
+            print(f"  bottom value: {10 ** (a * panel['y_bottom'] + b):.3g}")
+        else:
+            print(f"  calibration: value = {a:.6g} * y_pixel + {b:.6g}")
+            print(f"  top value: {a * panel['y_top'] + b:.3g}")
+            print(f"  bottom value: {a * panel['y_bottom'] + b:.3g}")
+        print(f"  scale_type: {scale_type}")
+        print(f"  inliers: {best_used['inliers']}")
+
+    return Calibration(
+        a=float(a),
+        b=float(b),
+        scale_type=scale_type,
+        orientation_flipped=orientation_flipped,
+        ocr_candidates=ocr,
+        inliers=best["inliers"],
+    )
+
+
+# ===================================================================
+# Cross-row propagation
+# ===================================================================
+
+def propagate_y_calibration_across_rows(
+    panels: list[dict],
+    verbose: bool = False,
+) -> list[dict]:
+    """Copy y-axis calibration from labelled panels to row-aligned siblings.
+
+    In horizontal multi-panel figures only the leftmost subplot may have
+    y-axis labels.  This copies its calibration to sibling panels at
+    the same y-pixel extents.
+    """
+    if not panels:
+        return panels
+
+    used: set[int] = set()
+    row_groups: list[list[dict]] = []
+
+    for i, p in enumerate(panels):
+        if i in used:
+            continue
+        group = [p]
+        used.add(i)
+        for j, q in enumerate(panels):
+            if j in used:
+                continue
+            same_row = abs(p["y_top"] - q["y_top"]) < 18 and abs(p["y_bottom"] - q["y_bottom"]) < 18
+            if same_row:
+                group.append(q)
+                used.add(j)
+        row_groups.append(group)
+
+    for group in row_groups:
+        refs = [p for p in group if p["y_calibration"]["scale_type"] == "ocr_value"]
+        if not refs:
+            continue
+
+        ref = refs[0]
+        ref_cal = ref["y_calibration"]
+        ref_top_val = ref_cal["a"] * ref["y_top"] + ref_cal["b"]
+        ref_bottom_val = ref_cal["a"] * ref["y_bottom"] + ref_cal["b"]
+
+        for p in group:
+            if p["y_calibration"]["scale_type"] == "ocr_value":
+                continue
+
+            a = (ref_bottom_val - ref_top_val) / max(1, p["y_bottom"] - p["y_top"])
+            b = ref_top_val - a * p["y_top"]
+
+            p["y_calibration"] = {
+                **p["y_calibration"],
+                "a": float(a),
+                "b": float(b),
+                "scale_type": f"propagated_from_panel_{ref['panel_id']}",
+            }
+
+            if verbose:
+                print(f"Panel {p['panel_id']}: propagated y calibration from panel {ref['panel_id']}")
+
+    return panels

trace_digitiser/cli.py

@@ -0,0 +1,123 @@
+"""Command-line interface for trace_digitiser.
+
+Usage::
+
+    trace-digitiser figure.jpg --layout stacked --rows 2 --cols 1
+
+    # Batch mode
+    trace-digitiser figures/*.jpg --layout auto --output-dir results/
+
+    # Generate synthetic test figures
+    trace-digitiser --generate-test-figures --output-dir test_figures/
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(
+        prog="trace-digitiser",
+        description="Template-free digitisation of raster scientific line plots.",
+    )
+
+    p.add_argument(
+        "images",
+        nargs="*",
+        type=Path,
+        help="Input image file(s) to digitise.",
+    )
+    p.add_argument(
+        "--layout",
+        dest="layout_mode",
+        default="auto",
+        choices=["auto", "single", "stacked", "horizontal", "grid"],
+        help="Panel layout hint (default: auto).",
+    )
+    p.add_argument("--rows", dest="expected_rows", type=int, default=None, help="Expected number of panel rows.")
+    p.add_argument("--cols", dest="expected_cols", type=int, default=None, help="Expected number of panel columns.")
+    p.add_argument("--panels", dest="expected_panels", type=int, default=None, help="Expected total panel count.")
+    p.add_argument(
+        "--output-dir",
+        "-o",
+        dest="output_dir",
+        type=Path,
+        default=None,
+        help="Output directory for CSV files and diagnostics.",
+    )
+    p.add_argument(
+        "--save-diagnostics",
+        action="store_true",
+        help="Write diagnostic overlay images.",
+    )
+    p.add_argument(
+        "--generate-test-figures",
+        action="store_true",
+        help="Generate synthetic test plots instead of digitising.",
+    )
+    p.add_argument(
+        "--seed",
+        type=int,
+        default=123,
+        help="Random seed for synthetic test generation (default: 123).",
+    )
+
+    return p
+
+
+def main(argv: list[str] | None = None) -> None:
+    parser = _build_parser()
+    args = parser.parse_args(argv)
+
+    # --- Generate test figures -------------------------------------------
+    if args.generate_test_figures:
+        from .synthetic import generate_noisy_example_plots
+
+        out = args.output_dir or Path("test_figures")
+        paths, truth = generate_noisy_example_plots(output_dir=out, seed=args.seed, show=False)
+        print("Generated test figures:")
+        for p in paths:
+            print(f"  {p}")
+        print(f"Ground truth: {truth}")
+        return
+
+    # --- Digitise images -------------------------------------------------
+    if not args.images:
+        parser.print_help()
+        sys.exit(1)
+
+    from . import digitise
+
+    for image_path in args.images:
+        if not image_path.exists():
+            print(f"WARNING: {image_path} not found, skipping.", file=sys.stderr)
+            continue
+
+        print(f"\n{'=' * 60}")
+        print(f"Processing: {image_path}")
+
+        result = digitise(
+            image_path,
+            layout_mode=args.layout_mode,
+            expected_rows=args.expected_rows,
+            expected_cols=args.expected_cols,
+            expected_panels=args.expected_panels,
+            output_dir=args.output_dir,
+            show_debug=False,
+            save_diagnostics=args.save_diagnostics,
+        )
+
+        n_panels = len(result.panels)
+        n_traces = result.trace_data["trace_id"].nunique() if not result.trace_data.empty else 0
+        print(f"  Panels: {n_panels}, Traces: {n_traces}")
+        print(f"  Trace CSV:    {result.trace_csv_path}")
+        print(f"  Metadata CSV: {result.metadata_csv_path}")
+        if result.summary_csv_path:
+            print(f"  Summary CSV:  {result.summary_csv_path}")
+
+
+if __name__ == "__main__":
+    main()