Update app.py

app.py

@@ -1,43 +1,49 @@
-import io, os, re, json
+# Streamlit Invoice Extraction — Hugging Face Donut (no local .pth) + Tesseract tables
+# - Uses a pretrained model from HF Hub (default: naver-clova-ix/donut-base-finetuned-sroie)
+# - Extracts key fields via Donut JSON if available, else regex fallback
+# - Extracts line items via Tesseract word boxes + geometry heuristics
+# - Works on HF Spaces without any custom checkpoints
+
+import os, io, re, json
 from typing import List, Tuple, Dict
-import os, sys
-sys.path.append(os.path.abspath("pdrt"))  # <— add this
-import models as pdrt_models              # <— from your vendored repo
 
 import numpy as np
 import pandas as pd
 from PIL import Image, ImageOps, ImageFilter
 
 import streamlit as st
-import torch
-import torchvision.transforms as T
 
-# --- word detector (Tesseract) ---
+# OCR for word boxes (detection only) + pdf to images
 import pytesseract
 from pytesseract import Output
-
-# --- PDF -> images ---
 from pdf2image import convert_from_bytes
 
-# ---- import the repo's models ----
-# Install via requirements.txt (git+https URL) OR copy repo files into root.
-# The repo defines model classes: Swin_CTC, VED
-import models as pdrt_models  # from dparres/Pretrained-Document-Recognition-Transformers
-
-st.set_page_config(page_title="Invoice OCR (ViT recognizer + Tesseract detector)", layout="wide")
-
-# ========================= UI SIDEBAR =========================
-st.sidebar.header("Model")
-arch = st.sidebar.selectbox("Architecture", ["Swin_CTC", "VED"], index=0)
-ckpt_path = st.sidebar.text_input("Checkpoint path (inside Space)", value="checkpoints/pdrt_weights.pth")
-alphabet = st.sidebar.text_input("Alphabet (ordered classes, exclude CTC blank)", value="0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz-_/.,:;()[]{}#+*&%$@!?\"' ")
-img_h = st.sidebar.number_input("Recognizer input height", 64, 256, 128, 8)
-img_w = st.sidebar.number_input("Recognizer input width", 128, 2048, 512, 16)
-det_lang = st.sidebar.text_input("Tesseract lang(s) for detection only", value="eng")
+# HF Donut (pretrained, downloaded automatically)
+import torch
+from transformers import DonutProcessor, VisionEncoderDecoderModel
+
+st.set_page_config(page_title="Invoice Extraction — Donut (HF) + Tesseract tables", layout="wide")
+
+# ----------------------------- Sidebar -----------------------------
+st.sidebar.header("Model (Hugging Face)")
+model_id = st.sidebar.text_input(
+    "HF model id",
+    value="naver-clova-ix/donut-base-finetuned-sroie",  # good default for receipts/invoices (SROIE)
+    help="Examples: naver-clova-ix/donut-base-finetuned-sroie, naver-clova-ix/donut-base-finetuned-docvqa"
+)
+task_prompt = st.sidebar.text_input(
+    "Task prompt (for Donut models expecting prompts)",
+    value="<s_cord-v2>",  # SROIE/cord-style models typically ignore or use default; harmless to keep
+    help="Some Donut checkpoints use task-specific prompts; keep or adjust as needed."
+)
+det_lang = st.sidebar.text_input("Tesseract language(s) — detection only", value="eng")
 show_boxes = st.sidebar.checkbox("Show word boxes", value=False)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# ========================= UTILITIES =========================
+st.sidebar.markdown("---")
+st.sidebar.caption("Tip: If your model outputs JSON (e.g., SROIE), we’ll parse it for key fields. Otherwise we’ll regex from generated text.")
+
+# ----------------------------- Utilities -----------------------------
 def load_pages(file_bytes: bytes, name: str) -> List[Image.Image]:
     name = (name or "").lower()
     if name.endswith(".pdf"):
@@ -50,94 +56,39 @@ def preprocess_for_detection(img: Image.Image) -> Image.Image:
     g = g.filter(ImageFilter.UnsharpMask(radius=1, percent=150, threshold=3))
     return g
 
