app.py

"""Server-side demo for the commercial Nutrient Form Field VLM and optional FF-DETR localization."""
from __future__ import annotations

import json
import os
from pathlib import Path

import gradio as gr
import spaces
import torch
from PIL import Image, ImageDraw

VLM_REPO = os.environ.get("FORM_FIELD_VLM_REPO", "nutrientdocs/form-field-vlm-private")
DETECTOR_REPO = os.environ.get("FORM_FIELD_DETECTOR_REPO", "nutrientdocs/ffdetr-rfdetr-medium-hf")
TOKEN = os.environ.get("HF_TOKEN")
PROMPT = ("Detect every form field in this page. For each, return a JSON object with these keys: box "
          "(the widget bounding box as [x0,y0,x1,y1] on a 0-1000 grid), type (one of: text, "
          "choice_radio, choice_checkbox, choice_select, signature), label (the field's text label), "
          "group_id (a shared id for radio buttons belonging to the same group; null if the field is not "
          "part of a group). Return a JSON list of these objects.")
SCHEMA = {"type": "array", "items": {"type": "object", "additionalProperties": False,
    "properties": {"box": {"type": "array", "items": {"type": "integer"}, "minItems": 4, "maxItems": 4},
        "type": {"type": "string", "enum": ["text", "choice_radio", "choice_checkbox", "choice_select", "signature"]},
        "label": {"type": "string"}, "group_id": {"type": ["string", "null"]}},
    "required": ["box", "type", "label", "group_id"]}}
COLORS = {"text": "#2563eb", "choice_checkbox": "#16a34a", "choice_radio": "#ea580c",
          "choice_select": "#9333ea", "signature": "#dc2626"}
STATE = {}
APP_DIR = Path(__file__).resolve().parent
EXAMPLES_DIR = APP_DIR / "examples"
if not EXAMPLES_DIR.is_dir():
    EXAMPLES_DIR = APP_DIR.parent / "examples"
EXAMPLES = [
    ("City permit", "city-of-mesquite-front-carport-application.jpg"),
    ("Insurance claim", "great-west-life-notice-of-claim.jpg"),
    ("Police information check", "north-vancouver-rcmp-police-information-check.jpg"),
    ("Employment application", "wyatt-management-employment-application.jpg"),
    ("Agricultural membership", "honiton-agricultural-association-membership-application.jpg"),
    ("Spanish medical information", "spanish-medical-information-form.jpg"),
    ("U.S. direct deposit", "us-direct-deposit-sign-up-form-1199a.jpg"),
    ("Australian aviation application", "australian-aviation-reference-number-application.jpg"),
    ("Employee information", "global-electronics-employee-information-form.png"),
]


def _load():
    """Load once at startup through ZeroGPU's emulated CUDA placement."""
    if STATE:
        return STATE
    try:
        from transformers import AutoImageProcessor, AutoModelForImageTextToText, AutoProcessor
        from transformers import RfDetrForObjectDetection, RfDetrImageProcessor
        kw = {"token": TOKEN} if TOKEN else {}
        print(f"[startup] loading VLM processor from {VLM_REPO}", flush=True)
        STATE["processor"] = AutoProcessor.from_pretrained(VLM_REPO, **kw)
        print(f"[startup] loading VLM weights from {VLM_REPO}", flush=True)
        STATE["vlm"] = AutoModelForImageTextToText.from_pretrained(
            VLM_REPO, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, **kw).eval().to("cuda")
        print(f"[startup] loading detector from {DETECTOR_REPO}", flush=True)
        try:
            STATE["detector_processor"] = AutoImageProcessor.from_pretrained(DETECTOR_REPO, **kw)
        except Exception:
            STATE["detector_processor"] = RfDetrImageProcessor()
        STATE["detector_processor"].size = {"height": 1216, "width": 1216}
        STATE["detector"] = RfDetrForObjectDetection.from_pretrained(DETECTOR_REPO, **kw).eval().to("cuda")
        print("[startup] models ready", flush=True)
    except Exception as exc:
        import traceback
        traceback.print_exc()
        STATE["error"] = exc
    return STATE


def _iou(a, b):
    x0, y0, x1, y1 = max(a[0], b[0]), max(a[1], b[1]), min(a[0]+a[2], b[0]+b[2]), min(a[1]+a[3], b[1]+b[3])
    inter = max(0, x1-x0) * max(0, y1-y0); union = a[2]*a[3] + b[2]*b[3] - inter
    return inter / union if union else 0.0


def _ios(a, b):
    x0, y0, x1, y1 = max(a[0], b[0]), max(a[1], b[1]), min(a[0]+a[2], b[0]+b[2]), min(a[1]+a[3], b[1]+b[3])
    inter = max(0, x1-x0) * max(0, y1-y0); smaller = min(a[2]*a[3], b[2]*b[3])
    return inter / smaller if smaller else 0.0


def _coarse(t):
    return "ChoiceButton" if t in {"choice_checkbox", "choice_radio", "choice_select", "ChoiceButton"} else "Signature" if t in {"signature", "Signature"} else "TextBox"


