first commit

ppocr/modeling/heads/rec_abinet_head.py

@@ -0,0 +1,299 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+This code is refer from:
+https://github.com/FangShancheng/ABINet/tree/main/modules
+"""
+
+import math
+import paddle
+from paddle import nn
+import paddle.nn.functional as F
+from paddle.nn import LayerList
+from ppocr.modeling.heads.rec_nrtr_head import TransformerBlock, PositionalEncoding
+
+
+class BCNLanguage(nn.Layer):
+    def __init__(
+        self,
+        d_model=512,
+        nhead=8,
+        num_layers=4,
+        dim_feedforward=2048,
+        dropout=0.0,
+        max_length=25,
+        detach=True,
+        num_classes=37,
+    ):
+        super().__init__()
+
+        self.d_model = d_model
+        self.detach = detach
+        self.max_length = max_length + 1  # additional stop token
+        self.proj = nn.Linear(num_classes, d_model, bias_attr=False)
+        self.token_encoder = PositionalEncoding(
+            dropout=0.1, dim=d_model, max_len=self.max_length
+        )
+        self.pos_encoder = PositionalEncoding(
+            dropout=0, dim=d_model, max_len=self.max_length
+        )
+
+        self.decoder = nn.LayerList(
+            [
+                TransformerBlock(
+                    d_model=d_model,
+                    nhead=nhead,
+                    dim_feedforward=dim_feedforward,
+                    attention_dropout_rate=dropout,
+                    residual_dropout_rate=dropout,
+                    with_self_attn=False,
+                    with_cross_attn=True,
+                )
+                for i in range(num_layers)
+            ]
+        )
+
+        self.cls = nn.Linear(d_model, num_classes)
+
+    def forward(self, tokens, lengths):
+        """
+        Args:
+            tokens: (B, N, C) where N is length, B is batch size and C is classes number
+            lengths: (B,)
+        """
+        if self.detach:
+            tokens = tokens.detach()
+        embed = self.proj(tokens)  # (B, N, C)
+        embed = self.token_encoder(embed)  # (B, N, C)
+        padding_mask = _get_mask(lengths, self.max_length)
+        zeros = paddle.zeros_like(embed)  # (B, N, C)
+        query = self.pos_encoder(zeros)
+        for decoder_layer in self.decoder:
+            query = decoder_layer(query, embed, cross_mask=padding_mask)
+        output = query  # (B, N, C)
+
+        logits = self.cls(output)  # (B, N, C)
+
+        return output, logits
+
+
+def encoder_layer(in_c, out_c, k=3, s=2, p=1):
+    return nn.Sequential(
+        nn.Conv2D(in_c, out_c, k, s, p), nn.BatchNorm2D(out_c), nn.ReLU()
+    )
+
+
+def decoder_layer(
+    in_c, out_c, k=3, s=1, p=1, mode="nearest", scale_factor=None, size=None
+):
+    align_corners = False if mode == "nearest" else True
+    return nn.Sequential(
+        nn.Upsample(
+            size=size, scale_factor=scale_factor, mode=mode, align_corners=align_corners
+        ),
+        nn.Conv2D(in_c, out_c, k, s, p),
+        nn.BatchNorm2D(out_c),
+        nn.ReLU(),
+    )
+
+
+class PositionAttention(nn.Layer):
+    def __init__(
+        self,
+        max_length,
+        in_channels=512,
+        num_channels=64,
+        h=8,
+        w=32,
+        mode="nearest",
+        **kwargs,
+    ):
+        super().__init__()
+        self.max_length = max_length
+        self.k_encoder = nn.Sequential(
+            encoder_layer(in_channels, num_channels, s=(1, 2)),
+            encoder_layer(num_channels, num_channels, s=(2, 2)),
+            encoder_layer(num_channels, num_channels, s=(2, 2)),
+            encoder_layer(num_channels, num_channels, s=(2, 2)),
+        )
+        self.k_decoder = nn.Sequential(
+            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
+            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
+            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
+            decoder_layer(num_channels, in_channels, size=(h, w), mode=mode),
+        )
+
+        self.pos_encoder = PositionalEncoding(
+            dropout=0, dim=in_channels, max_len=max_length
+        )
+        self.project = nn.Linear(in_channels, in_channels)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        k, v = x, x
+
+        # calculate key vector
+        features = []
+        for i in range(0, len(self.k_encoder)):
+            k = self.k_encoder[i](k)
+            features.append(k)
+        for i in range(0, len(self.k_decoder) - 1):
+            k = self.k_decoder[i](k)
+            # print(k.shape, features[len(self.k_decoder) - 2 - i].shape)
+            k = k + features[len(self.k_decoder) - 2 - i]
+        k = self.k_decoder[-1](k)
+
+        # calculate query vector
+        # TODO q=f(q,k)
+        zeros = paddle.zeros((B, self.max_length, C), dtype=x.dtype)  # (B, N, C)
+        q = self.pos_encoder(zeros)  # (B, N, C)
+        q = self.project(q)  # (B, N, C)
+
+        # calculate attention
+        attn_scores = q @ k.flatten(2)  # (B, N, (H*W))
+        attn_scores = attn_scores / (C**0.5)
+        attn_scores = F.softmax(attn_scores, axis=-1)
+
+        v = v.flatten(2).transpose([0, 2, 1])  # (B, (H*W), C)
+        attn_vecs = attn_scores @ v  # (B, N, C)
+
+        return attn_vecs, attn_scores.reshape([0, self.max_length, H, W])
+
+
+class ABINetHead(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        d_model=512,
+        nhead=8,
+        num_layers=3,
+        dim_feedforward=2048,
+        dropout=0.1,
+        max_length=25,
+        use_lang=False,
+        iter_size=1,
+        image_size=(32, 128),
+    ):
+        super().__init__()
+        self.max_length = max_length + 1
+        h, w = image_size[0] // 4, image_size[1] // 4
+        self.pos_encoder = PositionalEncoding(dropout=0.1, dim=d_model, max_len=h * w)
+        self.encoder = nn.LayerList(
+            [
+                TransformerBlock(
+                    d_model=d_model,
+                    nhead=nhead,
+                    dim_feedforward=dim_feedforward,
+                    attention_dropout_rate=dropout,
+                    residual_dropout_rate=dropout,
+                    with_self_attn=True,
+                    with_cross_attn=False,
+                )
+                for i in range(num_layers)
+            ]
+        )
+        self.decoder = PositionAttention(
+            max_length=max_length + 1, mode="nearest", h=h, w=w  # additional stop token
+        )
+        self.out_channels = out_channels
+        self.cls = nn.Linear(d_model, self.out_channels)
+        self.use_lang = use_lang
+        if use_lang:
+            self.iter_size = iter_size
+            self.language = BCNLanguage(
+                d_model=d_model,
+                nhead=nhead,
+                num_layers=4,
+                dim_feedforward=dim_feedforward,
+                dropout=dropout,
+                max_length=max_length,
+                num_classes=self.out_channels,
+            )
+            # alignment
+            self.w_att_align = nn.Linear(2 * d_model, d_model)
+            self.cls_align = nn.Linear(d_model, self.out_channels)
+
+    def forward(self, x, targets=None):
+        x = x.transpose([0, 2, 3, 1])
+        _, H, W, C = x.shape
+        feature = x.flatten(1, 2)
+        feature = self.pos_encoder(feature)
+        for encoder_layer in self.encoder:
+            feature = encoder_layer(feature)
+        feature = feature.reshape([0, H, W, C]).transpose([0, 3, 1, 2])
+        v_feature, attn_scores = self.decoder(feature)  # (B, N, C), (B, C, H, W)
+        vis_logits = self.cls(v_feature)  # (B, N, C)
+        logits = vis_logits
+        vis_lengths = _get_length(vis_logits)
+        if self.use_lang:
+            align_logits = vis_logits
+            align_lengths = vis_lengths
+            all_l_res, all_a_res = [], []
+            for i in range(self.iter_size):
+                tokens = F.softmax(align_logits, axis=-1)
+                lengths = align_lengths
+                lengths = paddle.clip(
+                    lengths, 2, self.max_length
+                )  # TODO:move to language model
+                l_feature, l_logits = self.language(tokens, lengths)
+
+                # alignment
+                all_l_res.append(l_logits)
+                fuse = paddle.concat((l_feature, v_feature), -1)
+                f_att = F.sigmoid(self.w_att_align(fuse))
+                output = f_att * v_feature + (1 - f_att) * l_feature
+                align_logits = self.cls_align(output)  # (B, N, C)
+
+                align_lengths = _get_length(align_logits)
+                all_a_res.append(align_logits)
+            if self.training:
+                return {"align": all_a_res, "lang": all_l_res, "vision": vis_logits}
+            else:
+                logits = align_logits
+        if self.training:
+            return logits
+        else:
+            return F.softmax(logits, -1)
+
+
+def _get_length(logit):
+    """Greed decoder to obtain length from logit"""
+    out = logit.argmax(-1) == 0
+    abn = out.any(-1)
+    out_int = out.cast("int32")
+    out = (out_int.cumsum(-1) == 1) & out
+    out = out.cast("int32")
+    out = out.argmax(-1)
+    out = out + 1
+    len_seq = paddle.zeros_like(out) + logit.shape[1]
+    out = paddle.where(abn, out, len_seq)
+    return out
+
+
+def _get_mask(length, max_length):
+    """Generate a square mask for the sequence. The masked positions are filled with float('-inf').
+    Unmasked positions are filled with float(0.0).
+    """
+    length = length.unsqueeze(-1)
+    B = length.shape[0]
+    grid = paddle.arange(0, max_length).unsqueeze(0).tile([B, 1])
+    zero_mask = paddle.zeros([B, max_length], dtype="float32")
+    inf_mask = paddle.full([B, max_length], "-inf", dtype="float32")
+    diag_mask = paddle.diag(
+        paddle.full([max_length], "-inf", dtype=paddle.float32), offset=0, name=None
+    )
+    mask = paddle.where(grid >= length, inf_mask, zero_mask)
+    mask = mask.unsqueeze(1) + diag_mask
+    return mask.unsqueeze(1)