-@st.cache_resource
-def load_pdrt(arch_name: str, ckpt: str, num_classes: int):
-    if arch_name == "Swin_CTC":
-        model = pdrt_models.Swin_CTC(num_classes=num_classes)
-    elif arch_name == "VED":
-        model = pdrt_models.VED(num_classes=num_classes)
-    else:
-        raise ValueError("Unknown model")
-    state = torch.load(ckpt, map_location="cpu")
-    model.load_state_dict(state, strict=False)
-    model.eval().to(device)
-    return model
-
-def build_transform(img_h: int, img_w: int):
-    return T.Compose([
-        T.Grayscale(num_output_channels=3),  # keep 3ch if encoder expects RGB
-        T.Resize((img_h, img_w)),
-        T.ToTensor(),
-        T.Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5]),
-    ])
-
-def greedy_ctc_decode(logits: torch.Tensor, alphabet: str) -> str:
-    """
-    logits: (B, T, C) or (T, B, C). We map argmax to chars, collapse repeats, remove blank.
-    We assume blank_id = len(alphabet).
-    """
-    if logits.dim() == 3 and logits.shape[0] != 1 and logits.shape[1] == 1:
-        # rare shape, just permute if needed
-        pass
-    if logits.shape[0] == 1:
-        logits = logits.squeeze(0)  # (T, C)
-    elif logits.shape[1] == 1:
-        logits = logits[:,0,:]      # (T, C)
-    probs = logits.softmax(-1)
-    ids = probs.argmax(-1).tolist()
-    blank_id = len(alphabet)
-    out = []
-    prev = None
-    for i in ids:
-        if i != prev and i != blank_id:
-            out.append(alphabet[i] if i < len(alphabet) else "")
-        prev = i
-    return "".join(out)
-
-def recognize_word_crops(model, crops: List[Image.Image], tfm, arch_name: str, alphabet: str) -> List[str]:
-    texts = []
+@st.cache_resource(show_spinner=True)
+def load_donut(_model_id: str):
+    processor = DonutProcessor.from_pretrained(_model_id)
+    model = VisionEncoderDecoderModel.from_pretrained(_model_id)
+    model.to(device)
+    model.eval()
+    return processor, model
+
+def donut_infer(img: Image.Image, processor: DonutProcessor, model: VisionEncoderDecoderModel, prompt: str):
+    # Donut expects RGB PIL Image; processor handles resizing/normalization
+    inputs = processor(images=img, text=prompt, return_tensors="pt").to(device)
     with torch.no_grad():
-        for im in crops:
-            x = tfm(im).unsqueeze(0).to(device)
-            y = model(x)
-            if arch_name == "Swin_CTC":
-                # expect CTC logits [B, T, C] or [T, B, C]
-                if y.dim() == 3 and y.shape[0] == 1:   # [1, T, C]
-                    logits = y[0]                      # [T, C]
-                elif y.dim() == 3 and y.shape[1] == 1: # [T, 1, C]
-                    logits = y[:,0,:]
-                else:
-                    logits = y
-                txt = greedy_ctc_decode(logits, alphabet)
-            else:
-                # VED: if returns token ids/logits, plug your repo's decoding here.
-                # Fallback: argmax over last dim per step and map ids to alphabet (no blank).
-                if y.dim() == 3 and y.shape[0] == 1:
-                    y = y[0]
-                ids = y.argmax(-1).tolist()
-                txt = "".join(alphabet[i] if i < len(alphabet) else "" for i in ids).strip()
-            texts.append(txt)
-    return texts
-
-def detect_words(img: Image.Image, lang="eng") -> pd.DataFrame:
-    df = pytesseract.image_to_data(img, lang=lang, output_type=Output.DATAFRAME)
-    df = df.dropna(subset=["text"]).reset_index(drop=True)
-    df["x2"] = df["left"] + df["width"]
-    df["y2"] = df["top"] + df["height"]
-    return df[df["conf"] > -1]
-
-def crop_words(img: Image.Image, df: pd.DataFrame) -> List[Tuple[Image.Image, Dict]]:
-    crops, metas = [], []
-    for _, r in df.iterrows():
-        if str(r["text"]).strip() == "":
-            continue
-        box = (int(r["left"]), int(r["top"]), int(r["x2"]), int(r["y2"]))
-        c = img.crop(box)
-        crops.append(c)
-        metas.append({"box": box})
-    return crops, metas
-
-# ---------------- key fields & table (same logic as earlier Tesseract app) ----------------
+        outputs = model.generate(
+            **inputs,
+            max_length=1024,
+            num_beams=1,
+            early_stopping=True,
+        )
+    # decode
+    seq = processor.batch_decode(outputs, skip_special_tokens=True)[0]
+    # Donut models often emit JSON; try to parse
+    parsed = None
+    try:
+        # strip whitespace garbage around JSON
+        start = seq.find("{")
+        end = seq.rfind("}")
+        if start != -1 and end != -1 and end > start:
+            parsed = json.loads(seq[start:end+1])
+    except Exception:
+        parsed = None
+    return seq, parsed
+
+# ----------------------------- Key fields & line items -----------------------------
 CURRENCY = r"(?P<curr>USD|CAD|EUR|GBP|\$|C\$|€|£)?"
 MONEY = rf"{CURRENCY}\s?(?P<amt>\d{{1,3}}(?:[,]\d{{3}})*(?:[.]\d{{2}})?)"
 DATE = r"(?P<date>(?:\d{4}[-/]\d{1,2}[-/]\d{1,2})|(?:\d{1,2}[-/]\d{1,2}[-/]\d{2,4})|(?:[A-Za-z]{3,9}\s+\d{1,2},\s*\d{2,4}))"
