| """Page orientation classifier: 0 / 90 / 180 / 270 degrees, voted over text-dense patches.""" |
|
|
| from torch import nn |
|
|
| from .common import ConvBNAct |
|
|
| PATCH = 192 |
| ANGLES = (0, 90, 180, 270) |
|
|
|
|
| class OrientationNet(nn.Module): |
| def __init__(self, widths=(16, 32, 64, 88, 120)): |
| super().__init__() |
| layers, cin = [], 1 |
| for i, w in enumerate(widths): |
| layers += [ConvBNAct(cin, w, s=2), ConvBNAct(w, w, k=3, groups=w), ConvBNAct(w, w, k=1)] |
| cin = w |
| self.features = nn.Sequential(*layers) |
| self.pool = nn.AdaptiveAvgPool2d(1) |
| self.fc = nn.Sequential(nn.Dropout(0.2), nn.Linear(cin, len(ANGLES))) |
|
|
| def forward(self, x): |
| """x: (B, 1, 192, 192) in [0, 1]. Returns logits (B, 4).""" |
| return self.fc(self.pool(self.features(x - 0.5)).flatten(1)) |
|
|