def _union(a, b):
    x0, y0, x1, y1 = min(a[0], b[0]), min(a[1], b[1]), max(a[0]+a[2], b[0]+b[2]), max(a[1]+a[3], b[1]+b[3])
    return [x0, y0, x1-x0, y1-y0]


def _merge_duplicates(items, threshold=.5):
    kept = []
    for item in sorted(items, key=lambda x: -(x["box"][2]*x["box"][3])):
        hit = next((k for k in kept if k.get("type") == item.get("type") and _ios(k["box"], item["box"]) >= threshold), None)
        if hit is None: kept.append(dict(item))
        else: hit["box"] = _union(hit["box"], item["box"])
    return kept


def _correlate(det, vlm, threshold=.5):
    """FF-DETR owns boxes; Nutrient VLM owns labels, fine types, and radio links."""
    used, enriched = set(), []
    det_types = {"TextBox": "text", "ChoiceButton": "choice_checkbox", "Signature": "signature"}
    for d in det:
        claims = []
        for i, v in enumerate(vlm):
            if i in used: continue
            overlap, containment = _iou(d["box"], v["box"]), _ios(d["box"], v["box"])
            compatible = _coarse(v.get("type")) == _coarse(d.get("type"))
            if overlap >= threshold or (containment >= threshold and compatible):
                claims.append((compatible, max(overlap, containment), i, v))
        if not claims:
            enriched.append({**d, "type": det_types.get(d["type"], d["type"]), "source": "FF-DETR"}); continue
        claims.sort(key=lambda x: (x[0], x[1]), reverse=True); chosen = claims[0][3]; used.add(claims[0][2])
        for _, _, i, v in claims[1:]:
            if _iou(chosen["box"], v["box"]) >= threshold: used.add(i)
        enriched.append({"box": d["box"], "type": chosen.get("type") or d["type"],
                         "label": chosen.get("label") or "", "group_id": chosen.get("group_id"),
                         "score": d.get("score"), "source": "Hybrid"})
    return enriched + _merge_duplicates([{**v, "source": "VLM"} for i, v in enumerate(vlm) if i not in used])


# NOTE: page tiling was removed from this demo (2026-07) — it created too many PRECISION issues. Splitting
# dense pages into an overlapping grid and running the VLM per crop split/duplicated boxes at tile seams,
# and the per-tile offset + cross-tile de-dup merge misaligned the final boxes. A single full-page pass is
# marginally weaker on very dense forms but far more geometrically reliable, and much faster (one generate,
# not N — which also lets the ZeroGPU duration drop under the free-tier cap). See DEADENDS.md. Tiling still
# lives in the eval harness (src/formfield_ft/tiling.py) where seam merging is *scored*, not shown to a user.


def _compile_grammar(model, processor, schema):
    import xgrammar as xgr
    vocab = getattr(model.config, "vocab_size", None) or len(processor.tokenizer)
    info = xgr.TokenizerInfo.from_huggingface(processor.tokenizer, vocab_size=vocab)
    return xgr.GrammarCompiler(info).compile_json_schema(json.dumps(schema))


def _generate(model, processor, image, prompt, schema, grammar):
    import xgrammar as xgr
    original = image.size
    if max(image.size) != 1216:
        scale = 1216/max(image.size); image = image.resize((max(1,round(image.width*scale)),max(1,round(image.height*scale))))
    messages = [{"role":"user","content":[{"type":"image"},{"type":"text","text":prompt}]}]
    chat = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = processor(text=[chat], images=[image], return_tensors="pt").to("cuda")
    lp = [xgr.contrib.hf.LogitsProcessor(grammar)]
    with torch.no_grad(): generated = model.generate(**inputs, max_new_tokens=2048, do_sample=False, logits_processor=lp)
    text = processor.tokenizer.decode(generated[0,inputs["input_ids"].shape[1]:], skip_special_tokens=True)
    try: raw = json.loads(text[text.find("["):text.rfind("]")+1])
    except Exception: raw = []
    out=[]
    for item in raw if isinstance(raw,list) else []:
        try:
            x0,y0,x1,y1=item["box"]; w,h=original
            out.append({**item,"box":[x0*w/1000,y0*h/1000,(x1-x0)*w/1000,(y1-y0)*h/1000],"source":"VLM"})
        except Exception: pass
    return out


@spaces.GPU(duration=100)   # single full-page pass; small reservation so low-ZeroGPU-quota (free) visitors
                            # can still start a run. max_new_tokens=2048 keeps a cold 4B load+generate inside 100s.