@@ -147,7 +98,7 @@ TOTAL_PAT = rf"(?:\b(total(?:\s*amount)?|amount\s*due|grand\s*total)\b.*?{MONEY}
 SUBTOTAL_PAT = rf"(?:\bsub\s*total\b.*?{MONEY})"
 TAX_PAT = rf"(?:\b(tax|gst|vat|hst)\b.*?{MONEY})"
 
-def parse_fields(fulltext: str):
+def parse_fields_regex(fulltext: str):
     t = re.sub(r"[ \t]+", " ", fulltext)
     t = re.sub(r"\n{2,}", "\n", t)
     out = {"invoice_number":None,"invoice_date":None,"po_number":None,"subtotal":None,"tax":None,"total":None,"currency":None}
@@ -155,22 +106,87 @@ def parse_fields(fulltext: str):
     m = re.search(PO_PAT, t, re.I);            out["po_number"] = m.group("po") if m else None
     m = re.search(rf"(invoice\s*date[:\-\s]*){DATE}", t, re.I)
     out["invoice_date"] = (m.group("date") if m else (re.search(DATE, t, re.I).group("date") if re.search(DATE, t, re.I) else None))
-    m = re.search(SUBTOTAL_PAT, t, re.I|re.S); 
+    m = re.search(SUBTOTAL_PAT, t, re.I|re.S);
     if m: out["subtotal"], out["currency"] = m.group("amt").replace(",",""), m.group("curr") or out["currency"]
     m = re.search(TAX_PAT, t, re.I|re.S);
     if m: out["tax"], out["currency"] = m.group("amt").replace(",",""), m.group("curr") or out["currency"]
     m = re.search(TOTAL_PAT, t, re.I|re.S);
-    if m: out["total"], out["currency"] = m.group("amt").replace(",",""), m.group("curr") or out["currency"]
+    if m:
+        out["total"], out["currency"] = m.group("amt").replace(",", ""), m.group("curr") or out["currency"]
     if out["currency"] in ["$", "C$", "€", "£"]:
         out["currency"] = {"$":"USD", "C$":"CAD", "€":"EUR", "£":"GBP"}[out["currency"]]
     return out
 
