| |
| """ |
| detect.py -- run docforensics on a page image and draw the regions. |
| |
| python detect.py page.png [out.png] |
| |
| Needs: onnxruntime, opencv-python, numpy. |
| """ |
| import sys |
| import cv2 |
| import numpy as np |
| import onnxruntime as ort |
|
|
| CLASSES = ["text", "heading", "list", "table", "picture", "caption", "formula", |
| "header", "footer", "footnote", "seal", "form", "field", "checkbox", |
| "signature"] |
| SIGNALS = ["skew", "blur", "noise", "bleed", "warp", "tears", "handwriting", "native"] |
| COLOURS = {"text": (75, 180, 60), "heading": (75, 25, 230), "list": (25, 225, 255), |
| "table": (200, 130, 0), "picture": (48, 130, 245), "caption": (180, 30, 145), |
| "field": (90, 220, 60), "checkbox": (230, 0, 0), "form": (0, 140, 255)} |
| CONF, IOU = 0.18, 0.5 |
|
|
|
|
| def letterbox(img, side=640): |
| """Match training preprocessing exactly: fit long side, pad with 114.""" |
| h, w = img.shape[:2] |
| s = side / max(h, w) |
| small = cv2.resize(img, (int(w * s), int(h * s))) |
| canvas = np.full((side, side, 3), 114, np.uint8) |
| canvas[:small.shape[0], :small.shape[1]] = small |
| return canvas |
|
|
|
|
| def nms(boxes, scores, thr): |
| order = scores.argsort()[::-1] |
| keep = [] |
| while order.size: |
| i = order[0] |
| keep.append(i) |
| if order.size == 1: |
| break |
| xx1 = np.maximum(boxes[i, 0], boxes[order[1:], 0]) |
| yy1 = np.maximum(boxes[i, 1], boxes[order[1:], 1]) |
| xx2 = np.minimum(boxes[i, 2], boxes[order[1:], 2]) |
| yy2 = np.minimum(boxes[i, 3], boxes[order[1:], 3]) |
| inter = np.maximum(0, xx2 - xx1) * np.maximum(0, yy2 - yy1) |
| a = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) |
| order = order[1:][inter / (a[i] + a[order[1:]] - inter) <= thr] |
| return keep |
|
|
|
|
| def detect(path, model="docforensics-layout-s.onnx"): |
| raw = cv2.imread(path) |
| if raw is None: |
| sys.exit(f"cannot read {path}") |
| img = letterbox(raw) |
| x = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0 |
| x = ((x - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]) |
| x = x.transpose(2, 0, 1)[None].astype(np.float32) |
|
|
| sess = ort.InferenceSession(model, providers=["CPUExecutionProvider"]) |
| det, signals, _maps = sess.run( |
| None, {"images": x, "evidence": np.zeros((1, 4), np.float32)}) |
|
|
| |
| |
| d = det[0].T |
| cx, cy, bw, bh = d[:, 0], d[:, 1], d[:, 2], d[:, 3] |
| xyxy = np.stack([cx - bw / 2, cy - bh / 2, cx + bw / 2, cy + bh / 2], 1) |
| scores, labels = d[:, 4:].max(1), d[:, 4:].argmax(1) |
|
|
| out = [] |
| for c in np.unique(labels[scores > CONF]): |
| m = (labels == c) & (scores > CONF) |
| for i in nms(xyxy[m], scores[m], IOU): |
| out.append((xyxy[m][i], scores[m][i], CLASSES[c])) |
| return img, out, dict(zip(SIGNALS, signals[0].round(3).tolist())) |
|
|
|
|
| if __name__ == "__main__": |
| if len(sys.argv) < 2: |
| sys.exit(__doc__) |
| img, regions, signals = detect(sys.argv[1]) |
| for box, score, name in regions: |
| c = COLOURS.get(name, (150, 150, 150)) |
| cv2.rectangle(img, (int(box[0]), int(box[1])), |
| (int(box[2]), int(box[3])), c, 2) |
| cv2.putText(img, f"{name} {score:.2f}", (int(box[0]), max(10, int(box[1]) - 4)), |
| cv2.FONT_HERSHEY_SIMPLEX, 0.35, c, 1) |
| print(f"{len(regions)} regions") |
| for k, v in signals.items(): |
| print(f" {k:12s} {v:.3f}") |
| dst = sys.argv[2] if len(sys.argv) > 2 else "out.png" |
| cv2.imwrite(dst, img) |
| print(f"wrote {dst}") |
|
|
