File size: 869 Bytes
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 | """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) # clockwise rotation the page has undergone
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))
|