+def normalize_kv_from_donut(parsed: dict):
+    """Try to map common Donut outputs to our schema."""
+    txt = json.dumps(parsed).lower()
+    # heuristic mapping for typical SROIE/receipt keys
+    candidates = {
+        "invoice_number": ["invoice_number","invoice no","invoice_no","invoice","inv_no"],
+        "invoice_date":   ["date","invoice_date","bill_date"],
+        "po_number":      ["po_number","po","purchase_order"],
+        "subtotal":       ["subtotal","sub_total"],
+        "tax":            ["tax","gst","vat","hst"],
+        "total":          ["total","amount_total","amount_due","grand_total"]
+    }
+    out = {k: None for k in ["invoice_number","invoice_date","po_number","subtotal","tax","total","currency"]}
+    # simple search: pick first occurrence
+    def search_keys(obj, key_list):
+        # breadth-first scan
+        if isinstance(obj, dict):
+            for k, v in obj.items():
+                if any(kk in k.lower() for kk in key_list):
+                    return v
+                found = search_keys(v, key_list)
+                if found is not None:
+                    return found
+        elif isinstance(obj, list):
+            for it in obj:
+                found = search_keys(it, key_list)
+                if found is not None:
+                    return found
+        return None
+
+    for outk, key_list in candidates.items():
+        val = search_keys(parsed, key_list)
+        if isinstance(val, (dict, list)):
+            val = None  # keep it simple; Donut sometimes nests values
+        if isinstance(val, str):
+            out[outk] = val.strip()
+    # currency guess:
+    curr = re.search(r"(USD|CAD|EUR|GBP|\$|C\$|€|£)", json.dumps(parsed, ensure_ascii=False), re.I)
+    if curr:
+        sym = curr.group(1)
+        out["currency"] = {"$":"USD","C$":"CAD","€":"EUR","£":"GBP"}.get(sym, sym.upper())
+    return out
+
+def detect_words(img: Image.Image, lang="eng") -> pd.DataFrame:
+    df = pytesseract.image_to_data(img, lang=lang, output_type=Output.DATAFRAME)
+    df = df.dropna(subset=["text"]).reset_index(drop=True)
+    df["x2"] = df["left"] + df["width"]
+    df["y2"] = df["top"] + df["height"]
+    return df[df["conf"] > -1]
+
+def crop_words(img: Image.Image, df: pd.DataFrame) -> List[Tuple[Image.Image, Dict]]:
+    crops, metas = [], []
+    for _, r in df.iterrows():
+        if str(r["text"]).strip() == "":
+            continue
+        box = (int(r["left"]), int(r["top"]), int(r["x2"]), int(r["y2"]))
+        c = img.crop(box)
+        crops.append(c)
+        metas.append({"box": box})
+    return crops, metas
+
 HEAD_CANDIDATES = ["description","item","qty","quantity","price","unit","rate","amount","total"]
 def items_from_wordgrid(df: pd.DataFrame) -> pd.DataFrame:
-    # Group into lines
+    if df.empty:
+        return pd.DataFrame()
     df = df.copy()
     df["cx"] = df["left"] + 0.5*df["width"]
     df["cy"] = df["top"]  + 0.5*df["height"]
+
+    # group lines
     lines = []
     for (b,p,l), g in df.groupby(["block_num","par_num","line_num"]):
         text = " ".join([t for t in g["text"].astype(str) if t.strip()])
@@ -180,7 +196,7 @@ def items_from_wordgrid(df: pd.DataFrame) -> pd.DataFrame:
                 "text": text.lower(),
                 "top": g["top"].min(), "bottom": (g["top"]+g["height"]).max(),
                 "left": g["left"].min(), "right": (g["left"]+g["width"]).max(),
-                "words": g.sort_values("cx")[["cx","left","top","width","height"]].values.tolist()
+                "words": g.sort_values("left")[["left","top","width","height","text"]].values.tolist()
             })
     L = pd.DataFrame(lines)
     if L.empty: return pd.DataFrame()
@@ -190,36 +206,80 @@ def items_from_wordgrid(df: pd.DataFrame) -> pd.DataFrame:
     H = headers.iloc[0]
     header_y = H["bottom"] + 4
 
-    # choose column centers from header words positions
-    # we reuse df within header band
-    header_band = df[(df["top"]>=H["top"]-5) & ((df["top"]+df["height"])<=H["bottom"]+5)]
-    header_band = header_band.sort_values("left")
-    col_x = header_band["left"].tolist()
-    if len(col_x)<2: return pd.DataFrame()
-    # region below header until totals
-    below = df[df["top"]>header_y].copy()
+    # derive column anchors from header words positions
+    df_header = detect_words(img=None, lang="eng")  # placeholder to keep signature consistent
+
+    # get header band words
+    # reconstruct header band from original DF
+    # (we need original df back here; easier: pass it in as closure var)
+    # → we'll adapt: compute from global last_df if present
+    return_df = pd.DataFrame()
+    return return_df
+
+# We’ll implement a simpler, robust table extractor to avoid closure complexity:
+def items_from_words_simple(tsv: pd.DataFrame) -> pd.DataFrame:
+    # find header line
+    L = []
+    for (b,p,l), g in tsv.groupby(["block_num","par_num","line_num"]):
+        text = " ".join([w for w in g["text"].astype(str).tolist() if w.strip()])
+        if text.strip():
+            L.append({
+                "block_num": b, "par_num": p, "line_num": l,
+                "text": text.lower(),
+                "top": g["top"].min(), "bottom": (g["top"]+g["height"]).max(),
+                "left": g["left"].min(), "right": (g["left"]+g["width"]).max()
+            })
+    lines = pd.DataFrame(L)
+    if lines.empty:
+        return pd.DataFrame()
+
+    def score_header(s: str):
+        return sum(1 for h in HEAD_CANDIDATES if h in s)
+
+    lines["header_score"] = lines["text"].apply(score_header)
+    hdrs = lines[lines["header_score"] >= 2].sort_values(["header_score","top"], ascending=[False,True])
+    if hdrs.empty:
+        return pd.DataFrame()
+    H = hdrs.iloc[0]
+    header_top, header_bottom = H["top"], H["bottom"]
+
+    # header words
+    header_words = tsv[(tsv["top"] >= header_top - 5) & ((tsv["top"] + tsv["height"]) <= header_bottom + 5)]
+    header_words = header_words.sort_values("left")
+    if header_words.empty:
+        return pd.DataFrame()
+    xs = header_words["left"].tolist()
+
+    # items region
+    below = tsv[tsv["top"] > header_bottom + 5].copy()
     totals_mask = below["text"].str.lower().str.contains(r"(sub\s*total|amount\s*due|total|grand\s*total|balance)", regex=True, na=False)
     if totals_mask.any():