def _infer_gpu(image, prompt, schema):
    state=STATE; vlm=state["vlm"]; processor=state["processor"]
    detector=state["detector"]; dip=state["detector_processor"]
    det_inputs=dip(images=[image],return_tensors="pt").to("cuda")
    with torch.no_grad(): raw=detector(**det_inputs)
    result=dip.post_process_object_detection(raw,target_sizes=[(image.height,image.width)],threshold=.5)[0]
    names={0:"TextBox",1:"ChoiceButton",2:"Signature"}; det=[]
    for box,label,score in zip(result["boxes"],result["labels"],result["scores"]):
        if int(label) in names:
            x0,y0,x1,y1=[float(v) for v in box];det.append({"box":[x0,y0,x1-x0,y1-y0],"type":names[int(label)],"label":"","group_id":None,"score":float(score),"source":"FF-DETR"})
    grammar=_compile_grammar(vlm,processor,schema)
    fields=_merge_duplicates(_generate(vlm,processor,image,prompt,schema,grammar))
    return det,fields,_correlate(det,fields)


def _draw(image, fields):
    out=image.copy();draw=ImageDraw.Draw(out)
    for f in fields:
        x,y,w,h=f["box"];color=COLORS.get(f.get("type"),"#d946ef");draw.rectangle((x,y,x+w,y+h),outline=color,width=3);draw.text((x+2,max(0,y-13)),f.get("type","field"),fill=color)
    return out


def _render(state, view):
    if not state:
        return None, [], "Run detection before switching output views."
    fields = state[view.lower().replace("ff-detr", "det")]
    rows_out = [[i+1, f.get("source"), f.get("type"), f.get("label") or "", f.get("group_id") or "",
                 "" if f.get("score") is None else round(f["score"], 3), [round(v) for v in f["box"]]]
                for i, f in enumerate(fields)]
    status = f"{len(fields)} fields · {view} view"
    return _draw(state["image"], fields), rows_out, status


def run(image, view, prompt, schema_text):
    if image is None:return {},None,[],"Upload a form page."
    try:schema=json.loads(schema_text)
    except Exception as exc:return {},None,[],f"Invalid JSON schema: {exc}"
    state=_load()
    if "error" in state:return {},None,[],f"Model unavailable ({type(state['error']).__name__}). The Space needs private-repo access."
    image=image.convert("RGB");det,vlm,hybrid=_infer_gpu(image,prompt,schema)
    result_state={"image":image,"det":det,"vlm":vlm,"hybrid":hybrid}
    return result_state,*_render(result_state,view)


INTRO = """# Detect and understand form fields
Upload an **empty digital form page** to extract field boxes, fine field types, text labels, and radio-group links as schema-constrained JSON. **VLM** shows the Nutrient model alone. **Hybrid** optionally replaces matched VLM boxes with tighter FF-DETR detections while retaining the VLM's semantics.

→ [model](https://huggingface.co/nutrientdocs/form-field-vlm) · [leaderboard](https://huggingface.co/spaces/nutrientdocs/form-field-vlm-leaderboard) · [benchmark](https://huggingface.co/datasets/nutrientdocs/form-field-vlm-benchmark)
"""
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 demo 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.
"""

_load()

with gr.Blocks(title="Form Field VLM") as demo:
    gr.Markdown(INTRO)
    with gr.Row():source=gr.Image(type="pil",label="Form page",height=520);result=gr.Image(type="pil",label="Detected fields",height=520)
    gr.Examples(
        examples=[[str(EXAMPLES_DIR / filename)] for _label, filename in EXAMPLES],
        example_labels=[label for label, _filename in EXAMPLES],
        inputs=[source],
        label="Try an example form",
        examples_per_page=9,
        run_on_click=False,
    )
    gr.Markdown('<sub>Benchmark examples derive from <a href="https://github.com/jbarrow/commonforms" target="_blank" rel="noopener">CommonForms</a>, CC-BY-4.0.</sub>')
    with gr.Row():view=gr.Radio(["Hybrid","VLM","FF-DETR"],value="Hybrid",label="Output");submit=gr.Button("Detect fields",variant="primary")
    gr.Markdown('<sub>For production latency, we recommend vLLM-backed inference; see the <a href="https://huggingface.co/spaces/nutrientdocs/form-field-vlm-leaderboard" target="_blank" rel="noopener">leaderboard</a> for measured vLLM run latency.</sub>')
    with gr.Accordion("Prompt and JSON schema",open=False):
        with gr.Row():prompt=gr.Textbox(value=PROMPT,lines=12,label="Prompt");schema=gr.Textbox(value=json.dumps(SCHEMA,indent=2),lines=12,label="JSON schema")
    result_state=gr.State({})
    status=gr.Markdown();table=gr.Dataframe(headers=["#","source","type","label","group_id","score","box"],wrap=True)
    submit.click(run,[source,view,prompt,schema],[result_state,result,table,status],show_progress="full")
    view.change(_render,[result_state,view],[result,table,status],show_progress="hidden")
    gr.Markdown('**FF-DETR attribution.** Hybrid and FF-DETR views use <a href="https://huggingface.co/jbarrow/FFDetr" target="_blank" rel="noopener">FF-DETR / CommonForms by jbarrow</a>, Apache-2.0, only for object detection and box localization. FF-DETR does not produce labels, fine field types, values, or radio-group links; those come from the Nutrient VLM.')
    gr.Markdown(ABOUT)

if __name__=="__main__":
    demo.queue(default_concurrency_limit=1, max_size=8).launch(theme=gr.themes.Soft(), ssr_mode=False)