# author: Xinyuzhou
# email: Xinyuzhou@sjtu.edu.cn
# date: 2025-12-24
#
# The code is for ResNet50 backbone.

import os
import logging
from typing import Union
import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
from metrics.registry import BACKBONE

logger = logging.getLogger(__name__)

@BACKBONE.register_module(module_name="resnet50")
class ResNet50(nn.Module):
    def __init__(self, resnet_config):
        super(ResNet50, self).__init__()
        """ Constructor
        Args:
            resnet_config: configuration file with the dict format
        """
        self.num_classes = resnet_config["num_classes"]
        inc = resnet_config["inc"]
        self.mode = resnet_config["mode"]

        # Load the pretrained ResNet50 weights from torchvision
        # New API for torchvision >= 0.13:
        #   resnet = torchvision.models.resnet50(weights=torchvision.models.ResNet50_Weights.IMAGENET1K_V1)
        # The legacy API is still available:
        resnet = torchvision.models.resnet50(pretrained=True)

        # If you want to support inc != 3, uncomment this and adjust accordingly
        # resnet.conv1 = nn.Conv2d(inc, 64, kernel_size=7, stride=2, padding=3, bias=False)

        # Remove the final avgpool and fc layers, keeping feature maps up to layer4
        self.resnet = torch.nn.Sequential(*list(resnet.children())[:-2])

        # The final feature dimension of ResNet50 is 2048
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(2048, self.num_classes)

        if self.mode == 'adjust_channel':
            self.adjust_channel = nn.Sequential(
                nn.Conv2d(2048, 2048, 1, 1),
                nn.BatchNorm2d(2048),
                nn.ReLU(inplace=True),
            )

    def features(self, inp):
        x = self.resnet(inp)
        if self.mode == 'adjust_channel':
            x = self.adjust_channel(x)
        return x

    def classifier(self, features):
        x = self.avgpool(features)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

    def forward(self, inp):
        x = self.features(inp)
        out = self.classifier(x)
        return out