Spaces:
Sleeping
Sleeping
| """OpenHTR visual encoder with local and global feature mixing.""" | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from torch.nn.init import kaiming_normal_, ones_, trunc_normal_, zeros_ | |
| from .layers import DropPath, Identity | |
| __all__ = ['OpenHTR_Encoder'] | |
| class ConvBNLayer(nn.Module): | |
| def __init__( | |
| self, | |
| in_channels, | |
| out_channels, | |
| kernel_size=3, | |
| stride=1, | |
| padding=0, | |
| bias=False, | |
| groups=1, | |
| act=nn.GELU, | |
| ): | |
| super().__init__() | |
| self.conv = nn.Conv2d( | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| kernel_size=kernel_size, | |
| stride=stride, | |
| padding=padding, | |
| groups=groups, | |
| bias=bias, | |
| ) | |
| self.norm = nn.BatchNorm2d(out_channels) | |
| self.act = act() | |
| def forward(self, inputs): | |
| out = self.conv(inputs) | |
| out = self.norm(out) | |
| out = self.act(out) | |
| return out | |
| class LastStage(nn.Module): | |
| def __init__(self, in_channels, out_channels, last_drop, out_char_num=0): | |
| super().__init__() | |
| self.last_conv = nn.Linear(in_channels, out_channels, bias=False) | |
| self.hardswish = nn.Hardswish() | |
| self.dropout = nn.Dropout(p=last_drop) | |
| def forward(self, x, sz): | |
| x = x.reshape(-1, sz[0], sz[1], x.shape[-1]) | |
| x = x.mean(1) | |
| x = self.last_conv(x) | |
| x = self.hardswish(x) | |
| x = self.dropout(x) | |
| return x, [1, sz[1]] | |
| class Feat2D(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| def forward(self, x, sz): | |
| C = x.shape[-1] | |
| x = x.transpose(1, 2).reshape(-1, C, sz[0], sz[1]) | |
| return x, sz | |
| def build_rope_2d(head_dim, H, W, device, dtype=torch.float32): | |
| """2D RoPE angles for a h-major flattened grid (H*W, w fastest). | |
| Returns cos/sin of shape [N, head_dim // 2] for the interleaved | |
| rotation convention (token dims are rotated in (2i, 2i+1) pairs). | |
| The head dim is split evenly: first half of the rotation frequencies | |
| encode the H axis, second half the W axis. | |
| """ | |
| half = head_dim // 2 | |
| assert half % 2 == 0, 'head_dim must be divisible by 4 for 2D RoPE' | |
| inv_freq = 1.0 / (10000.0**(torch.arange(0, half, 2, device=device, | |
| dtype=dtype) / half)) | |
| h_pos = torch.arange(H, device=device, dtype=dtype) | |
| w_pos = torch.arange(W, device=device, dtype=dtype) | |
| ang_h = h_pos[:, None] * inv_freq[None, :] # [H, half//2] | |
| ang_w = w_pos[:, None] * inv_freq[None, :] # [W, half//2] | |
| ang_h = ang_h.repeat_interleave(W, dim=0) # [H*W, half//2] (h-major) | |
| ang_w = ang_w.repeat(H, 1) # [H*W, half//2] | |
| angles = torch.cat([ang_h, ang_w], dim=1) # [N, half] | |
| return angles.cos(), angles.sin() | |
| def rotate_interleaved(x, cos, sin): | |
| """RoPE rotation, interleaved-pair convention. | |
| x: [B, heads, N, D]; cos/sin: [N, D//2] (broadcastable). | |
| """ | |
| x1 = x[..., 0::2] | |
| x2 = x[..., 1::2] | |
| out1 = x1 * cos - x2 * sin | |
| out2 = x1 * sin + x2 * cos | |
| return torch.stack([out1, out2], dim=-1).flatten(-2) | |
| class RMSNorm(nn.Module): | |
| def __init__(self, dim, eps=1e-6): | |
| super().__init__() | |
| self.eps = eps | |
| self.weight = nn.Parameter(torch.ones(dim)) | |
| def forward(self, x): | |
| dtype = x.dtype | |
| x = x.float() | |
| x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) | |
| return self.weight * x.to(dtype) | |
| class SwiGLU(nn.Module): | |
| def __init__(self, dim, hidden_dim, drop=0.0): | |
| super().__init__() | |
| self.w1 = nn.Linear(dim, hidden_dim, bias=False) # gate | |
| self.w3 = nn.Linear(dim, hidden_dim, bias=False) # up | |
| self.w2 = nn.Linear(hidden_dim, dim, bias=False) # down | |
| self.drop = nn.Dropout(drop) | |
| def forward(self, x): | |
| return self.drop(self.w2(F.silu(self.w1(x)) * self.w3(x))) | |
| class ConvNeXtBlock(nn.Module): | |
| """ConvNeXt block operating on [B, C, H, W] feature maps.""" | |
| def __init__(self, dim, kernel_size=7, mlp_ratio=4, drop_path=0.0): | |
| super().__init__() | |
| self.dwconv = nn.Conv2d(dim, | |
| dim, | |
| kernel_size, | |
| 1, | |
| kernel_size // 2, | |
| groups=dim) | |
| self.norm = nn.LayerNorm(dim, eps=1e-6) | |
| hidden = int(dim * mlp_ratio) | |
| self.pw1 = nn.Linear(dim, hidden) | |
| self.act = nn.GELU() | |
| self.pw2 = nn.Linear(hidden, dim) | |
| self.drop_path = DropPath(drop_path) if drop_path > 0.0 else Identity() | |
| def forward(self, x): | |
| B, C, H, W = x.shape | |
| y = self.dwconv(x).permute(0, 2, 3, 1) # B H W C | |
| y = self.norm(y) | |
| y = self.pw2(self.act(self.pw1(y))) | |
| y = y.permute(0, 3, 1, 2) | |
| return x + self.drop_path(y) | |
| class GlobalAttentionBlock(nn.Module): | |
| """Pre-norm attention block with 2D RoPE, RMSNorm and SwiGLU. | |
| Input/output: [B, C, H, W]; attention runs on the flattened (h-major) | |
| sequence so global mixing covers the whole line. | |
| """ | |
| def __init__(self, | |
| dim, | |
| num_heads, | |
| mlp_ratio=2.67, | |
| drop=0.0, | |
| attn_drop=0.0, | |
| drop_path=0.0, | |
| qkv_bias=True): | |
| super().__init__() | |
| self.num_heads = num_heads | |
| self.head_dim = dim // num_heads | |
| assert self.head_dim * num_heads == dim | |
| assert self.head_dim % 4 == 0, 'head_dim must be divisible by 4' | |
| self.norm1 = RMSNorm(dim) | |
| self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) | |
| self.proj = nn.Linear(dim, dim) | |
| self.proj_drop = nn.Dropout(drop) | |
| self.attn_drop = attn_drop | |
| self.drop_path = DropPath(drop_path) if drop_path > 0.0 else Identity() | |
| self.norm2 = RMSNorm(dim) | |
| hidden = max(64, int(dim * mlp_ratio) // 64 * 64) | |
| self.mlp = SwiGLU(dim, hidden, drop=drop) | |
| def forward(self, x): | |
| # x: [B, C, H, W] | |
| B, C, H, W = x.shape | |
| y = x.flatten(2).transpose(1, 2) # B N C | |
| y2 = self.norm1(y) | |
| qkv = self.qkv(y2).reshape(B, -1, 3, self.num_heads, | |
| self.head_dim).permute(2, 0, 3, 1, 4) | |
| q, k, v = qkv.unbind(0) | |
| cos, sin = build_rope_2d(self.head_dim, H, W, x.device, | |
| q.dtype) | |
| cos = cos.to(q.dtype)[None, None] # 1 1 N hd/2 | |
| sin = sin.to(q.dtype)[None, None] | |
| q = rotate_interleaved(q, cos, sin) | |
| k = rotate_interleaved(k, cos, sin) | |
| y = F.scaled_dot_product_attention( | |
| q, k, v, dropout_p=self.attn_drop if self.training else 0.0) | |
| y = y.transpose(1, 2).reshape(B, -1, C) | |
| y = self.proj_drop(self.proj(y)) | |
| x = x.flatten(2).transpose(1, 2) + self.drop_path(y) | |
| x = x + self.drop_path(self.mlp(self.norm2(x))) | |
| return x.transpose(1, 2).reshape(B, C, H, W) | |
| class EncoderStage(nn.Module): | |
| def __init__(self, | |
| dim=64, | |
| out_dim=256, | |
| depth=3, | |
| mixer=['Conv'] * 3, | |
| kernel_sizes=[7] * 3, | |
| sub_k=[2, 1], | |
| num_heads=2, | |
| mlp_ratio=4, | |
| attn_mlp_ratio=2.67, | |
| qkv_bias=True, | |
| drop_rate=0.0, | |
| attn_drop_rate=0.0, | |
| drop_path=[0.1] * 3, | |
| downsample=None, | |
| **kwargs): | |
| super().__init__() | |
| self.blocks = nn.ModuleList() | |
| for i in range(depth): | |
| if mixer[i] == 'Conv': | |
| self.blocks.append( | |
| ConvNeXtBlock(dim=dim, | |
| kernel_size=kernel_sizes[i], | |
| mlp_ratio=mlp_ratio, | |
| drop_path=drop_path[i])) | |
| else: | |
| self.blocks.append( | |
| GlobalAttentionBlock( | |
| dim=dim, | |
| num_heads=num_heads, | |
| mlp_ratio=attn_mlp_ratio, | |
| drop=drop_rate, | |
| attn_drop=attn_drop_rate, | |
| drop_path=drop_path[i], | |
| qkv_bias=qkv_bias, | |
| )) | |
| if downsample: | |
| self.downsample = EncoderDownsample(dim, out_dim, stride=sub_k) | |
| else: | |
| self.downsample = None | |
| def forward(self, x, sz): | |
| for blk in self.blocks: | |
| x = blk(x) | |
| if self.downsample is not None: | |
| x, sz = self.downsample(x, sz) | |
| return x, sz | |
| class EncoderDownsample(nn.Module): | |
| """Strided 3x3 conv + RMSNorm, operating on [B, C, H, W] (stays 4D).""" | |
| def __init__(self, in_channels, out_channels, stride=[2, 1]): | |
| super().__init__() | |
| self.conv = nn.Conv2d(in_channels, | |
| out_channels, | |
| kernel_size=3, | |
| stride=stride, | |
| padding=1) | |
| self.norm = RMSNorm(out_channels) | |
| def forward(self, x, sz): | |
| x = self.conv(x) | |
| C, H, W = x.shape[1:] | |
| y = self.norm(x.permute(0, 2, 3, 1)) # B H W C | |
| return y.permute(0, 3, 1, 2).contiguous(), [H, W] | |
| class OpenHTR_Encoder(nn.Module): | |
| """Encode local strokes and global context as a sequence of visual tokens.""" | |
| def __init__(self, | |
| max_sz=[32, 128], | |
| in_channels=3, | |
| out_channels=192, | |
| depths=[3, 6, 3], | |
| dims=[64, 128, 256], | |
| mixer=[['Conv'] * 3, ['Conv'] * 3 + ['Global'] * 3, | |
| ['Global'] * 3], | |
| sub_k=[[1, 1], [2, 1], [1, 1]], | |
| num_heads=[2, 4, 8], | |
| mlp_ratio=4, | |
| attn_mlp_ratio=2.67, | |
| qkv_bias=True, | |
| drop_rate=0.0, | |
| last_drop=0.1, | |
| attn_drop_rate=0.0, | |
| drop_path_rate=0.1, | |
| last_stage=False, | |
| feat2d=False, | |
| num_convs=None, # unused, kept for config compat | |
| kernel_sizes=[[7] * 3, [7] * 3 + [7] * 3, [7] * 3], | |
| **kwargs): | |
| super().__init__() | |
| num_stages = len(depths) | |
| self.num_features = dims[-1] | |
| # Two stride-2 convolutions reduce the input resolution by a factor of four. | |
| self.pope = nn.Sequential( | |
| ConvBNLayer(in_channels=in_channels, | |
| out_channels=dims[0] // 2, | |
| kernel_size=3, | |
| stride=2, | |
| padding=1, | |
| act=nn.GELU), | |
| ConvBNLayer(in_channels=dims[0] // 2, | |
| out_channels=dims[0], | |
| kernel_size=3, | |
| stride=2, | |
| padding=1, | |
| act=nn.GELU), | |
| ) | |
| dpr = np.linspace(0, drop_path_rate, sum(depths)) | |
| self.stages = nn.ModuleList() | |
| for i_stage in range(num_stages): | |
| stage = EncoderStage( | |
| dim=dims[i_stage], | |
| out_dim=dims[i_stage + 1] if i_stage < num_stages - 1 else 0, | |
| depth=depths[i_stage], | |
| mixer=mixer[i_stage], | |
| kernel_sizes=kernel_sizes[i_stage] if len( | |
| kernel_sizes[i_stage]) == len(mixer[i_stage]) else [7] * | |
| len(mixer[i_stage]), | |
| sub_k=sub_k[i_stage], | |
| num_heads=num_heads[i_stage], | |
| mlp_ratio=mlp_ratio, | |
| attn_mlp_ratio=attn_mlp_ratio, | |
| qkv_bias=qkv_bias, | |
| drop=drop_rate, | |
| attn_drop=attn_drop_rate, | |
| drop_path=dpr[sum(depths[:i_stage]):sum(depths[:i_stage + 1])], | |
| downsample=False if i_stage == num_stages - 1 else True, | |
| ) | |
| self.stages.append(stage) | |
| self.out_channels = self.num_features | |
| self.last_stage = last_stage | |
| self.last_stage_layer = None | |
| self.feat2d_layer = None | |
| if last_stage: | |
| self.out_channels = out_channels | |
| self.last_stage_layer = LastStage(self.num_features, out_channels, | |
| last_drop) | |
| if feat2d: | |
| self.feat2d_layer = Feat2D() | |
| self.apply(self._init_weights) | |
| def _init_weights(self, m: nn.Module): | |
| if isinstance(m, nn.Linear): | |
| trunc_normal_(m.weight, mean=0, std=0.02) | |
| if m.bias is not None: | |
| zeros_(m.bias) | |
| if isinstance(m, nn.LayerNorm): | |
| zeros_(m.bias) | |
| ones_(m.weight) | |
| if isinstance(m, nn.Conv2d): | |
| kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') | |
| def no_weight_decay(self): | |
| return {'patch_embed', 'downsample', 'pos_embed'} | |
| def forward(self, x): | |
| if len(x.shape) == 5: | |
| x = x.flatten(0, 1) | |
| sz = x.shape[2:] | |
| x = self.pope(x) | |
| sz = [sz[0] // 4, sz[1] // 4] | |
| for stage in self.stages: | |
| x, sz = stage(x, sz) | |
| # All EncoderStage blocks keep [B, C, H, W]; convert once at the end. | |
| if self.last_stage_layer is not None: | |
| x, sz = self.last_stage_layer(x.flatten(2).transpose(1, 2), sz) | |
| elif self.feat2d_layer is not None: | |
| pass # already [B, C, H, W] | |
| else: | |
| x = x.flatten(2).transpose(1, 2) # B N C | |
| return x | |