| import json, re, sys, collections |
| BOX = re.compile(r"<box><(\d+)><(\d+)><(\d+)><(\d+)></box>") |
| def iou(a, b): |
| ix = max(0, min(a[2], b[2]) - max(a[0], b[0])); iy = max(0, min(a[3], b[3]) - max(a[1], b[1])) |
| i = ix * iy; u = (a[2]-a[0])*(a[3]-a[1]) + (b[2]-b[0])*(b[3]-b[1]) - i |
| return i / u if u > 0 else 0.0 |
| def parse(text, dedupe=0.85, cap=30): |
| kept = [] |
| for m in BOX.finditer(text): |
| b = tuple(int(x) for x in m.groups()) |
| if b[2] <= b[0] or b[3] <= b[1]: continue |
| if dedupe and any(iou(b, o) > dedupe for o in kept): break |
| kept.append(b) |
| if len(kept) >= cap: break |
| return kept |
| def score(preds, S, thr=0.5, **kw): |
| tp = fp = fn = 0; ious = []; npred = []; stopped = 0; secs = []; exact = 0; n = 0 |
| for s in S: |
| p = preds.get(str(s["id"])) |
| if p is None: continue |
| n += 1; secs.append(p["sec"]); stopped += ("<|im_end|>" in p["text"]) |
| pb = parse(p["text"], **kw); npred.append(len(pb)); used = set() |
| for b in pb: |
| best, bj = 0, -1 |
| for j, g in enumerate(s["boxes"]): |
| if j in used: continue |
| v = iou(b, g) |
| if v > best: best, bj = v, j |
| if best >= thr: tp += 1; used.add(bj); ious.append(best) |
| else: fp += 1 |
| fn += len(s["boxes"]) - len(used); exact += (len(pb) == len(s["boxes"]) and len(used) == len(pb)) |
| P = tp / (tp + fp) if tp + fp else 0; R = tp / (tp + fn) if tp + fn else 0 |
| return {"n": n, "P": round(P, 3), "R": round(R, 3), "F1": round(2*P*R/(P+R), 3) if P+R else 0, |
| "mIoU(matched)": round(sum(ious)/len(ious), 3) if ious else 0, "pred/img": round(sum(npred)/n, 2), |
| "exact_imgs": exact, "sec/img": round(sum(secs)/n, 2)} |
| if __name__ == "__main__": |
| S = json.load(open(sys.argv[1])); files = sys.argv[2:] |
| allp = {} |
| for f in files: |
| d = json.load(open(f)) |
| allp.update(d if all(isinstance(v, dict) and "text" not in v for v in d.values()) else {"board": d}) |
| for name, preds in allp.items(): |
| print("%-18s" % name, score(preds, S), "| F1@0.75", score(preds, S, thr=0.75)["F1"]) |
|
|
