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_()
[((10, 18, 10, 72), 'torch.nn.MaxPool2d', 'torch.nn.MaxPool2d', (), '', False, 'import torch\n'), ((12, 24, 12, 99), 'torch.nn.Conv2d', 'torch.nn.Conv2d', (), '', False, 'import torch\n'), ((13, 24, 13, 86), 'torch.nn.Conv2d', 'torch.nn.Conv2d', ({(13, 40, 13, 57): 'self.out_channels', (13, 59, 13, 76): 'self.out_channels', (13, 78, 13, 79): '3', (13, 81, 13, 82): '1', (13, 84, 13, 85): '1'}, {}), '(self.out_channels, self.out_channels, 3, 1, 1)', False, 'import torch\n'), ((15, 24, 15, 99), 'torch.nn.Conv2d', 'torch.nn.Conv2d', (), '', False, 'import torch\n'), ((16, 24, 16, 86), 'torch.nn.Conv2d', 'torch.nn.Conv2d', ({(16, 40, 16, 57): 'self.out_channels', (16, 59, 16, 76): 'self.out_channels', (16, 78, 16, 79): '3', (16, 81, 16, 82): '1', (16, 84, 16, 85): '1'}, {}), '(self.out_channels, self.out_channels, 3, 1, 1)', False, 'import torch\n'), ((18, 24, 18, 99), 'torch.nn.Conv2d', 'torch.nn.Conv2d', (), '', False, 'import torch\n'), ((19, 24, 19, 86), 'torch.nn.Conv2d', 'torch.nn.Conv2d', ({(19, 40, 19, 57): 'self.out_channels', (19, 59, 19, 76): 'self.out_channels', (19, 78, 19, 79): '3', (19, 81, 19, 82): '1', (19, 84, 19, 85): '1'}, {}), '(self.out_channels, self.out_channels, 3, 1, 1)', False, 'import torch\n'), ((61, 15, 61, 69), 'torch.nn.functional.interpolate', 'functional.interpolate', (), '', False, 'from torch.nn import functional\n')]