add cafpn

ppocr/modeling/necks/__init__.py

@@ -16,7 +16,7 @@ __all__ = ['build_neck']
 
 
 def build_neck(config):
-    from .db_fpn import DBFPN
+    from .db_fpn import DBFPN, CAFPN
     from .east_fpn import EASTFPN
     from .sast_fpn import SASTFPN
     from .rnn import SequenceEncoder
@@ -26,8 +26,8 @@ def build_neck(config):
     from .fce_fpn import FCEFPN
     from .pren_fpn import PRENFPN
     support_dict = [
-        'FPN', 'FCEFPN', 'DBFPN', 'EASTFPN', 'SASTFPN', 'SequenceEncoder',
-        'PGFPN', 'TableFPN', 'PRENFPN'
+        'FPN', 'FCEFPN', 'DBFPN', 'CAFPN', 'EASTFPN', 'SASTFPN',
+        'SequenceEncoder', 'PGFPN', 'TableFPN', 'PRENFPN'
     ]
 
     module_name = config.pop('name')

ppocr/modeling/necks/db_fpn.py

@@ -20,6 +20,7 @@ import paddle
 from paddle import nn
 import paddle.nn.functional as F
 from paddle import ParamAttr
+from ppocr.backbones.det_mobilenet_v3 import SEModule
 
 
 class DBFPN(nn.Layer):
@@ -106,3 +107,75 @@ class DBFPN(nn.Layer):
 
         fuse = paddle.concat([p5, p4, p3, p2], axis=1)
         return fuse
+
+
+class CALayer(nn.Layer):
+    def __init__(self, in_channels, out_channels, kernel_size, shortcut=True):
+        super(CALayer, self).__init__()
+        weight_attr = paddle.nn.initializer.KaimingUniform()
+        self.in_conv = nn.Conv2D(
+            in_channels=in_channels,
+            out_channels=self.out_channels,
+            kernel_size=kernel_size,
+            padding=int(kernel_size // 2),
+            weight_attr=ParamAttr(initializer=weight_attr),
+            bias_attr=False)
+        self.se_block = SEModule(self.out_channels)
+        self.shortcut = shortcut
+
+    def forward(self, ins):
+        x = self.in_conv(ins)
+        if self.shortcut:
+            out = x + self.se_block(x)
+        else:
+            out = self.se_block(x)
+        return out
+
+
+class CAFPN(nn.Layer):
+    def __init__(self, in_channels, out_channels, shortcut, **kwargs):
+        super(CAFPN, self).__init__()
+
+        self.ins_convs = []
+        self.inp_convs = []
+
+        for i in range(len(in_channels)):
+            self.ins_conv.append(
+                CALayer(
+                    in_channels[i],
+                    out_channels,
+                    kernel_size=1,
+                    shortcut=shortcut))
+            self.inp_conv.append(
+                CALayer(
+                    out_channels,
+                    out_channels // 4,
+                    kernel_size=3,
+                    shortcut=shortcut))
+
+    def forward(self, x):
+        c2, c3, c4, c5 = x
+
+        in5 = self.ins_conv[3](c5)
+        in4 = self.ins_conv[2](c4)
+        in3 = self.ins_conv[1](c3)
+        in2 = self.ins_conv[0](c2)
+
+        out4 = in4 + F.upsample(
+            in5, scale_factor=2, mode="nearest", align_mode=1)  # 1/16
+        out3 = in3 + F.upsample(
+            out4, scale_factor=2, mode="nearest", align_mode=1)  # 1/8
+        out2 = in2 + F.upsample(
+            out3, scale_factor=2, mode="nearest", align_mode=1)  # 1/4
+
+        p5 = self.inp_conv[3](in5)
+        p4 = self.inp_conv[2](out4)
+        p3 = self.inp_conv[1](out3)
+        p2 = self.inp_conv[0](out2)
+
+        p5 = F.upsample(p5, scale_factor=8, mode="nearest", align_mode=1)
+        p4 = F.upsample(p4, scale_factor=4, mode="nearest", align_mode=1)
+        p3 = F.upsample(p3, scale_factor=2, mode="nearest", align_mode=1)
+
+        fuse = paddle.concat([p5, p4, p3, p2], axis=1)
+        return fuse