1 files changed, 92 insertions, 3 deletions
diff --git a/text_recognizer/networks/encoders/efficientnet/efficientnet.py b/text_recognizer/networks/encoders/efficientnet/efficientnet.py
index d953c10..98d58fd 100644
--- a/text_recognizer/networks/encoders/efficientnet/efficientnet.py
+++ b/text_recognizer/networks/encoders/efficientnet/efficientnet.py
@@ -1,9 +1,98 @@
 """Efficient net."""
+from typing import Tuple
+
 from torch import nn, Tensor
 
+from .mbconv import MBConvBlock
+from .utils import (
+    block_args,
+    calculate_output_image_size,
+    get_same_padding_conv2d,
+    round_filters,
+    round_repeats,
+)
+
 
 class EfficientNet(nn.Module):
-    def __init__(
-        self,
-    ) -> None:
+    archs = {
+        #     width,depth0res,dropout
+        "b0": (1.0, 1.0, 0.2),
+        "b1": (1.0, 1.1, 0.2),
+        "b2": (1.1, 1.2, 0.3),
+        "b3": (1.2, 1.4, 0.3),
+        "b4": (1.4, 1.8, 0.4),
+        "b5": (1.6, 2.2, 0.4),
+        "b6": (1.8, 2.6, 0.5),
+        "b7": (2.0, 3.1, 0.5),
+        "b8": (2.2, 3.6, 0.5),
+        "l2": (4.3, 5.3, 0.5),
+    }
+
+    def __init__(self, arch: str, image_size: Tuple[int, int]) -> None:
         super().__init__()
+        assert arch in self.archs, f"{arch} not a valid efficient net architecure!"
+        self.arch = self.archs[arch]
+        self.image_size = image_size
+        self._conv_stem: nn.Sequential = None
+        self._blocks: nn.Sequential = None
+        self._conv_head: nn.Sequential = None
+        self._build()
+
+    def _build(self) -> None:
+        _block_args = block_args()
+        in_channels = 1  # BW
+        out_channels = round_filters(32, self.arch)
+        Conv2d = get_same_padding_conv2d(image_size=self.image_size)
+        self._conv_stem = nn.Sequential(
+            Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=3,
+                stride=2,
+                bias=False,
+            ),
+            nn.BatchNorm2d(num_features=out_channels, momentum=bn_momentum, eps=bn_eps),
+            nn.Mish(inplace=True),
+        )
+        image_size = calculate_output_image_size(self.image_size, 2)
+        self._blocks = nn.ModuleList([])
+        for args in _block_args:
+            args.in_channels = round_filters(args.in_channels, self.arch)
+            args.out_channels = round_filters(args.out_channels, self.arch)
+            args.num_repeat = round_repeats(args.num_repeat, self.arch)
+
+            self._blocks.append(
+                MBConvBlock(
+                    **args,
+                    bn_momentum=bn_momentum,
+                    bn_eps=bn_eps,
+                    image_size=image_size,
+                )
+            )
+            image_size = calculate_output_image_size(image_size, args.stride)
+            if args.num_repeat > 1:
+                args.in_channels = args.out_channels
+                args.stride = 1
+            for _ in range(args.num_repeat - 1):
+                self._blocks.append(
+                    MBConvBlock(
+                        **args,
+                        bn_momentum=bn_momentum,
+                        bn_eps=bn_eps,
+                        image_size=image_size,
+                    )
+                )
+
+        in_channels = args.out_channels
+        out_channels = round_filters(1280, self.arch)
+        Conv2d = get_same_padding_conv2d(image_size=image_size)
+        self._conv_head = nn.Sequential(
+            Conv2d(in_channels, out_channels, kernel_size=1, bias=False),
+            nn.BatchNorm2d(num_features=out_channels, momentum=bn_momentum, eps=bn_eps),
+        )
+
+    def extract_features(self, x: Tensor) -> Tensor:
+        x = self._conv_stem(x)
+
+    def forward(self, x: Tensor) -> Tensor:
+        pass