File size: 3,653 Bytes
4dce182 d28e3c4 4dce182 d28e3c4 4dce182 d28e3c4 4dce182 d28e3c4 4dce182 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | """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 height = ratio * box height
KERNEL_END_TRIM = 0.25 # kernel width = box width - 2 * trim * box height
BLOCK_SHRINK = 0.3 # block kernel = block rectangle shrunk by this many line heights per side
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) # /2
self.s1 = _stage(16, widths[0], depths[0], 2, expand) # /4
self.s2 = _stage(widths[0], widths[1], depths[1], 2, expand) # /8
self.s3 = _stage(widths[1], widths[2], depths[2], 2, expand) # /16
self.s4 = _stage(widths[2], widths[3], depths[3], 2, expand) # /32
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), # /2
nn.ConvTranspose2d(32, 2, 2, 2), # /1: lines, blocks
)
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)
|