from __future__ import absolute_import
import torch
from torch.nn import functional

class FPN(torch.nn.Module):
    def __init__(self, out_channels):
        super(FPN, self).__init__()
        self.out_channels = out_channels

        self.P5 = torch.nn.MaxPool2d(kernel_size=1, stride=2, padding=0)

        self.P4_conv1 = torch.nn.Conv2d(512, self.out_channels, kernel_size=1, stride=1, padding=0)
        self.P4_conv2 = torch.nn.Conv2d(self.out_channels, self.out_channels, 3, 1, 1)

        self.P3_conv1 = torch.nn.Conv2d(512, self.out_channels, kernel_size=1, stride=1, padding=0)
        self.P3_conv2 = torch.nn.Conv2d(self.out_channels, self.out_channels, 3, 1, 1)

        self.P2_conv1 = torch.nn.Conv2d(256, self.out_channels, kernel_size=1, stride=1, padding=0)
        self.P2_conv2 = torch.nn.Conv2d(self.out_channels, self.out_channels, 3, 1, 1)

        normal_init(self.P4_conv1, 0, 0.01)
        normal_init(self.P4_conv2, 0, 0.01)
        normal_init(self.P3_conv1, 0, 0.01)
        normal_init(self.P3_conv2, 0, 0.01)
        normal_init(self.P2_conv1, 0, 0.01)
        normal_init(self.P2_conv2, 0, 0.01)


    def forward(self, C2, C3, C4):

        p4_out = self.P4_conv1(C4)

        p5_out = self.P5(p4_out)

        p3_out = self._upsample_add(p4_out, self.P3_conv1(C3))
        p2_out = self._upsample_add(p3_out, self.P2_conv1(C2))

        p4_out = self.P4_conv2(p4_out)
        p3_out = self.P3_conv2(p3_out)
        p2_out = self.P2_conv2(p2_out)

        return p2_out, p3_out, p4_out, p5_out

    def _upsample_add(self, x, y):
        '''Upsample and add two feature maps.
        Args:
          x: (Variable) top feature map to be upsampled.
          y: (Variable) lateral feature map.
        Returns:
          (Variable) added feature map.
        Note in PyTorch, when input size is odd, the upsampled feature map
        with `F.upsample(..., scale_factor=2, mode='nearest')`
        maybe not equal to the lateral feature map size.
        e.g.
        original input size: [N,_,15,15] ->
        conv2d feature map size: [N,_,8,8] ->
        upsampled feature map size: [N,_,16,16]
        So we choose bilinear upsample which supports arbitrary output sizes.
        '''
        _,_,H,W = y.size()
        return functional.interpolate(x, size=(H,W), mode='bilinear') + y

def normal_init(m, mean, stddev, truncated=False):
    """
    weight initalizer: truncated normal and random normal.
    """
    # x is a parameter
    if truncated:
        m.weight.data.normal_().fmod_(2).mul_(stddev).add_(mean)  # not a perfect approximation
    else:
        m.weight.data.normal_(mean, stddev)
        m.bias.data.zero_()
[((247, 301), 'torch.nn.MaxPool2d', 'torch.nn.MaxPool2d', ([], {'kernel_size': '(1)', 'stride': '(2)', 'padding': '(0)'}), '(kernel_size=1, stride=2, padding=0)\n', (265, 301), False, 'import torch\n'), ((327, 402), 'torch.nn.Conv2d', 'torch.nn.Conv2d', (['(512)', 'self.out_channels'], {'kernel_size': '(1)', 'stride': '(1)', 'padding': '(0)'}), '(512, self.out_channels, kernel_size=1, stride=1, padding=0)\n', (342, 402), False, 'import torch\n'), ((427, 489), 'torch.nn.Conv2d', 'torch.nn.Conv2d', (['self.out_channels', 'self.out_channels', '(3)', '(1)', '(1)'], {}), '(self.out_channels, self.out_channels, 3, 1, 1)\n', (442, 489), False, 'import torch\n'), ((515, 590), 'torch.nn.Conv2d', 'torch.nn.Conv2d', (['(512)', 'self.out_channels'], {'kernel_size': '(1)', 'stride': '(1)', 'padding': '(0)'}), '(512, self.out_channels, kernel_size=1, stride=1, padding=0)\n', (530, 590), False, 'import torch\n'), ((615, 677), 'torch.nn.Conv2d', 'torch.nn.Conv2d', (['self.out_channels', 'self.out_channels', '(3)', '(1)', '(1)'], {}), '(self.out_channels, self.out_channels, 3, 1, 1)\n', (630, 677), False, 'import torch\n'), ((703, 778), 'torch.nn.Conv2d', 'torch.nn.Conv2d', (['(256)', 'self.out_channels'], {'kernel_size': '(1)', 'stride': '(1)', 'padding': '(0)'}), '(256, self.out_channels, kernel_size=1, stride=1, padding=0)\n', (718, 778), False, 'import torch\n'), ((803, 865), 'torch.nn.Conv2d', 'torch.nn.Conv2d', (['self.out_channels', 'self.out_channels', '(3)', '(1)', '(1)'], {}), '(self.out_channels, self.out_channels, 3, 1, 1)\n', (818, 865), False, 'import torch\n'), ((2263, 2318), 'torch.nn.functional.interpolate', 'functional.interpolate', (['x'], {'size': '(H, W)', 'mode': '"""bilinear"""'}), "(x, size=(H, W), mode='bilinear')\n", (2285, 2318), False, 'from torch.nn import functional\n')]