import torch, torchvision
import torch.nn as nn
import torch.nn.functional as F
from collections import OrderedDict
from torchvision.models import ResNet50_Weights
# 没有自定义库
# Resnet50FPN  和 CountRegressor 两个类
# weights_normal_init  和 weights_xavier_init 两个函数
# 类中的函数叫方法

# 定义 Resnet50FPN 类，继承自 PyTorch 的 nn.Module
# 类 函数 方法 类里面的函数叫方法

# 这个网络是真简单，就是把 ResNet-50 的前四层作为 conv1，然后分别提取第五、第六、第七层作为 conv2、conv3、conv4
# 只是返回了 conv3 和 conv4 的特征图，没有做其他操作

# 返回特征图 然后回归
class Resnet50FPN(nn.Module):
    # 类的初始化方法
    def __init__(self):
        # 调用父类 nn.Module 的初始化方法
        super(Resnet50FPN, self).__init__()

        # 以下三行代码被注释掉，它们分别使用不同的方式加载预训练的 ResNet-50 模型
        # 使用 torchvision 库加载预训练的 ResNet-50 模型，参数 pretrained=True 表示下载并使用 ImageNet 上预训练的权重
        # self.resnet = torchvision.models.resnet50(pretrained=True)
        # 使用自定义的 ResNet50_Weights.DEFAULT 作为权重，这可能是一个特定的权重配置
        # self.resnet = torchvision.models.resnet50(weights=ResNet50_Weights.DEFAULT)
        # 选择使用 ImageNet1K_V1 预训练权重的 ResNet-50 模型
        self.resnet = torchvision.models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)        
        
        # 获取 resnet 模型的所有子模块
        children = list(self.resnet.children())
        # 分别提取 ResNet-50 中的各个层

        # 提取前四层作为 conv1
        self.conv1 = nn.Sequential(*children[:4])
        self.conv2 = children[4] # 提取第五层作为 conv2
        self.conv3 = children[5] # 提取第六层作为 conv3
        self.conv4 = children[6] # 提取第七层作为 conv4
    # 定义前向传播方法
    def forward(self, im_data):        
        feat = OrderedDict()  # 创建一个有序字典 feat 用于存储特征图
        feat_map = self.conv1(im_data)  # 通过 conv1 层获取特征图
        feat_map = self.conv2(feat_map) #  将 conv1 的输出传递给 conv2       
        feat_map3 = self.conv3(feat_map) # 将 conv2 的输出传递给 conv3
        feat_map4 = self.conv4(feat_map3)  # 将 conv3 的输出传递给 conv4
        feat['map3'] = feat_map3 # 将不同层的特征图存储到 feat 字典中
        feat['map4'] = feat_map4
        return feat

