streamlit_viewer.py

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import altair as alt
import pandas as pd
import streamlit as st


ROOT = Path(__file__).resolve().parent
CONFIG_PATH = ROOT / "benchmark_config.json"
INPUTS_DIR = ROOT / "inputs"
RESULTS_DIR = ROOT / "results"
MODEL_ORDER = ["retab-micro", "retab-small", "retab-large"]
TYPED_IOU_THRESHOLD = 0.8
BOUNDARY_TOLERANCE = 1
POLITAX_COLOR_RANGE = [
    "#4e79a7",
    "#f28e2b",
    "#e15759",
    "#76b7b2",
    "#59a14f",
    "#edc949",
    "#af7aa1",
    "#ff9da7",
    "#9c755f",
    "#bab0ab",
    "#1f77b4",
    "#ff7f0e",
    "#2ca02c",
    "#d62728",
    "#9467bd",
    "#8c564b",
    "#e377c2",
    "#7f7f7f",
    "#bcbd22",
    "#17becf",
    "#393b79",
    "#637939",
    "#8c6d31",
    "#843c39",
    "#7b4173",
    "#3182bd",
    "#31a354",
    "#756bb1",
    "#636363",
    "#e6550d",
    "#6baed6",
    "#74c476",
    "#9e9ac8",
    "#969696",
    "#fd8d3c",
    "#9ecae1",
    "#a1d99b",
    "#bcbddc",
    "#bdbdbd",
    "#fdae6b",
]


def load_json(path: Path) -> Any:
    return json.loads(path.read_text())


@st.cache_data(ttl=10)
def load_inputs() -> dict[str, Any]:
    config = load_json(CONFIG_PATH)
    documents = [load_json(path) for path in sorted(INPUTS_DIR.glob("*.json"))]
    return {**config, "documents": documents}


def result_runs() -> list[str]:
    runs = []
    if (RESULTS_DIR / "article_snapshot").exists():
        runs.append("article_snapshot")
    live_dir = RESULTS_DIR / "live"
    if live_dir.exists():
        runs.extend(f"live/{p.name}" for p in sorted(live_dir.iterdir()) if p.is_dir())
    return runs


def run_dir(run_name: str) -> Path:
    if run_name == "article_snapshot":
        return RESULTS_DIR / "article_snapshot"
    if run_name.startswith("live/"):
        return RESULTS_DIR / run_name
    raise ValueError(run_name)


def load_result_index(run_name: str) -> dict[tuple[str, str], Path]:
    base = run_dir(run_name)
    index: dict[tuple[str, str], Path] = {}
    for path in sorted(base.glob("*/*.json")):
        split = load_json(path)
        filename = split.get("file", {}).get("filename")
        model = split.get("model")
        if filename and model:
            index[(filename, model)] = path
    return index


def segments_from_output(output: list[dict[str, Any]]) -> list[dict[str, Any]]:
    segments: list[dict[str, Any]] = []
    for split in output:
        pages = sorted(int(page) for page in split.get("pages", []))
        if not pages:
            continue
        start = pages[0]
        previous = pages[0]
        for page in pages[1:]:
            if page == previous + 1:
                previous = page
                continue
            segments.append(
                {
                    "type": split["name"],
                    "start_page": start,
                    "end_page": previous,
                }
            )
            start = page
            previous = page
        segments.append(
            {
                "type": split["name"],
                "start_page": start,
                "end_page": previous,
            }
        )
    return sorted(
        segments,
        key=lambda segment: (int(segment["start_page"]), int(segment["end_page"]), segment["type"]),
    )


def page_labels(segments: list[dict[str, Any]], page_count: int) -> list[str]:
    labels = [""] * (page_count + 1)
    for segment in segments:
        start = max(1, int(segment["start_page"]))
        end = min(page_count, int(segment["end_page"]))
        for page in range(start, end + 1):
            labels[page] = segment["type"]
    return labels


def f1(matches: int, predicted: int, truth: int) -> float:
    if predicted == 0 and truth == 0:
        return 1.0
    denominator = predicted + truth
    return 0.0 if denominator == 0 else (2 * matches) / denominator


def segment_iou(left: dict[str, Any], right: dict[str, Any]) -> float:
    start = max(int(left["start_page"]), int(right["start_page"]))
    end = min(int(left["end_page"]), int(right["end_page"]))
    intersection = max(0, end - start + 1)
    if intersection == 0:
        return 0.0
    left_length = int(left["end_page"]) - int(left["start_page"]) + 1
    right_length = int(right["end_page"]) - int(right["start_page"]) + 1
    return intersection / (left_length + right_length - intersection)


