| """ |
| SegFormer (MiT-B0 scale) for RockGlacierNet — transformer baseline. |
| |
| Faithful compact implementation of Xie et al. (2021) "SegFormer": |
| - Hierarchical Transformer encoder (4 stages, overlap patch embed) |
| - Efficient self-attention (spatial-reduction) + Mix-FFN (depthwise conv) |
| - Lightweight all-MLP decoder that fuses the 4 multi-scale features |
| |
| Input : (128, 128, N_BANDS) |
| Output: (128, 128, 1) sigmoid — binary glacier mask |
| |
| Encoder dims : 32 -> 64 -> 160 -> 256 (MiT-B0) |
| Decoder dim : 256 |
| """ |
|
|
| import tensorflow as tf |
| from tensorflow.keras import layers, Model |
|
|
|
|
| def _overlap_patch_embed(x, out_ch, patch, stride, name): |
| x = layers.Conv2D(out_ch, patch, strides=stride, padding='same', name=f'{name}_proj')(x) |
| x = layers.LayerNormalization(epsilon=1e-6, name=f'{name}_ln')(x) |
| return x |
|
|
|
|
| def _efficient_attention(x, dim, num_heads, sr_ratio, name): |
| """Spatial-reduction self-attention on a (B, H, W, C) tensor.""" |
| H, W = x.shape[1], x.shape[2] |
| q = layers.Reshape((H * W, dim), name=f'{name}_qseq')(x) |
| if sr_ratio > 1: |
| kv = layers.Conv2D(dim, sr_ratio, strides=sr_ratio, padding='same', |
| name=f'{name}_sr')(x) |
| kv = layers.LayerNormalization(epsilon=1e-6, name=f'{name}_srln')(kv) |
| h2, w2 = kv.shape[1], kv.shape[2] |
| kv = layers.Reshape((h2 * w2, dim), name=f'{name}_kvseq')(kv) |
| else: |
| kv = q |
| attn = layers.MultiHeadAttention(num_heads=num_heads, key_dim=max(dim // num_heads, 1), |
| name=f'{name}_mha')(q, kv) |
| attn = layers.Reshape((H, W, dim), name=f'{name}_out')(attn) |
| return attn |
|
|
|
|
| def _mix_ffn(x, dim, expansion, name): |
| h = layers.Dense(dim * expansion, name=f'{name}_fc1')(x) |
| h = layers.DepthwiseConv2D(3, padding='same', name=f'{name}_dw')(h) |
| h = layers.Activation('gelu', name=f'{name}_gelu')(h) |
| h = layers.Dense(dim, name=f'{name}_fc2')(h) |
| return h |
|
|
|
|
| def _transformer_block(x, dim, num_heads, sr_ratio, name): |
| a = layers.LayerNormalization(epsilon=1e-6, name=f'{name}_ln1')(x) |
| a = _efficient_attention(a, dim, num_heads, sr_ratio, f'{name}_attn') |
| x = layers.Add(name=f'{name}_add1')([x, a]) |
| f = layers.LayerNormalization(epsilon=1e-6, name=f'{name}_ln2')(x) |
| f = _mix_ffn(f, dim, 4, f'{name}_ffn') |
| x = layers.Add(name=f'{name}_add2')([x, f]) |
| return x |
|
|
|
|
| def build_segformer(input_shape=(128, 128, 12), num_classes=1): |
| """SegFormer MiT-B0. Signature matches the other build_* functions.""" |
| dims = [32, 64, 160, 256] |
| heads = [1, 2, 5, 8] |
| sr_ratios = [8, 4, 2, 1] |
| depths = [2, 2, 2, 2] |
| patch_cfg = [(7, 4), (3, 2), (3, 2), (3, 2)] |
| dec_dim = 256 |
|
|
| inp = layers.Input(input_shape, name='input') |
| x = inp |
| feats = [] |
| for i in range(4): |
| p, s = patch_cfg[i] |
| x = _overlap_patch_embed(x, dims[i], p, s, f'stage{i+1}_embed') |
| for d in range(depths[i]): |
| x = _transformer_block(x, dims[i], heads[i], sr_ratios[i], f'stage{i+1}_blk{d}') |
| x = layers.LayerNormalization(epsilon=1e-6, name=f'stage{i+1}_norm')(x) |
| feats.append(x) |
|
|
| |
| target_hw = feats[0].shape[1] |
| ups = [] |
| for i, f in enumerate(feats): |
| u = layers.Dense(dec_dim, name=f'dec_proj{i}')(f) |
| scale = target_hw // f.shape[1] |
| if scale > 1: |
| u = layers.UpSampling2D(size=scale, interpolation='bilinear', |
| name=f'dec_up{i}')(u) |
| ups.append(u) |
| x = layers.Concatenate(name='dec_concat')(ups) |
| x = layers.Conv2D(dec_dim, 1, name='dec_fuse')(x) |
| x = layers.BatchNormalization(name='dec_bn')(x) |
| x = layers.Activation('relu', name='dec_relu')(x) |
| x = layers.Dropout(0.1, name='dec_drop')(x) |
|
|
| x = layers.UpSampling2D(size=input_shape[0] // target_hw, interpolation='bilinear', |
| name='final_up')(x) |
| out = layers.Conv2D(num_classes, 1, activation='sigmoid', name='output')(x) |
| return Model(inp, out, name='SegFormer') |
|
|
