yomitoku 0.4.0.post1.dev0__py3-none-any.whl

yomitoku/models/dbnet_plus.py

@@ -0,0 +1,272 @@
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+from huggingface_hub import PyTorchModelHubMixin
+from torchvision.models._utils import IntermediateLayerGetter
+
+from .layers.dbnet_feature_attention import ScaleFeatureSelection
+
+
+class BackboneBase(nn.Module):
+    def __init__(self, backbone: nn.Module):
+        super().__init__()
+        return_layers = {
+            "layer1": "layer1",
+            "layer2": "layer2",
+            "layer3": "layer3",
+            "layer4": "layer4",
+        }
+
+        self.body = IntermediateLayerGetter(
+            backbone, return_layers=return_layers
+        )
+
+    def forward(self, tensor):
+        xs = self.body(tensor)
+        return xs
+
+
+class Backbone(BackboneBase):
+    """ResNet backbone with frozen BatchNorm."""
+
+    def __init__(self, name="resnet50", dilation=True):
+        backbone = getattr(torchvision.models, name)(
+            replace_stride_with_dilation=[False, False, dilation],
+            pretrained=False,
+        )
+        super().__init__(backbone)
+
+
+class DBNetDecoder(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        hidden_dim,
+        adaptive=False,
+        serial=False,
+        smooth=False,
+        k=50,
+    ):
+        super().__init__()
+        self.d_model = hidden_dim
+        self.n_layers = len(in_channels)
+        self.k = k
+        self.training = True
+        self.input_proj = nn.ModuleDict(
+            {
+                "layer1": nn.Conv2d(
+                    in_channels[0], self.d_model, 1, bias=False
+                ),
+                "layer2": nn.Conv2d(
+                    in_channels[1], self.d_model, 1, bias=False
+                ),
+                "layer3": nn.Conv2d(
+                    in_channels[2], self.d_model, 1, bias=False
+                ),
+                "layer4": nn.Conv2d(
+                    in_channels[3], self.d_model, 1, bias=False
+                ),
+            }
+        )
+
+        self.upsample_2x = nn.Upsample(
+            scale_factor=2, mode="bilinear", align_corners=False
+        )
+
+        self.out_proj = nn.ModuleDict(
+            {
+                "layer1": nn.Conv2d(
+                    self.d_model, self.d_model // 4, 3, padding=1, bias=False
+                ),
+                "layer2": nn.Sequential(
+                    nn.Conv2d(
+                        self.d_model,
+                        self.d_model // 4,
+                        3,
+                        padding=1,
+                        bias=False,
+                    ),
+                    nn.Upsample(
+                        scale_factor=2, mode="bilinear", align_corners=False
+                    ),
+                ),
+                "layer3": nn.Sequential(
+                    nn.Conv2d(
+                        self.d_model,
+                        self.d_model // 4,
+                        3,
+                        padding=1,
+                        bias=False,
+                    ),
+                    nn.Upsample(
+                        scale_factor=4, mode="bilinear", align_corners=False
+                    ),
+                ),
+                "layer4": nn.Sequential(
+                    nn.Conv2d(
+                        self.d_model,
+                        self.d_model // 4,
+                        3,
+                        padding=1,
+                        bias=False,
+                    ),
+                    nn.Upsample(
+                        scale_factor=4, mode="bilinear", align_corners=False
+                    ),
+                ),
+            }
+        )
+
+        self.binarize = nn.Sequential(
+            nn.Conv2d(
+                self.d_model, self.d_model // 4, 3, padding=1, bias=False
+            ),
+            nn.BatchNorm2d(self.d_model // 4),
+            nn.ReLU(inplace=True),
+            nn.ConvTranspose2d(self.d_model // 4, self.d_model // 4, 2, 2),
+            nn.BatchNorm2d(self.d_model // 4),
+            nn.ReLU(inplace=True),
+            nn.ConvTranspose2d(self.d_model // 4, 1, 2, 2),
+            nn.Sigmoid(),
+        )
+
+        self.adaptive = adaptive
+        self.serial = serial
+        if self.adaptive:
+            self.thresh = self._init_thresh(
+                self.d_model,
+                serial=serial,
+                smooth=smooth,
+                bias=False,
+            )
+            self.thresh.apply(self.weights_init)
+
+        self.binarize.apply(self.weights_init)
+
+        for layer in self.input_proj.values():
+            layer.apply(self.weights_init)
+
+        for layer in self.out_proj.values():
+            layer.apply(self.weights_init)
+
+        self.concat_attention = ScaleFeatureSelection(
+            self.d_model,
+            self.d_model // 4,
+            attention_type="scale_channel_spatial",
+        )
+
+    def weights_init(self, m):
+        classname = m.__class__.__name__
+        if classname.find("Conv") != -1:
+            nn.init.kaiming_normal_(m.weight.data)
+        elif classname.find("BatchNorm") != -1:
+            m.weight.data.fill_(1.0)
+            m.bias.data.fill_(1e-4)
+
+    def _init_thresh(
+        self, inner_channels, serial=False, smooth=False, bias=False
+    ):
+        in_channels = inner_channels
+        if serial:
+            in_channels += 1
+        self.thresh = nn.Sequential(
+            nn.Conv2d(
+                in_channels, inner_channels // 4, 3, padding=1, bias=bias
+            ),
+            nn.BatchNorm2d(inner_channels // 4),
+            nn.ReLU(inplace=True),
+            self._init_upsample(
+                inner_channels // 4,
+                inner_channels // 4,
+                smooth=smooth,
+                bias=bias,
+            ),
+            nn.BatchNorm2d(inner_channels // 4),
+            nn.ReLU(inplace=True),
+            self._init_upsample(
+                inner_channels // 4, 1, smooth=smooth, bias=bias
+            ),
+            nn.Sigmoid(),
+        )
+        return self.thresh
+
+    def _init_upsample(
+        self, in_channels, out_channels, smooth=False, bias=False
+    ):
+        if smooth:
+            inter_out_channels = out_channels
+            if out_channels == 1:
+                inter_out_channels = in_channels
+            module_list = [
+                nn.Upsample(scale_factor=2, mode="nearest"),
+                nn.Conv2d(in_channels, inter_out_channels, 3, 1, 1, bias=bias),
+            ]
+            if out_channels == 1:
+                module_list.append(
+                    nn.Conv2d(
+                        in_channels,
+                        out_channels,
+                        kernel_size=1,
+                        stride=1,
+                        padding=0,
+                        bias=False,
+                    )
+                )
+
+            return nn.Sequential(*module_list)
+        else:
+            return nn.ConvTranspose2d(in_channels, out_channels, 2, 2)
+
+    def step_function(self, x, y):
+        return torch.reciprocal(1 + torch.exp(-self.k * (x - y)))
+
+    def forward(self, features):
+        for layer, feature in features.items():
+            features[layer] = self.input_proj[layer](feature)
+
+        layers = ["layer4", "layer3", "layer2", "layer1"]
+        for i in range(self.n_layers - 1):
+            feature_bottom = features[layers[i]]
+            feature_top = features[layers[i + 1]]
+
+            bh, bw = feature_bottom.shape[-2:]
+            th, tw = feature_top.shape[-2:]
+
+            if bh != th or bw != tw:
+                feature_bottom = F.interpolate(
+                    feature_bottom,
+                    size=(th, tw),
+                    mode="bilinear",
+                    align_corners=False,
+                )
+
+            features[layers[i + 1]] = feature_bottom + feature_top
+
+        fp = []
+        for layer, feature in features.items():
+            fp.append(self.out_proj[layer](feature))
+        fuse = torch.cat(fp[::-1], dim=1)
+        fuse = self.concat_attention(fuse, fp[::-1])
+
+        binary = self.binarize(fuse)
+        result = OrderedDict(binary=binary)
+        return result
+
+
+class DBNet(nn.Module, PyTorchModelHubMixin):
+    def __init__(
+        self,
+        cfg,
+    ):
+        super().__init__()
+        self.cfg = cfg
+        self.backbone = Backbone(**cfg.backbone)
+        self.decoder = DBNetDecoder(**cfg.decoder)
+
+    def forward(self, tensor):
+        features = self.backbone(tensor)
+        xs = self.decoder(features)
+        return xs

yomitoku/models/layers/activate.py

@@ -0,0 +1,38 @@
+import torch.nn as nn
+
+
+def get_activation(act: str, inplace: bool = True):
+    """get activation"""
+    if act is None:
+        return nn.Identity()
+
+    elif isinstance(act, nn.Module):
+        return act
+
+    act = act.lower()
+
+    if act == "silu" or act == "swish":
+        m = nn.SiLU()
+
+    elif act == "relu":
+        m = nn.ReLU()
+
+    elif act == "leaky_relu":
+        m = nn.LeakyReLU()
+
+    elif act == "silu":
+        m = nn.SiLU()
+
+    elif act == "gelu":
+        m = nn.GELU()
+
+    elif act == "hardsigmoid":
+        m = nn.Hardsigmoid()
+
+    else:
+        raise RuntimeError("")
+
+    if hasattr(m, "inplace"):
+        m.inplace = inplace
+
+    return m

yomitoku/models/layers/dbnet_feature_attention.py

@@ -0,0 +1,160 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ScaleChannelAttention(nn.Module):
+    def __init__(self, in_planes, out_planes, num_features, init_weight=True):
+        super(ScaleChannelAttention, self).__init__()
+        self.avgpool = nn.AdaptiveAvgPool2d(1)
+        print(self.avgpool)
+        self.fc1 = nn.Conv2d(in_planes, out_planes, 1, bias=False)
+        self.bn = nn.BatchNorm2d(out_planes)
+        self.fc2 = nn.Conv2d(out_planes, num_features, 1, bias=False)
+        if init_weight:
+            self._initialize_weights()
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            if isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        global_x = self.avgpool(x)
+        global_x = self.fc1(global_x)
+        global_x = F.relu(self.bn(global_x))
+        global_x = self.fc2(global_x)
+        global_x = F.softmax(global_x, 1)
+        return global_x
+
+
+class ScaleChannelSpatialAttention(nn.Module):
+    def __init__(self, in_planes, out_planes, num_features, init_weight=True):
+        super(ScaleChannelSpatialAttention, self).__init__()
+        self.channel_wise = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(in_planes, out_planes, 1, bias=False),
+            # nn.BatchNorm2d(out_planes),
+            nn.ReLU(),
+            nn.Conv2d(out_planes, in_planes, 1, bias=False),
+        )
+        self.spatial_wise = nn.Sequential(
+            # Nx1xHxW
+            nn.Conv2d(1, 1, 3, bias=False, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(1, 1, 1, bias=False),
+            nn.Sigmoid(),
+        )
+        self.attention_wise = nn.Sequential(
+            nn.Conv2d(in_planes, num_features, 1, bias=False), nn.Sigmoid()
+        )
+        if init_weight:
+            self._initialize_weights()
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            if isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        # global_x = self.avgpool(x)
+        # shape Nx4x1x1
+        global_x = self.channel_wise(x).sigmoid()
+        # shape: NxCxHxW
+        global_x = global_x + x
+        # shape:Nx1xHxW
+        x = torch.mean(global_x, dim=1, keepdim=True)
+        global_x = self.spatial_wise(x) + global_x
+        global_x = self.attention_wise(global_x)
+        return global_x
+
+
+class ScaleSpatialAttention(nn.Module):
+    def __init__(self, in_planes, out_planes, num_features, init_weight=True):
+        super(ScaleSpatialAttention, self).__init__()
+        self.spatial_wise = nn.Sequential(
+            # Nx1xHxW
+            nn.Conv2d(1, 1, 3, bias=False, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(1, 1, 1, bias=False),
+            nn.Sigmoid(),
+        )
+        self.attention_wise = nn.Sequential(
+            nn.Conv2d(in_planes, num_features, 1, bias=False), nn.Sigmoid()
+        )
+        if init_weight:
+            self._initialize_weights()
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            if isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        global_x = torch.mean(x, dim=1, keepdim=True)
+        global_x = self.spatial_wise(global_x) + x
+        global_x = self.attention_wise(global_x)
+        return global_x
+
+
+class ScaleFeatureSelection(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        inter_channels,
+        out_features_num=4,
+        attention_type="scale_spatial",
+    ):
+        super(ScaleFeatureSelection, self).__init__()
+        self.in_channels = in_channels
+        self.inter_channels = inter_channels
+        self.out_features_num = out_features_num
+        self.conv = nn.Conv2d(in_channels, inter_channels, 3, padding=1)
+        self.type = attention_type
+        if self.type == "scale_spatial":
+            self.enhanced_attention = ScaleSpatialAttention(
+                inter_channels, inter_channels // 4, out_features_num
+            )
+        elif self.type == "scale_channel_spatial":
+            self.enhanced_attention = ScaleChannelSpatialAttention(
+                inter_channels, inter_channels // 4, out_features_num
+            )
+        elif self.type == "scale_channel":
+            self.enhanced_attention = ScaleChannelAttention(
+                inter_channels, inter_channels // 2, out_features_num
+            )
+
+    def _initialize_weights(self, m):
+        classname = m.__class__.__name__
+        if classname.find("Conv") != -1:
+            nn.init.kaiming_normal_(m.weight.data)
+        elif classname.find("BatchNorm") != -1:
+            m.weight.data.fill_(1.0)
+            m.bias.data.fill_(1e-4)
+
+    def forward(self, concat_x, features_list):
+        concat_x = self.conv(concat_x)
+        score = self.enhanced_attention(concat_x)
+        assert len(features_list) == self.out_features_num
+        if self.type not in ["scale_channel_spatial", "scale_spatial"]:
+            shape = features_list[0].shape[2:]
+            score = F.interpolate(score, size=shape, mode="bilinear")
+        x = []
+        for i in range(self.out_features_num):
+            x.append(score[:, i : i + 1] * features_list[i])
+        return torch.cat(x, dim=1)

yomitoku/models/layers/parseq_transformer.py

@@ -0,0 +1,218 @@
+# Scene Text Recognition Model Hub
+# Copyright 2022 Darwin Bautista
+#
+# 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
+#
+#     https://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.
+
+import math
+from typing import Optional
+
+import torch
+from timm.models.vision_transformer import PatchEmbed, VisionTransformer
+from torch import Tensor
+from torch import nn as nn
+from torch.nn import functional as F
+from torch.nn.modules import transformer
+
+
+class DecoderLayer(nn.Module):
+    """A Transformer decoder layer supporting two-stream attention (XLNet)
+    This implements a pre-LN decoder, as opposed to the post-LN default in PyTorch.
+    """
+
+    def __init__(
+        self,
+        embed_dim,
+        num_heads,
+        mlp_ratio=2048,
+        dropout=0.1,
+        activation="gelu",
+        layer_norm_eps=1e-5,
+    ):
+        super().__init__()
+        dim_feedforward = embed_dim * mlp_ratio
+        self.self_attn = nn.MultiheadAttention(
+            embed_dim, num_heads, dropout=dropout, batch_first=True
+        )
+        self.cross_attn = nn.MultiheadAttention(
+            embed_dim, num_heads, dropout=dropout, batch_first=True
+        )
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(embed_dim, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, embed_dim)
+
+        self.norm1 = nn.LayerNorm(embed_dim, eps=layer_norm_eps)
+        self.norm2 = nn.LayerNorm(embed_dim, eps=layer_norm_eps)
+        self.norm_q = nn.LayerNorm(embed_dim, eps=layer_norm_eps)
+        self.norm_c = nn.LayerNorm(embed_dim, eps=layer_norm_eps)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+        self.activation = transformer._get_activation_fn(activation)
+
+    def __setstate__(self, state):
+        if "activation" not in state:
+            state["activation"] = F.gelu
+        super().__setstate__(state)
+
+    def forward_stream(
+        self,
+        tgt: Tensor,
+        tgt_norm: Tensor,
+        tgt_kv: Tensor,
+        memory: Tensor,
+        tgt_mask: Optional[Tensor],
+        tgt_key_padding_mask: Optional[Tensor],
+    ):
+        """Forward pass for a single stream (i.e. content or query)
+        tgt_norm is just a LayerNorm'd tgt. Added as a separate parameter for efficiency.
+        Both tgt_kv and memory are expected to be LayerNorm'd too.
+        memory is LayerNorm'd by ViT.
+        """
+        tgt2, sa_weights = self.self_attn(
+            tgt_norm,
+            tgt_kv,
+            tgt_kv,
+            attn_mask=tgt_mask,
+            key_padding_mask=tgt_key_padding_mask,
+        )
+        tgt = tgt + self.dropout1(tgt2)
+
+        tgt2, ca_weights = self.cross_attn(self.norm1(tgt), memory, memory)
+        tgt = tgt + self.dropout2(tgt2)
+
+        tgt2 = self.linear2(
+            self.dropout(self.activation(self.linear1(self.norm2(tgt))))
+        )
+        tgt = tgt + self.dropout3(tgt2)
+        return tgt, sa_weights, ca_weights
+
+    def forward(
+        self,
+        query,
+        content,
+        memory,
+        query_mask: Optional[Tensor] = None,
+        content_mask: Optional[Tensor] = None,
+        content_key_padding_mask: Optional[Tensor] = None,
+        update_content: bool = True,
+    ):
+        query_norm = self.norm_q(query)
+        content_norm = self.norm_c(content)
+        query = self.forward_stream(
+            query,
+            query_norm,
+            content_norm,
+            memory,
+            query_mask,
+            content_key_padding_mask,
+        )[0]
+        if update_content:
+            content = self.forward_stream(
+                content,
+                content_norm,
+                content_norm,
+                memory,
+                content_mask,
+                content_key_padding_mask,
+            )[0]
+        return query, content
+
+
+class Decoder(nn.Module):
+    __constants__ = ["norm"]
+
+    def __init__(self, norm, cfg):
+        super().__init__()
+        decoder_layer = DecoderLayer(
+            embed_dim=cfg.embed_dim,
+            num_heads=cfg.num_heads,
+            mlp_ratio=cfg.mlp_ratio,
+        )
+
+        self.layers = transformer._get_clones(decoder_layer, cfg.depth)
+        self.num_layers = cfg.depth
+        self.norm = norm
+
+    def forward(
+        self,
+        query,
+        content,
+        memory,
+        query_mask: Optional[Tensor] = None,
+        content_mask: Optional[Tensor] = None,
+        content_key_padding_mask: Optional[Tensor] = None,
+    ):
+        for i, mod in enumerate(self.layers):
+            last = i == len(self.layers) - 1
+            query, content = mod(
+                query,
+                content,
+                memory,
+                query_mask,
+                content_mask,
+                content_key_padding_mask,
+                update_content=not last,
+            )
+        query = self.norm(query)
+        return query
+
+
+class Encoder(VisionTransformer):
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.0,
+        embed_layer=PatchEmbed,
+    ):
+        super().__init__(
+            img_size,
+            patch_size,
+            in_chans,
+            embed_dim=embed_dim,
+            depth=depth,
+            num_heads=num_heads,
+            mlp_ratio=mlp_ratio,
+            qkv_bias=qkv_bias,
+            drop_rate=drop_rate,
+            attn_drop_rate=attn_drop_rate,
+            drop_path_rate=drop_path_rate,
+            embed_layer=embed_layer,
+            num_classes=0,  # These
+            global_pool="",  # disable the
+            class_token=False,  # classifier head.
+        )
+
+    def forward(self, x):
+        # Return all tokens
+        return self.forward_features(x)
+
+
+class TokenEmbedding(nn.Module):
+    def __init__(self, charset_size: int, embed_dim: int):
+        super().__init__()
+        self.embedding = nn.Embedding(charset_size, embed_dim)
+        self.embed_dim = embed_dim
+
+    def forward(self, tokens: torch.Tensor):
+        return math.sqrt(self.embed_dim) * self.embedding(tokens)