# 定义 CountRegressor 类，继承自 PyTorch 的 nn.Module
# 目的是将输入的特征图转换为一个预测的密度图
# 图像的物体计数问题
# 将输入特征图转换为最终的预测密度图
# 所以是在预测图？？？怎么实现的物体计数，输出也不是很懂。
class CountRegressor(nn.Module):
    # 类的初始化方法
    def __init__(self, input_channels,pool='mean'):

        # 调用父类 nn.Module 的初始化方法
        super(CountRegressor, self).__init__()

        # 属性 pool 用于指定池化操作的类型，可以是 'mean' 或 'max'
        self.pool = pool

        # 定义回归器网络结构，使用 nn.Sequential 来顺序添加层
        # 卷积 激活 采样 卷积 激活 采样... 最后一层卷积激活
        self.regressor = nn.Sequential(
            nn.Conv2d(input_channels, 196, 7, padding=3),   # 第一个卷积层
            nn.ReLU(),  # 激活函数
            nn.UpsamplingBilinear2d(scale_factor=2),   # 上采样

            nn.Conv2d(196, 128, 5, padding=2),
            nn.ReLU(),
            nn.UpsamplingBilinear2d(scale_factor=2),

            nn.Conv2d(128, 64, 3, padding=1),
            nn.ReLU(),
            nn.UpsamplingBilinear2d(scale_factor=2),

            nn.Conv2d(64, 32, 1),
            nn.ReLU(),

            nn.Conv2d(32, 1, 1), # 最后一层卷积层, 输出通道数为 1
            # in_channels=32，out_channels=1，kernel_size=1
            nn.ReLU(),
        )

    # 前向传播方法
    def forward(self, im):
        # 获取输入样本的数量
        num_sample =  im.shape[0]

        # 如果输入只有一个样本
        if num_sample == 1:
            # 移除批次维度，并通过网络获取输出
            output = self.regressor(im.squeeze(0))

            # 根据pool属性选择池化操作
            if self.pool == 'mean':
                output = torch.mean(output, dim=(0),keepdim=True)  
                return output
            elif self.pool == 'max':
                output, _ = torch.max(output, 0,keepdim=True)
                return output
        else:
            for i in range(0,num_sample):
                # 对每个样本执行网络前向传播
                output = self.regressor(im[i])

                # 根据pool属性选择池化操作
                if self.pool == 'mean':
                    output = torch.mean(output, dim=(0),keepdim=True)
                elif self.pool == 'max':
                    output, _ = torch.max(output, 0,keepdim=True)
                
                # 如果是第一个样本，直接赋值Output
                if i == 0:
                    Output = output
                # 否则，将输出添加到Output列表中
                else:
                    Output = torch.cat((Output,output),dim=0)
            
            # 返回所有样本的输出
            return Output

# 正态分布初始化模型参数
def weights_normal_init(model, dev=0.01):

    # 如果传入的 model 是一个列表，则遍历列表中的每个模型
    if isinstance(model, list):
        for m in model:
            weights_normal_init(m, dev)  #【递归调用】 对列表中的每个模型递归调用初始化函数
    else:
        # 如果传入的 model 不是一个列表，则对单个模型进行操作
        # 遍历模型的所有子模块
        for m in model.modules():

            # 如果当前模块是 nn.Conv2d 类型（二维卷积层）
            if isinstance(m, nn.Conv2d):   

                # 对权重进行正态分布初始化，均值为 0.0，标准差为 dev             
                m.weight.data.normal_(0.0, dev)

                # 如果存在偏置项，将其初始化为 0.0
                if m.bias is not None:
                    m.bias.data.fill_(0.0)
            
            # 如果当前模块是 nn.Linear 类型（全连接层）
            elif isinstance(m, nn.Linear):

                # 对全连接层的权重也进行正态分布初始化，均值为 0.0，标准差为 dev
                m.weight.data.normal_(0.0, dev)

                # 全连接层可能有偏置项，这里没有显式地初始化偏置，因为 nn.Linear 默认偏置为 None
                # 代码中没有对 nn.BatchNorm2d、nn.ReLU 或其他类型的层进行特殊处理，因为这些层可能不需要权重初始化，或者初始化方式可能不同。

# Xavier初始化
def weights_xavier_init(m):
    #  检查传入的模块 m 是否是二维卷积层 nn.Conv2d
    if isinstance(m, nn.Conv2d):

        # 使用 Xavier 正态分布初始化权重
        # xavier_normal_ 函数是 PyTorch 中用于实现 Xavier 初始化的函数
        # gain 参数用于调整权重的缩放因子，
        # calculate_gain('relu') 所使用的激活函数（在这个例子中是 ReLU）来计算缩放因子
        torch.nn.init.xavier_normal_(m.weight, 
                                     gain=nn.init.calculate_gain('relu'))
        
        # 如果卷积层有偏置项
        if m.bias is not None:

            # 使用 zeros_ 函数将偏置初始化为 0
            torch.nn.init.zeros_(m.bias)
            
# 应该是试了两个初始化方法 这个py文件没有返回值，都是定义的类和函数，类里的函数是方法