app.py

"""form-field-v1 detector demo — nano/small/medium run as verified ONNX via ONNX Runtime. Field STATE
(filled/empty, checked/unchecked, signed/blank) comes from the medium 6-class head when available, and an
ink-presence check that rescues filled fields the head misses (and supplies state for the 3-class models).
ZeroGPU (CUDA) with CPU fallback."""
import os, json
import numpy as np
import cv2
from PIL import Image, ImageDraw, ImageFont
from huggingface_hub import hf_hub_download
import onnxruntime as ort
import gradio as gr
try:
    import spaces
except Exception:                      # local dev without the spaces shim
    class _S:
        def GPU(self, *a, **k):
            def deco(f): return f
            return deco
    spaces = _S()

TOKEN = os.environ.get("HF_TOKEN")
CLASSES = ["Text", "ChoiceButton", "Signature"]
COLORS = {"Text": (29, 95, 168), "ChoiceButton": (22, 121, 79), "Signature": (176, 84, 20)}
# medium is a 6-class coarse3_state head: raw index -> (coarse type, state). nano/small are 3-class (no state).
RF_STATE = {0: ("ChoiceButton", "checked"), 1: ("ChoiceButton", "unchecked"),
            2: ("Signature", "blank"),       3: ("Signature", "signed"),
            4: ("Text", "empty"),            5: ("Text", "filled")}
FILLED_OF = {"Text": "filled", "ChoiceButton": "checked", "Signature": "signed"}
EMPTY_OF = {"Text": "empty", "ChoiceButton": "unchecked", "Signature": "blank"}
POS = {"filled", "checked", "signed"}
MODELS = {
    "nano · open · 0.9M":        {"repo": "nutrientdocs/form-field-v1-nano",           "kind": "yolox",  "size": 640},
    "small · commercial · 8.9M": {"repo": "nutrientdocs/form-field-v1-small-private",  "kind": "yolox",  "size": 896},
    "medium · commercial · 34M": {"repo": "nutrientdocs/form-field-v1-medium-private", "kind": "rfdetr", "size": 1216},
}
_SESS = {}
_MEAN = np.array([0.485, 0.456, 0.406], np.float32); _STD = np.array([0.229, 0.224, 0.225], np.float32)


def has_ink(pil_crop):
    """Training-free ink presence: local-adaptive contrast (any polarity -> white/colored/dark/gradient fills) +
    morphological removal of the box frame / underline (so a mark touching the border survives) + component filter."""
    a = np.asarray(pil_crop.convert("RGB"), np.uint8); H, W = a.shape[:2]
    if H < 6 or W < 6: return False
    gray = cv2.cvtColor(a, cv2.COLOR_RGB2GRAY).astype(np.int16)
    blk = max(9, (min(H, W)) | 1)
    mean = cv2.blur(gray.astype(np.float32), (blk, blk))
    mask = (np.abs(gray - mean) > 25).astype(np.uint8)
    if mask.sum() == 0: return False
    hk = cv2.getStructuringElement(cv2.MORPH_RECT, (max(8, int(0.55 * W)), 1))
    vk = cv2.getStructuringElement(cv2.MORPH_RECT, (1, max(8, int(0.55 * H))))
    lines = cv2.morphologyEx(mask, cv2.MORPH_OPEN, hk) | cv2.morphologyEx(mask, cv2.MORPH_OPEN, vk)
    mask = cv2.bitwise_and(mask, cv2.bitwise_not(lines))
    n, _lab, stats, _c = cv2.connectedComponentsWithStats(mask, connectivity=8)
    content = 0; kept = 0; floor = max(6, int(0.0008 * H * W))
    for i in range(1, n):
        if stats[i, 4] >= floor: content += stats[i, 4]; kept += 1
    return (content / (H * W)) >= 0.004 and kept >= 1


