import torch
import torch.nn as nn
from torchvision.models import resnet50, ResNet50_Weights
from torchvision.models._utils import IntermediateLayerGetter
from torchvision.ops import FeaturePyramidNetwork
from torchvision.ops.feature_pyramid_network import LastLevelMaxPool
class ResNet50FPNBackbone(nn.Module):
    def __init__(self, pretrained=True, out_channels=256):
        super().__init__()
        weights = ResNet50_Weights.DEFAULT if pretrained else None
        resnet = resnet50(weights=weights)
        self.body = IntermediateLayerGetter(
            resnet,
            return_layers={
                "layer1": "c2",
                "layer2": "c3",
                "layer3": "c4",
                "layer4": "c5"
            }
        )
        self.fpn = FeaturePyramidNetwork(
            in_channels_list=[
                256,
                512,
                1024,
                2048
            ],
            out_channels=out_channels,
            extra_blocks=LastLevelMaxPool()
        )
        self.out_channels = out_channels
    def forward(self, x):
        features = self.body(x)
        features = self.fpn(features)
        return features