| """Детектор строк: компактный DB (differentiable binarization). |
| |
| Пиксельная морфология упирается в пороги: один набор чисел разбирает газету и |
| рвёт книгу. Детектор учится на рамках строк, поэтому переносится на любую |
| вёрстку, лишь бы строка была строкой. |
| |
| Почему именно DB, а не «просто сегментация текста»: в плотном газетном наборе |
| соседние строки и колонки сливаются в одно пятно. DB учит карту по *сжатым* |
| рамкам — между строками остаётся зазор — и отдельно учит порог, по которому |
| пятно снова разворачивается до полной строки. |
| |
| Веса backbone не берутся предобученными: download.pytorch.org из пода |
| недоступен, а бесплатной разметки из born-digital PDF хватает для обучения |
| с нуля. |
| """ |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| def conv_bn(cin, cout, k=3, s=1): |
| return nn.Sequential( |
| nn.Conv2d(cin, cout, k, s, k // 2, bias=False), |
| nn.BatchNorm2d(cout), |
| nn.ReLU(inplace=True)) |
|
|
|
|
| class Block(nn.Module): |
| """Остаточный блок; понижение разрешения — первой свёрткой.""" |
|
|
| def __init__(self, cin, cout, stride=1): |
| super().__init__() |
| self.c1 = conv_bn(cin, cout, 3, stride) |
| self.c2 = nn.Sequential( |
| nn.Conv2d(cout, cout, 3, 1, 1, bias=False), |
| nn.BatchNorm2d(cout)) |
| self.skip = None |
| if stride != 1 or cin != cout: |
| self.skip = nn.Sequential( |
| nn.Conv2d(cin, cout, 1, stride, bias=False), |
| nn.BatchNorm2d(cout)) |
|
|
| def forward(self, x): |
| y = self.c2(self.c1(x)) |
| s = x if self.skip is None else self.skip(x) |
| return F.relu(y + s, inplace=True) |
|
|
|
|
| class LineDet(nn.Module): |
| """Вход: 1×H×W (0..1, чернила ближе к 1). Выход: карты в 1/2 разрешения. |
| |
| Шаг выхода 2, а не 4 как в оригинальном DB: в газете строка занимает |
| 10–14 px, при шаге 4 сжатая рамка вырождается в один пиксель и учить |
| нечему. |
| """ |
|
|
| def __init__(self, width=32): |
| super().__init__() |
| w = width |
| self.stem = nn.Sequential(conv_bn(1, w, 3, 2), Block(w, w)) |
| self.s2 = nn.Sequential(Block(w, w * 2, 2), Block(w * 2, w * 2)) |
| self.s3 = nn.Sequential(Block(w * 2, w * 4, 2), Block(w * 4, w * 4)) |
| self.s4 = nn.Sequential(Block(w * 4, w * 8, 2), Block(w * 8, w * 8)) |
|
|
| self.l1 = nn.Conv2d(w, w * 2, 1) |
| self.l2 = nn.Conv2d(w * 2, w * 2, 1) |
| self.l3 = nn.Conv2d(w * 4, w * 2, 1) |
| self.l4 = nn.Conv2d(w * 8, w * 2, 1) |
| self.smooth = conv_bn(w * 2 * 4, w * 4) |
|
|
| |
| self.prob = nn.Sequential(conv_bn(w * 4, w * 2), nn.Conv2d(w * 2, 1, 1)) |
| self.thr = nn.Sequential(conv_bn(w * 4, w * 2), nn.Conv2d(w * 2, 1, 1)) |
| self.k = 50.0 |
|
|
| def forward(self, x): |
| p1 = self.stem(x) |
| p2 = self.s2(p1) |
| p3 = self.s3(p2) |
| p4 = self.s4(p3) |
| c4 = self.l4(p4) |
| c3 = self.l3(p3) + F.interpolate(c4, scale_factor=2, mode="nearest") |
| c2 = self.l2(p2) + F.interpolate(c3, scale_factor=2, mode="nearest") |
| c1 = self.l1(p1) + F.interpolate(c2, scale_factor=2, mode="nearest") |
| f = torch.cat([ |
| c1, |
| F.interpolate(c2, size=c1.shape[-2:], mode="nearest"), |
| F.interpolate(c3, size=c1.shape[-2:], mode="nearest"), |
| F.interpolate(c4, size=c1.shape[-2:], mode="nearest")], 1) |
| f = self.smooth(f) |
| prob = self.prob(f) |
| thr = self.thr(f) |
| return prob, thr |
|
|
| def binary(self, prob, thr): |
| """Дифференцируемая бинаризация: мягкая ступенька вместо сравнения.""" |
| return torch.sigmoid(self.k * (torch.sigmoid(prob) - torch.sigmoid(thr))) |
|
|
|
|
| def dice_loss(pred, target, mask, eps=1e-6): |
| pred, target, mask = pred * mask, target * mask, mask |
| inter = (pred * target).sum() |
| return 1.0 - 2.0 * inter / (pred.sum() + target.sum() + eps) |
|
|
|
|
| def ohem_bce(logit, target, mask, ratio=3.0): |
| """BCE только по всем положительным и втрое большему числу худших отрицательных. |
| |
| Текст занимает малую долю страницы; без отбора отрицательных обучение |
| уезжает в «везде фон». |
| """ |
| pos = (target > 0.5) & (mask > 0.5) |
| neg = (target <= 0.5) & (mask > 0.5) |
| n_pos = int(pos.sum()) |
| n_neg = min(int(neg.sum()), max(int(n_pos * ratio), 1000)) |
| loss = F.binary_cross_entropy_with_logits(logit, target, reduction="none") |
| lp = loss[pos].sum() if n_pos else logit.sum() * 0.0 |
| ln = torch.topk(loss[neg], n_neg).values.sum() if n_neg else logit.sum() * 0.0 |
| return (lp + ln) / max(1, n_pos + n_neg) |
|
|
|
|
| def det_loss(model, prob_logit, thr_logit, gt_prob, gt_thr, mask_prob, mask_thr): |
| l_prob = ohem_bce(prob_logit, gt_prob, mask_prob) |
| l_thr = (torch.abs(torch.sigmoid(thr_logit) - gt_thr) * mask_thr).sum() \ |
| / mask_thr.sum().clamp(min=1) |
| binary = model.binary(prob_logit, thr_logit) |
| l_bin = dice_loss(binary, gt_prob, mask_prob) |
| return l_prob + 10.0 * l_thr + l_bin, (l_prob.item(), l_thr.item(), l_bin.item()) |
|
|