def _cb_marked(pil_crop, inset=0.22, thr=0.06):
    """Checkbox/radio state: measure ink strictly INSIDE the box frame (inset past the border) vs the paper. An
    empty box's interior is blank; a checked box's mark lives in the interior. Excluding the frame by construction
    means an empty box can't be mistaken for checked (frame ink never counts)."""
    a = np.asarray(pil_crop.convert("RGB"), np.uint8); H, W = a.shape[:2]
    if H < 6 or W < 6: return False
    g = cv2.cvtColor(a, cv2.COLOR_RGB2GRAY)
    bg = np.percentile(g, 90)                                  # paper / fill color
    mx = max(1, int(W * inset)); my = max(1, int(H * inset))
    core = g[my:H - my, mx:W - mx]
    if core.size < 4: return False
    return float((core.astype(np.int16) < bg - 50).mean()) >= thr


def _resolve_state(det, pil):
    """Hybrid: trust the head's positive (filled/checked/signed); otherwise let the ink check rescue strays and
    supply state for the 3-class models. One-directional — the heuristic never downgrades a positive.
    Checkboxes/radios use the interior-fill method; text/signature use the frame-removal component method."""
    t, hs = det["type"], det["state"]
    if hs in POS:
        return hs
    x, y, w, h = det["box"]
    crop = pil.crop((int(x), int(y), int(x + w), int(y + h)))
    ink = _cb_marked(crop) if t == "ChoiceButton" else has_ink(crop)
    return FILLED_OF[t] if ink else EMPTY_OF[t]


def _session(name):
    if name not in _SESS:
        cfg = MODELS[name]
        path = hf_hub_download(cfg["repo"], "model.onnx", token=TOKEN)
        provs = [("CUDAExecutionProvider", {"use_tf32": 0}), "CPUExecutionProvider"]
        _SESS[name] = (ort.InferenceSession(path, providers=provs), cfg)
    return _SESS[name]


def _nms(boxes, scores, iou=0.6):
    if not boxes: return []
    b = np.array(boxes); x1, y1, x2, y2 = b[:, 0], b[:, 1], b[:, 0] + b[:, 2], b[:, 1] + b[:, 3]
    a = (x2 - x1) * (y2 - y1); order = np.array(scores).argsort()[::-1]; keep = []
    while order.size:
        i = order[0]; keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]]); yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]]); yy2 = np.minimum(y2[i], y2[order[1:]])
        w = np.maximum(0, xx2 - xx1); h = np.maximum(0, yy2 - yy1); inter = w * h
        ov = inter / (a[i] + a[order[1:]] - inter + 1e-9); order = order[1:][ov <= iou]
    return keep


def _infer(name, pil, thr):
    sess, cfg = _session(name); W, H = pil.size; S = cfg["size"]
    dets = []
    iname = sess.get_inputs()[0].name
    if cfg["kind"] == "rfdetr":
        x = np.asarray(pil.resize((S, S), Image.BILINEAR), np.float32) / 255.0
        x = ((x - _MEAN) / _STD).transpose(2, 0, 1)[None]
        logits, boxes = sess.run(None, {iname: x}); logits, boxes = logits[0], boxes[0]
        prob = 1 / (1 + np.exp(-logits)); lab = prob.argmax(1); sc = prob.max(1)
        for q in range(len(lab)):
            if sc[q] < thr: continue
            typ, state = RF_STATE[int(lab[q])]         # 6-class head: keep the state (filled/empty, checked/…)
            cx, cy, bw, bh = boxes[q]
            dets.append({"type": typ, "state": state, "score": float(sc[q]),
                         "box": [(cx-bw/2)*W, (cy-bh/2)*H, bw*W, bh*H]})
    else:  # yolox: letterbox, decoded output [N,8] (3-class, no state head)
        r = min(S/W, S/H); nw, nh = int(W*r), int(H*r)
        canvas = np.full((S, S, 3), 114, np.float32)
        canvas[:nh, :nw] = np.asarray(pil.resize((nw, nh), Image.BILINEAR), np.float32)
        out = sess.run(None, {iname: canvas.transpose(2, 0, 1)[None]})[0][0]
        obj = out[:, 4]; cls = out[:, 5:8]; clab = cls.argmax(1); score = obj * cls.max(1)
        for c in (0, 1, 2):
            idx = [i for i in range(len(clab)) if clab[i] == c and score[i] >= thr]
            kb = [[float((out[i, 0]-out[i, 2]/2)/r), float((out[i, 1]-out[i, 3]/2)/r),
                   float(out[i, 2]/r), float(out[i, 3]/r)] for i in idx]
            ks = [float(score[i]) for i in idx]
            for j in _nms(kb, ks):
                dets.append({"type": CLASSES[c], "state": None, "score": ks[j], "box": kb[j]})
    return dets