def greedy_typed_iou_matches(
    predictions: list[dict[str, Any]],
    truth: list[dict[str, Any]],
) -> int:
    candidates: list[tuple[float, int, int]] = []
    for pred_index, prediction in enumerate(predictions):
        for truth_index, reference in enumerate(truth):
            if prediction["type"] != reference["type"]:
                continue
            score = segment_iou(prediction, reference)
            if score >= TYPED_IOU_THRESHOLD:
                candidates.append((score, pred_index, truth_index))

    candidates.sort(reverse=True)
    used_predictions: set[int] = set()
    used_truth: set[int] = set()
    matches = 0
    for _, pred_index, truth_index in candidates:
        if pred_index in used_predictions or truth_index in used_truth:
            continue
        used_predictions.add(pred_index)
        used_truth.add(truth_index)
        matches += 1
    return matches


def internal_boundaries(segments: list[dict[str, Any]], page_count: int) -> list[int]:
    boundaries = {
        int(segment["start_page"])
        for segment in segments
        if 1 < int(segment["start_page"]) <= page_count
    }
    return sorted(boundaries)


def boundary_matches(predicted: list[int], truth: list[int]) -> int:
    used_truth: set[int] = set()
    matches = 0
    for boundary in predicted:
        best_index: int | None = None
        best_distance = BOUNDARY_TOLERANCE + 1
        for truth_index, truth_boundary in enumerate(truth):
            if truth_index in used_truth:
                continue
            distance = abs(boundary - truth_boundary)
            if distance <= BOUNDARY_TOLERANCE and distance < best_distance:
                best_index = truth_index
                best_distance = distance
        if best_index is None:
            continue
        used_truth.add(best_index)
        matches += 1
    return matches


def score_segments(
    predictions: list[dict[str, Any]],
    truth: list[dict[str, Any]],
    page_count: int,
) -> pd.DataFrame:
    pred_labels = page_labels(predictions, page_count)
    truth_labels = page_labels(truth, page_count)
    correct_pages = sum(
        1
        for page in range(1, page_count + 1)
        if pred_labels[page] == truth_labels[page]
    )
    predicted_boundaries = internal_boundaries(predictions, page_count)
    truth_boundaries = internal_boundaries(truth, page_count)
    boundary_match_count = boundary_matches(predicted_boundaries, truth_boundaries)
    typed_iou_match_count = greedy_typed_iou_matches(predictions, truth)
    pred_count = len(predictions)
    truth_count = len(truth)
    oversegmentation_count = max(0, pred_count - truth_count)
    undersegmentation_count = max(0, truth_count - pred_count)

    return pd.DataFrame(
        [
            {
                "view": "prediction",
                "page_level_accuracy": round(correct_pages / max(1, page_count), 3),
                "typed_iou_f1": round(
                    f1(typed_iou_match_count, pred_count, truth_count),
                    3,
                ),
                "boundary_f1": round(
                    f1(
                        boundary_match_count,
                        len(predicted_boundaries),
                        len(truth_boundaries),
                    ),
                    3,
                ),
                "oversegmentation": round(
                    oversegmentation_count / max(1, truth_count),
                    3,
                ),
                "undersegmentation": round(
                    1 - (undersegmentation_count / max(1, truth_count)),
                    3,
                ),
            }
        ]
    )


def build_chart_df(
    views_spec: list[tuple[str, str, list[dict[str, Any]]]],
) -> pd.DataFrame:
    rows: list[dict[str, Any]] = []
    for view_id, label, segments in views_spec:
        for segment in segments:
            type_name = segment["type"]
            start = int(segment["start_page"])
            end = int(segment["end_page"])
            rows.append(
                {
                    "view": label,
                    "view_id": view_id,
                    "type": type_name,
                    "start_page": start,
                    "end_page": end,
                    "x": start - 0.5,
                    "x2": end + 0.5,
                    "n_pages": end - start + 1,
                }
            )
    return pd.DataFrame(rows)


def dataframe_kwargs() -> dict[str, Any]:
    return {"hide_index": True, "use_container_width": True}


def color_domain_for(
    subdocuments: list[dict[str, Any]],
    observed_types: set[str],
) -> list[str]:
    configured_types = [subdocument["name"] for subdocument in subdocuments]
    configured_type_set = set(configured_types)
    unknown_types = sorted(observed_types - configured_type_set)
    return configured_types + unknown_types