-        stop_y = below.loc[totals_mask,"top"].min()
-        below = below[below["top"]<stop_y-4]
+        stop_y = below.loc[totals_mask, "top"].min()
+        below = below[below["top"] < stop_y - 4]
+    if below.empty:
+        return pd.DataFrame()
+
+    # build rows by assigning words to nearest header x
     rows = []
     for (b,p,l), g in below.groupby(["block_num","par_num","line_num"]):
-        if g["text"].astype(str).str.strip().eq("").all(): continue
         g = g.sort_values("left")
-        # assign to nearest header word x
-        xs = np.array(col_x)
         buckets = {i:[] for i in range(len(xs))}
-        for _,w in g.iterrows():
-            idx = int(np.abs(xs - w["left"]).argmin())
+        for _, w in g.iterrows():
+            if not str(w["text"]).strip():
+                continue
+            idx = int(np.abs(np.array(xs) - w["left"]).argmin())
             buckets[idx].append(str(w["text"]))
-        vals = [" ".join(buckets.get(i,[])).strip() for i in range(len(xs))]
+        vals = [" ".join(buckets[i]).strip() for i in range(len(xs))]
         rows.append(vals)
-    if not rows: return pd.DataFrame()
+    if not rows:
+        return pd.DataFrame()
+
     df_rows = pd.DataFrame(rows).fillna("")
-    # try to name columns
+    # name columns heuristically
     names = []
-    for i, w in enumerate(header_band["text"].tolist()[:df_rows.shape[1]]):
-        wl = w.lower()
+    hdr_tokens = [t.lower() for t in header_words["text"].tolist()]
+    for i in range(df_rows.shape[1]):
+        wl = hdr_tokens[i] if i < len(hdr_tokens) else f"col_{i}"
         if "desc" in wl or wl in ["item","description"]:
             names.append("description")
         elif wl in ["qty","quantity"]:
@@ -235,62 +295,49 @@ def items_from_wordgrid(df: pd.DataFrame) -> pd.DataFrame:
     df_rows = df_rows[~(df_rows.fillna("").apply(lambda r: "".join(r.values), axis=1).str.strip()=="")]
     return df_rows.reset_index(drop=True)
 
-# ========================= APP =========================
-st.title("Invoice Extraction — ViT recognizer (dparres) + Tesseract detector")
+# ----------------------------- App -----------------------------
+st.title("Invoice Extraction — Donut (HF pretrained) + Tesseract tables")
 
 up = st.file_uploader("Upload an invoice (PDF/JPG/PNG)", type=["pdf","png","jpg","jpeg"])
 if not up:
     st.info("Upload a scanned invoice to begin.")
     st.stop()
 
-pages = load_pages(up.read(), up.name)
-
 # load model once
-num_classes = len(alphabet) + (1 if arch=="Swin_CTC" else 0)  # add CTC blank for Swin_CTC
-assert os.path.exists(ckpt_path), f"Checkpoint not found: {ckpt_path}"
-model = load_pdrt(arch, ckpt_path, num_classes)
-tfm = build_transform(img_h, img_w)
+with st.spinner(f"Loading model '{model_id}' from Hugging Face…"):
+    processor, donut_model = load_donut(model_id)
+
+pages = load_pages(up.read(), up.name)
 
 page_idx = 0
 if len(pages) > 1:
     page_idx = st.number_input("Page", 1, len(pages), 1) - 1
 img = pages[page_idx]
 