@spaces.GPU(duration=60)
def detect(image, model_name, threshold):
    if image is None:
        return None, "Upload a form page to detect its fields."
    pil = image.convert("RGB")
    try:
        dets = _infer(model_name, pil, float(threshold))
    except Exception as e:
        return None, f"⚠️ Model unavailable — this Space needs access to the model weights. ({str(e)[:120]})"
    out = pil.copy(); d = ImageDraw.Draw(out)
    try: font = ImageFont.load_default()
    except Exception: font = None
    rows = []
    for det in sorted(dets, key=lambda t: -t["score"]):
        cls, sc = det["type"], det["score"]; x, y, w, h = det["box"]
        state = _resolve_state(det, pil)               # hybrid: 6-class head + ink-presence stray-catch
        col = COLORS[cls]; d.rectangle([x, y, x+w, y+h], outline=col, width=3)
        d.text((x+2, max(0, y-11)), f"{cls}-{state} {sc:.2f}", fill=col, font=font)
        rows.append({"box": [round(x), round(y), round(w), round(h)], "type": cls,
                     "state": state, "score": round(sc, 3)})
    return out, json.dumps(rows, indent=2)


INTRO = """# form-field-v1 · detector demo
Detect **Text**, **Choice** (checkbox/radio) and **Signature** fields on empty, filled, and handwritten form pages —
and read each field's **state** (filled/empty, checked/unchecked, signed/blank). State comes from the model on
`medium`, plus an ink check that catches filled fields the model misses. Just upload a page.

- 🏆 [Leaderboard](https://huggingface.co/spaces/nutrientdocs/form-field-v1-leaderboard) ·
  📊 [Benchmark](https://huggingface.co/datasets/nutrientdocs/form-field-v1-benchmark) ·
  Models: [nano](https://huggingface.co/nutrientdocs/form-field-v1-nano) (open) ·
  [small](https://huggingface.co/nutrientdocs/form-field-v1-small) ·
  [medium](https://huggingface.co/nutrientdocs/form-field-v1-medium) (commercial)
"""
ABOUT = """## About the author
<a href="https://nutrient.io/"><img src="https://avatars2.githubusercontent.com/u/1527679?v=3&s=200" height="80" /></a>

This project is maintained and funded by [Nutrient](https://nutrient.io/) - The deterministic document infrastructure enterprises run their highest-stakes workflows on: replayable output, clear exceptions, and full audit trails on the messy, regulated documents where AI alone breaks.
"""

with gr.Blocks(title="form-field-v1 detector") as demo:
    gr.Markdown(INTRO)
    with gr.Row():
        with gr.Column():
            inp = gr.Image(type="pil", label="Form page")
            model = gr.Dropdown(list(MODELS), value="medium · commercial · 34M", label="Model")
            thr = gr.Slider(0.05, 0.9, value=0.35, step=0.05, label="Confidence threshold")
            btn = gr.Button("Detect fields", variant="primary")
        with gr.Column():
            outimg = gr.Image(type="pil", label="Detections")
            outjson = gr.Code(language="json", label="Fields", lines=8, max_lines=24)
    ex = [[os.path.join("examples", x)] for x in sorted(os.listdir("examples"))] if os.path.isdir("examples") else []
    if ex: gr.Examples(ex, inputs=[inp], run_on_click=False)
    btn.click(detect, [inp, model, thr], [outimg, outjson])
    gr.Markdown(ABOUT)

if __name__ == "__main__":
    demo.queue().launch(theme=gr.themes.Soft())