| """Text-segment detector: MobileNetV2-style backbone + FPN + segmentation head. |
| |
| Predicts two probability maps at input resolution: |
| 0. line kernels: each text segment's central region (half its height, with its ends trimmed by a |
| quarter height). Post-processing expands each kernel back to the full box, so adjacent lines stay |
| separated even when their line spacing is tight. |
| 1. block kernels: each layout block (paragraph, list item, heading, table, caption...) as the rectangle |
| around its lines shrunk by BLOCK_SHRINK line heights on every side, so neighbouring blocks stay apart. |
| Lines are grouped into paragraphs by the block region their centre falls in. |
| """ |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
|
|
| from .common import ConvBNAct |
|
|
| KERNEL_HEIGHT_RATIO = 0.5 |
| KERNEL_END_TRIM = 0.25 |
| BLOCK_SHRINK = 0.3 |
|
|
|
|
| class InvertedResidual(nn.Module): |
| def __init__(self, cin, cout, stride=1, expand=4): |
| super().__init__() |
| mid = cin * expand |
| self.use_res = stride == 1 and cin == cout |
| self.body = nn.Sequential( |
| ConvBNAct(cin, mid, k=1), |
| ConvBNAct(mid, mid, k=3, s=stride, groups=mid), |
| ConvBNAct(mid, cout, k=1, act=False), |
| ) |
|
|
| def forward(self, x): |
| y = self.body(x) |
| return x + y if self.use_res else y |
|
|
|
|
| def _stage(cin, cout, n, stride, expand): |
| blocks = [InvertedResidual(cin, cout, stride, expand)] |
| blocks += [InvertedResidual(cout, cout, 1, expand) for _ in range(n - 1)] |
| return nn.Sequential(*blocks) |
|
|
|
|
| class Detector(nn.Module): |
| def __init__(self, widths=(24, 48, 104, 192), depths=(2, 3, 4, 3), expand=4, fpn=64): |
| super().__init__() |
| self.stem = ConvBNAct(3, 16, s=2) |
| self.s1 = _stage(16, widths[0], depths[0], 2, expand) |
| self.s2 = _stage(widths[0], widths[1], depths[1], 2, expand) |
| self.s3 = _stage(widths[1], widths[2], depths[2], 2, expand) |
| self.s4 = _stage(widths[2], widths[3], depths[3], 2, expand) |
| self.lateral = nn.ModuleList(nn.Conv2d(w, fpn, 1) for w in widths) |
| self.smooth = nn.ModuleList(ConvBNAct(fpn, fpn // 4) for _ in widths) |
| self.head = nn.Sequential( |
| ConvBNAct(fpn, 64), |
| nn.ConvTranspose2d(64, 32, 2, 2), nn.BatchNorm2d(32), nn.ReLU(inplace=True), |
| nn.ConvTranspose2d(32, 2, 2, 2), |
| ) |
|
|
| def forward(self, x): |
| """x: (B, 3, H, W) in [0, 1], H and W multiples of 32. Returns logits (B, 2, H, W): line, block kernels.""" |
| x = self.stem(x - 0.5) |
| c2 = self.s1(x) |
| c3 = self.s2(c2) |
| c4 = self.s3(c3) |
| c5 = self.s4(c4) |
| p5 = self.lateral[3](c5) |
| p4 = self.lateral[2](c4) + F.interpolate(p5, size=c4.shape[-2:], mode="nearest") |
| p3 = self.lateral[1](c3) + F.interpolate(p4, size=c3.shape[-2:], mode="nearest") |
| p2 = self.lateral[0](c2) + F.interpolate(p3, size=c2.shape[-2:], mode="nearest") |
| size = p2.shape[-2:] |
| fused = torch.cat([ |
| self.smooth[0](p2), |
| F.interpolate(self.smooth[1](p3), size=size, mode="nearest"), |
| F.interpolate(self.smooth[2](p4), size=size, mode="nearest"), |
| F.interpolate(self.smooth[3](p5), size=size, mode="nearest"), |
| ], dim=1) |
| return self.head(fused) |
|
|