-col1, col2 = st.columns([1.1,1.3], gap="large")
+col1, col2 = st.columns([1.1, 1.3], gap="large")
 
 with col1:
     st.subheader("Preview")
     st.image(img, use_column_width=True)
     det_img = preprocess_for_detection(img)
-    with st.expander("Detection view"):
+    with st.expander("Detection view (preprocessed for boxes)"):
         st.image(det_img, use_column_width=True)
 
 with col2:
     st.subheader("OCR & Extraction")
-    # 1) detect words (boxes only)
-    det_df = detect_words(det_img, lang=det_lang)
 
-    # 2) crop & recognize each word via ViT recognizer
-    crops, metas = crop_words(det_img, det_df)
-    texts = recognize_word_crops(model, crops, tfm, arch, alphabet)
+    # 1) Donut extraction (key fields or full text)
+    with st.spinner("Running Donut…"):
+        seq, parsed = donut_infer(img, processor, donut_model, task_prompt)
 
-    # 3) stitch line-by-line using tesseract line indices
-    det_df = det_df.reset_index(drop=True)
-    det_df["pred"] = texts
-    grouped = det_df.groupby(["block_num","par_num","line_num"])
-    lines = []
-    for _, g in grouped:
-        g = g.sort_values("left")
-        line = " ".join([t for t in g["pred"].tolist() if t])
-        lines.append(line)
-    full_text = "\n".join([ln for ln in lines if ln.strip()])
-
-    if show_boxes:
-        st.caption("First 15 predicted words")
-        st.write(det_df[["left","top","width","height","text","pred"]].head(15))
+    # 2) Key fields
+    if parsed:
+        key_fields = normalize_kv_from_donut(parsed)
+        donut_payload = parsed
+    else:
+        key_fields = parse_fields_regex(seq)
+        donut_payload = {"generated_text": seq}
 
-    # 4) key fields
-    key_fields = parse_fields(full_text)
     k1,k2,k3 = st.columns(3)
     with k1:
         st.write(f"**Invoice #:** {key_fields.get('invoice_number') or '—'}")
@@ -304,21 +351,34 @@ with col2:
         cur = key_fields.get('currency') or ''
         st.write(f"**Total:** {tot} {cur}".strip())
 
-    # 5) line items (geometry heuristic)
-    items = items_from_wordgrid(det_df.assign(text=det_df["pred"]))
+    # 3) Tesseract line items (geometry heuristic)
+    with st.spinner("Detecting words with Tesseract (for table)…"):
+        tsv = pytesseract.image_to_data(det_img, lang=det_lang, output_type=Output.DATAFRAME)
+        tsv = tsv.dropna(subset=["text"]).reset_index(drop=True)
+        tsv["x2"] = tsv["left"] + tsv["width"]
+        tsv["y2"] = tsv["top"] + tsv["height"]
+
     st.markdown("**Line Items**")
+    items = items_from_words_simple(tsv)
     if items.empty:
         st.caption("No line items confidently detected.")
     else:
         st.dataframe(items, use_container_width=True)
 
-    # 6) downloads
+    # 4) Downloads
     result = {
-        "file": up.name, "page": page_idx+1,
+        "file": up.name,
+        "page": page_idx + 1,
         "key_fields": key_fields,
         "items": items.to_dict(orient="records") if not items.empty else [],
-        "full_text": full_text
+        "donut_raw": donut_payload,
     }
-    st.download_button("Download JSON", data=json.dumps(result, indent=2), file_name="invoice_extraction.json", mime="application/json")
+    st.download_button("Download JSON", data=json.dumps(result, indent=2),
+                       file_name="invoice_extraction.json", mime="application/json")
     if not items.empty:
-        st.download_button("Download Items CSV", data=items.to_csv(index=False), file_name="invoice_items.csv", mime="text/csv")
+        st.download_button("Download Items CSV", data=items.to_csv(index=False),
+                           file_name="invoice_items.csv", mime="text/csv")
+
+    if show_boxes:
+        st.caption("First 20 Tesseract word boxes")
+        st.dataframe(tsv[["left","top","width","height","text","conf"]].head(20), use_container_width=True)