def color_range_for(domain_size: int) -> list[str]:
    if domain_size <= len(POLITAX_COLOR_RANGE):
        return POLITAX_COLOR_RANGE[:domain_size]
    repeats = (domain_size // len(POLITAX_COLOR_RANGE)) + 1
    return (POLITAX_COLOR_RANGE * repeats)[:domain_size]


def sorted_model_names(models: set[str]) -> list[str]:
    known = [model for model in MODEL_ORDER if model in models]
    unknown = sorted(models - set(MODEL_ORDER))
    return known + unknown


def main() -> None:
    st.set_page_config(
        page_title="Ground truth vs Retab",
        layout="wide",
    )
    st.markdown(
        """
        <style>
        [data-testid="stMetricLabel"] {
            font-size: 0.8rem;
        }
        [data-testid="stMetricValue"] {
            font-size: 1.05rem;
            line-height: 1.25;
        }
        </style>
        """,
        unsafe_allow_html=True,
    )

    inputs = load_inputs()
    descriptions = {
        subdocument["name"]: subdocument.get("description", "")
        for subdocument in inputs["subdocuments"]
    }
    runs = result_runs()
    if not runs:
        st.error("No result JSONs found under results/.")
        st.stop()

    run_name = runs[0]
    index = load_result_index(run_name)

    with st.sidebar:
        st.header("Model")
        model_names = sorted_model_names({model for _, model in index})
        if not model_names:
            st.warning(f"No split result JSONs found for {run_name}.")
            st.stop()
        model = st.selectbox("Model", model_names, label_visibility="collapsed")

        st.header("Document")
        documents = [
            doc
            for doc in inputs["documents"]
            if (doc["document"], model) in index
        ]
        if not documents:
            st.warning(f"No documents found for model={model}.")
            st.stop()
        document_name = st.selectbox(
            "Document",
            [doc["document"] for doc in documents],
            index=0,
            label_visibility="collapsed",
        )

        input_doc = next(doc for doc in documents if doc["document"] == document_name)
        split = load_json(index[(document_name, model)])
        predictions = segments_from_output(split.get("output", []))
        truth = input_doc["ground_truth"]

    st.title("Ground truth vs Retab")

    page_count = int(input_doc["page_count"])
    cols = st.columns(3)
    cols[0].metric("Document", document_name)
    cols[1].metric("Pages", page_count)
    cols[2].metric("Model", model)

    scores_df = score_segments(predictions, truth, page_count)
    st.markdown("### Scores")
    st.dataframe(scores_df, **dataframe_kwargs())

    views_spec = [
        ("ground_truth", "ground truth", truth),
        ("prediction", model, predictions),
    ]
    chart_df = build_chart_df(views_spec)

    if chart_df.empty:
        st.warning("No segments to render with the current filter.")
        return

    n_pages_for_xaxis = max(chart_df["end_page"].max(), page_count)
    types_in_chart = sorted(chart_df["type"].unique())
    color_domain = color_domain_for(inputs["subdocuments"], set(types_in_chart))
    color_scale = alt.Scale(
        domain=color_domain,
        range=color_range_for(len(color_domain)),
    )
    view_sort_order = [label for _, label, _ in views_spec]
    bar_h = 38
    chart_h = max(120, len(views_spec) * (bar_h + 8) + 60)

    st.markdown("### Segment diagram")

    chart = (
        alt.Chart(chart_df)
        .mark_bar(stroke="black", strokeWidth=0.4, height=bar_h)
        .encode(
            y=alt.Y(
                "view:N",
                sort=view_sort_order,
                title=None,
                axis=alt.Axis(labelLimit=300),
            ),
            x=alt.X(
                "x:Q",
                title="page",
                scale=alt.Scale(domain=[0.5, n_pages_for_xaxis + 0.5]),
                axis=alt.Axis(format="d"),
            ),
            x2="x2:Q",
            color=alt.Color(
                "type:N",
                scale=color_scale,
                legend=alt.Legend(
                    title="type",
                    columns=2,
                    values=types_in_chart,
                    symbolStrokeWidth=0.7,
                ),
            ),
            tooltip=[
                alt.Tooltip("view:N", title="view"),
                alt.Tooltip("type:N", title="type"),
                alt.Tooltip("start_page:Q", title="start"),
                alt.Tooltip("end_page:Q", title="end"),
                alt.Tooltip("n_pages:Q", title="pages"),
            ],
        )
        .properties(height=chart_h, width="container")
    )
    st.altair_chart(chart, use_container_width=True)

    with st.expander("Subdocument type descriptions"):
        if types_in_chart:
            st.dataframe(
                pd.DataFrame(
                    [
                        {"type": type_name, "description": descriptions.get(type_name, "")}
                        for type_name in types_in_chart
                    ]
                ),
                **dataframe_kwargs(),
            )
        else:
            st.caption("No types in the current view.")

    with st.expander("Raw split JSON"):
        st.json(split)


if __name__ == "__main__":
    main()