| |
| |
| """Score a page read by a document reader: P(each table cell or text region is misread) and the likely error. |
| pip install torch safetensors numpy pillow (transformers only for a vision variant; this one has none) |
| |
| python usage.py page.png page.json # prints the 20 units most likely misread |
| |
| The page dict is one page of a decosa.docreader.v1 document: {"elements": [{id, kind, bbox: [x, y, w, h], text?, |
| table?: {rows, cols, cells: [{r, c, rs, cs, text, tok?}]}, source, conf_bp?, min_bp?, ...}]}, with bboxes in the pixels of |
| the page image you pass. Cell token confidences (cell["tok"]) are optional but help a lot (see the eval). |
| """ |
| import json |
| import sys |
|
|
| import numpy as np |
| import torch |
| from PIL import Image |
| from safetensors.torch import load_file |
|
|
| import features as F |
| from model import TYPES, ReadConf, encode_text |
|
|
| D = "." |
| cfg = json.load(open(f"{D}/readconf_config.json")) |
| vcfg = None |
| if cfg.get("vision"): |
| from transformers import Siglip2VisionConfig |
| vcfg = Siglip2VisionConfig.from_dict(cfg["vision"]["config"]) |
| m = ReadConf(cfg["n_feat"], vision_config=vcfg, vision_layers=(cfg.get("vision") or {}).get("layers", 0)) |
| m.load_state_dict(load_file(f"{D}/model.safetensors")) |
| m.eval() |
|
|
|
|
| @torch.no_grad() |
| def score(image: Image.Image, page: dict) -> list[dict]: |
| us = F.page_units(image.convert("RGB"), page) |
| if not us: |
| return [] |
| x = torch.tensor(np.array([u["x"] for u in us], np.float32)) |
| ch = encode_text([u["text"] for u in us]) |
| img = None |
| if m.vision_layers: |
| w = cfg["crop"][1] |
| a = np.full((len(us), F.CROP_H, w), 255, np.uint8) |
| for i, u in enumerate(us): |
| c = F.crop(image.convert("RGB"), u, max_w=w) |
| a[i, :, : c.width] = np.asarray(c, np.uint8) |
| img = torch.tensor(a) |
| lw, lt = m(x, ch, img) |
| p, tp = m.p_wrong(lw), lt.softmax(-1) |
| out = [] |
| for i, u in enumerate(us): |
| |
| out.append({"uid": u["uid"], "text": u["text"], "p_wrong": round(float(p[i]), 4), "type": TYPES[int(tp[i].argmax())]}) |
| return out |
|
|
|
|
| if __name__ == "__main__": |
| im = Image.open(sys.argv[1]) |
| page = json.load(open(sys.argv[2])) |
| if "pages" in page: |
| page = page["pages"][0] |
| for r in sorted(score(im, page), key=lambda r: -r["p_wrong"])[:20]: |
| print(r) |
|
|
