| """Оценка детектора строк на отложенных выпусках. |
| |
| Меряем то, что важно продукту: сколько строк найдено (полнота), сколько |
| выдумано (точность) и на скольких страницах порядок чтения совпал с вёрсткой. |
| """ |
| import argparse, json, os, sys, time |
|
|
| import numpy as np |
| import torch |
| from PIL import Image, ImageFilter |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from linedet import LineDet |
| from detect import lines_from_prob, iou_match |
|
|
|
|
| @torch.no_grad() |
| def run_page(model, im, dev, stride, max_side=2600): |
| """Карта вероятностей для целой страницы; длинная сторона ограничена.""" |
| s = min(1.0, max_side / max(im.width, im.height)) |
| if s < 1.0: |
| im = im.resize((int(im.width * s), int(im.height * s)), Image.BILINEAR) |
| w, h = im.width, im.height |
| pw, ph = (w + 31) // 32 * 32, (h + 31) // 32 * 32 |
| a = np.zeros((ph, pw), np.float32) |
| a[:h, :w] = 1.0 - np.asarray(im, np.float32) / 255.0 |
| x = torch.from_numpy(a)[None, None].to(dev) |
| with torch.amp.autocast("cuda", enabled=(dev == "cuda")): |
| prob, _ = model(x) |
| prob = torch.sigmoid(prob.float())[0, 0].cpu().numpy() |
| prob = prob[:h // stride, :w // stride] |
| return prob, s |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--ckpt", required=True) |
| ap.add_argument("--data", required=True) |
| ap.add_argument("--pages", type=int, default=0) |
| ap.add_argument("--skip", type=int, default=0, |
| help="пропустить первые N страниц: по ним отбирался чекпоинт") |
| ap.add_argument("--iou", type=float, default=0.5) |
| ap.add_argument("--max-side", type=int, default=2000) |
| ap.add_argument("--unclip", type=float, default=1.6) |
| ap.add_argument("--book-scale", type=float, default=4.0) |
| ap.add_argument("--book", action="store_true", |
| help="книжно-экранное условие: та же разметка, крупнее и бледнее") |
| ap.add_argument("--max-gap", type=float, default=2.0, |
| help="порог склейки кусков строки в долях высоты строки") |
| a = ap.parse_args() |
|
|
| dev = "cuda" if torch.cuda.is_available() else "cpu" |
| ck = torch.load(a.ckpt, map_location="cpu", weights_only=False) |
| model = LineDet(ck.get("width", 32)).to(dev).eval() |
| model.load_state_dict(ck["model"]) |
| stride = ck.get("stride", 2) |
|
|
| rows = [json.loads(l) for l in |
| open(os.path.join(a.data, "index.jsonl"), encoding="utf-8")] |
| rows = rows[a.skip:] |
| if a.pages: |
| rows = rows[:a.pages] |
|
|
| tot_rec = tot_prec = 0.0 |
| n_gt = n_pred = n_hit = 0 |
| ok_pairs = all_pairs = 0 |
| t0 = time.time() |
| for r in rows: |
| im = Image.open(os.path.join(a.data, r["f"])).convert("L") |
| pre, off = 1.0, (0, 0) |
| gt_src = r["b"] |
| if a.book: |
| |
| |
| |
| |
| |
| pre = a.book_scale |
| big = im.resize((int(im.width * pre), int(im.height * pre)), Image.BILINEAR) |
| cw = min(big.width, a.max_side) |
| ch = min(big.height, a.max_side) |
| |
| |
| cx = float(np.mean([(b[0] + b[2]) / 2 for b in r["b"]])) * pre |
| cy = float(np.mean([(b[1] + b[3]) / 2 for b in r["b"]])) * pre |
| ox = int(min(max(0.0, cx - cw / 2), big.width - cw)) |
| oy = int(min(max(0.0, cy - ch / 2), big.height - ch)) |
| im = big.crop((ox, oy, ox + cw, oy + ch)) |
| arr = np.asarray(im, np.float32) / 255.0 |
| im = Image.fromarray(((0.30 + arr * 0.60) * 255).astype(np.uint8)).filter( |
| ImageFilter.GaussianBlur(0.6)) |
| off = (ox, oy) |
| gt_src = [[b[0] * pre - ox, b[1] * pre - oy, b[2] * pre - ox, b[3] * pre - oy] |
| for b in r["b"]] |
| |
| |
| |
| keep = [] |
| for b in gt_src: |
| x0, y0 = max(0.0, b[0]), max(0.0, b[1]) |
| x1, y1 = min(float(cw), b[2]), min(float(ch), b[3]) |
| if x1 <= x0 or y1 <= y0: |
| continue |
| if (x1 - x0) * (y1 - y0) < 0.5 * (b[2] - b[0]) * (b[3] - b[1]): |
| continue |
| keep.append([x0, y0, x1, y1]) |
| gt_src = keep |
| pre = 1.0 |
| prob, s = run_page(model, im, dev, stride, a.max_side) |
| w_s, h_s = int(im.width * s), int(im.height * s) |
| pred, lh = lines_from_prob(prob, stride, w_s, h_s, a.max_gap, a.unclip) |
| gt = [tuple(int(v * pre * s) for v in b) for b in gt_src] |
| rec, prec, hit = iou_match(pred, gt, a.iou) |
| n_gt += len(gt); n_pred += len(pred); n_hit += hit |
|
|
| |
| ordered = pred |
| idx = [] |
| for b in ordered: |
| best, bi = 0.0, -1 |
| ba = max(1, (b[2]-b[0]) * (b[3]-b[1])) |
| for i, g in enumerate(gt): |
| ox = min(b[2], g[2]) - max(b[0], g[0]) |
| oy = min(b[3], g[3]) - max(b[1], g[1]) |
| if ox <= 0 or oy <= 0: |
| continue |
| v = ox * oy / ba |
| if v > best: |
| best, bi = v, i |
| if best > 0.25: |
| idx.append(bi) |
| ok_pairs += sum(1 for p, q in zip(idx, idx[1:]) if q >= p) |
| all_pairs += max(1, len(idx) - 1) |
|
|
| print(f"страниц {len(rows)}, строк в эталоне {n_gt}, найдено рамок {n_pred}") |
| print(f"IoU {a.iou}: полнота {n_hit/max(1,n_gt)*100:.1f}%, " |
| f"точность {n_hit/max(1,n_pred)*100:.1f}%") |
| print(f"порядок чтения: {ok_pairs/all_pairs*100:.1f}% соседних пар верно") |
| print(f"{(time.time()-t0)/len(rows)*1000:.0f} мс на страницу ({dev})") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
