Spaces:
Runtime error
Runtime error
File size: 6,423 Bytes
153628e |
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 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 |
# Copyright (C) 2021-2024, Mindee.
# This program is licensed under the Apache License 2.0.
# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
import math
from copy import deepcopy
from functools import partial
from typing import Any, Dict, List, Optional, Tuple
import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras.models import Sequential
from doctr.datasets import VOCABS
from ...utils import load_pretrained_params
from ..resnet.tensorflow import ResNet
__all__ = ["magc_resnet31"]
default_cfgs: Dict[str, Dict[str, Any]] = {
"magc_resnet31": {
"mean": (0.694, 0.695, 0.693),
"std": (0.299, 0.296, 0.301),
"input_shape": (32, 32, 3),
"classes": list(VOCABS["french"]),
"url": "https://doctr-static.mindee.com/models?id=v0.6.0/magc_resnet31-addbb705.zip&src=0",
},
}
class MAGC(layers.Layer):
"""Implements the Multi-Aspect Global Context Attention, as described in
<https://arxiv.org/pdf/1910.02562.pdf>`_.
Args:
----
inplanes: input channels
headers: number of headers to split channels
attn_scale: if True, re-scale attention to counteract the variance distibutions
ratio: bottleneck ratio
**kwargs
"""
def __init__(
self,
inplanes: int,
headers: int = 8,
attn_scale: bool = False,
ratio: float = 0.0625, # bottleneck ratio of 1/16 as described in paper
**kwargs,
) -> None:
super().__init__(**kwargs)
self.headers = headers # h
self.inplanes = inplanes # C
self.attn_scale = attn_scale
self.planes = int(inplanes * ratio)
self.single_header_inplanes = int(inplanes / headers) # C / h
self.conv_mask = layers.Conv2D(filters=1, kernel_size=1, kernel_initializer=tf.initializers.he_normal())
self.transform = Sequential(
[
layers.Conv2D(filters=self.planes, kernel_size=1, kernel_initializer=tf.initializers.he_normal()),
layers.LayerNormalization([1, 2, 3]),
layers.ReLU(),
layers.Conv2D(filters=self.inplanes, kernel_size=1, kernel_initializer=tf.initializers.he_normal()),
],
name="transform",
)
def context_modeling(self, inputs: tf.Tensor) -> tf.Tensor:
b, h, w, c = (tf.shape(inputs)[i] for i in range(4))
# B, H, W, C -->> B*h, H, W, C/h
x = tf.reshape(inputs, shape=(b, h, w, self.headers, self.single_header_inplanes))
x = tf.transpose(x, perm=(0, 3, 1, 2, 4))
x = tf.reshape(x, shape=(b * self.headers, h, w, self.single_header_inplanes))
# Compute shorcut
shortcut = x
# B*h, 1, H*W, C/h
shortcut = tf.reshape(shortcut, shape=(b * self.headers, 1, h * w, self.single_header_inplanes))
# B*h, 1, C/h, H*W
shortcut = tf.transpose(shortcut, perm=[0, 1, 3, 2])
# Compute context mask
# B*h, H, W, 1
context_mask = self.conv_mask(x)
# B*h, 1, H*W, 1
context_mask = tf.reshape(context_mask, shape=(b * self.headers, 1, h * w, 1))
# scale variance
if self.attn_scale and self.headers > 1:
context_mask = context_mask / math.sqrt(self.single_header_inplanes)
# B*h, 1, H*W, 1
context_mask = tf.keras.activations.softmax(context_mask, axis=2)
# Compute context
# B*h, 1, C/h, 1
context = tf.matmul(shortcut, context_mask)
context = tf.reshape(context, shape=(b, 1, c, 1))
# B, 1, 1, C
context = tf.transpose(context, perm=(0, 1, 3, 2))
# Set shape to resolve shape when calling this module in the Sequential MAGCResnet
batch, chan = inputs.get_shape().as_list()[0], inputs.get_shape().as_list()[-1]
context.set_shape([batch, 1, 1, chan])
return context
def call(self, inputs: tf.Tensor, **kwargs) -> tf.Tensor:
# Context modeling: B, H, W, C -> B, 1, 1, C
context = self.context_modeling(inputs)
# Transform: B, 1, 1, C -> B, 1, 1, C
transformed = self.transform(context)
return inputs + transformed
def _magc_resnet(
arch: str,
pretrained: bool,
num_blocks: List[int],
output_channels: List[int],
stage_downsample: List[bool],
stage_conv: List[bool],
stage_pooling: List[Optional[Tuple[int, int]]],
origin_stem: bool = True,
**kwargs: Any,
) -> ResNet:
kwargs["num_classes"] = kwargs.get("num_classes", len(default_cfgs[arch]["classes"]))
kwargs["input_shape"] = kwargs.get("input_shape", default_cfgs[arch]["input_shape"])
kwargs["classes"] = kwargs.get("classes", default_cfgs[arch]["classes"])
_cfg = deepcopy(default_cfgs[arch])
_cfg["num_classes"] = kwargs["num_classes"]
_cfg["classes"] = kwargs["classes"]
_cfg["input_shape"] = kwargs["input_shape"]
kwargs.pop("classes")
# Build the model
model = ResNet(
num_blocks,
output_channels,
stage_downsample,
stage_conv,
stage_pooling,
origin_stem,
attn_module=partial(MAGC, headers=8, attn_scale=True),
cfg=_cfg,
**kwargs,
)
# Load pretrained parameters
if pretrained:
load_pretrained_params(model, default_cfgs[arch]["url"])
return model
def magc_resnet31(pretrained: bool = False, **kwargs: Any) -> ResNet:
"""Resnet31 architecture with Multi-Aspect Global Context Attention as described in
`"MASTER: Multi-Aspect Non-local Network for Scene Text Recognition",
<https://arxiv.org/pdf/1910.02562.pdf>`_.
>>> import tensorflow as tf
>>> from doctr.models import magc_resnet31
>>> model = magc_resnet31(pretrained=False)
>>> input_tensor = tf.random.uniform(shape=[1, 224, 224, 3], maxval=1, dtype=tf.float32)
>>> out = model(input_tensor)
Args:
----
pretrained: boolean, True if model is pretrained
**kwargs: keyword arguments of the ResNet architecture
Returns:
-------
A feature extractor model
"""
return _magc_resnet(
"magc_resnet31",
pretrained,
[1, 2, 5, 3],
[256, 256, 512, 512],
[False] * 4,
[True] * 4,
[(2, 2), (2, 1), None, None],
False,
stem_channels=128,
**kwargs,
)
|