refine model.py

Author: Pandede

wpodnet/model.py

@@ -1,27 +1,28 @@
 import torch
-import torch.nn as nn
 
 
-class BasicConvBlock(nn.Module):
+class BasicConvBlock(torch.nn.Module):
     def __init__(self, in_channels: int, out_channels: int):
         super().__init__()
-        self.conv_layer = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
-        self.bn_layer = nn.BatchNorm2d(out_channels, momentum=0.99, eps=0.001)
-        self.act_layer = nn.ReLU(inplace=True)
+        self.conv_layer = torch.nn.Conv2d(
+            in_channels, out_channels, kernel_size=3, padding=1
+        )
+        self.bn_layer = torch.nn.BatchNorm2d(out_channels, momentum=0.99, eps=0.001)
+        self.act_layer = torch.nn.ReLU(inplace=True)
 
     def forward(self, x):
         x = self.conv_layer(x)
         x = self.bn_layer(x)
         return self.act_layer(x)
 
 
-class ResBlock(nn.Module):
+class ResBlock(torch.nn.Module):
     def __init__(self, channels: int):
         super().__init__()
         self.conv_block = BasicConvBlock(channels, channels)
-        self.sec_layer = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
-        self.bn_layer = nn.BatchNorm2d(channels, momentum=0.99, eps=0.001)
-        self.act_layer = nn.ReLU(inplace=True)
+        self.sec_layer = torch.nn.Conv2d(channels, channels, kernel_size=3, padding=1)
+        self.bn_layer = torch.nn.BatchNorm2d(channels, momentum=0.99, eps=0.001)
+        self.act_layer = torch.nn.ReLU(inplace=True)
 
     def forward(self, x):
         h = self.conv_block(x)
@@ -30,35 +31,45 @@ def forward(self, x):
         return self.act_layer(x + h)
 
 
-class WPODNet(nn.Module):
+class WPODNet(torch.nn.Module):
+    """
+    WPODNet in PyTorch.
+
+    The original architecture is built in Keras: https://github.com/sergiomsilva/alpr-unconstrained/blob/master/create-model.py
+    """
+
+    # https://github.com/sergiomsilva/alpr-unconstrained/blob/master/src/keras_utils.py#L43-L44
+    stride = 16  # net_stride
+    scale_factor = 7.75  # side
+
     def __init__(self):
         super().__init__()
-        self.backbone = nn.Sequential(
+        self.backbone = torch.nn.Sequential(
             BasicConvBlock(3, 16),
             BasicConvBlock(16, 16),
-            nn.MaxPool2d(2),
+            torch.nn.MaxPool2d(2),
             BasicConvBlock(16, 32),
             ResBlock(32),
-            nn.MaxPool2d(2),
+            torch.nn.MaxPool2d(2),
             BasicConvBlock(32, 64),
             ResBlock(64),
             ResBlock(64),
-            nn.MaxPool2d(2),
+            torch.nn.MaxPool2d(2),
             BasicConvBlock(64, 64),
             ResBlock(64),
             ResBlock(64),
-            nn.MaxPool2d(2),
+            torch.nn.MaxPool2d(2),
             BasicConvBlock(64, 128),
             ResBlock(128),
             ResBlock(128),
             ResBlock(128),
-            ResBlock(128)
+            ResBlock(128),
         )
-        self.prob_layer = nn.Conv2d(128, 2, kernel_size=3, padding=1)
-        self.bbox_layer = nn.Conv2d(128, 6, kernel_size=3, padding=1)
+        self.prob_layer = torch.nn.Conv2d(128, 2, kernel_size=3, padding=1)
+        self.bbox_layer = torch.nn.Conv2d(128, 6, kernel_size=3, padding=1)
 
         # Registry a dummy tensor for retrieve the attached device
-        self.register_buffer('dummy', torch.Tensor(), persistent=False)
+        self.register_buffer("dummy", torch.Tensor(), persistent=False)
 
     @property
     def device(self) -> torch.device: