# YOLOFPN 全称是 YOLOv5 Feature Pyramid Network
# 该模块是 YOLOv5 的 backbone，用于提取特征图，用于后续的检测，分割等任务
class YOLOFPN(nn.Module):
    """
    YOLOFPN module. Darknet 53 is the default backbone of this model.
    """

    def __init__(
        self,
        depth=53,
        in_features=["dark3", "dark4", "dark5"],
    ):
        super().__init__()

        self.backbone = Darknet(depth)  # backbone，Darknet53，用于提取特征图
        self.in_features = in_features  # in_features，用于提取特征图的层级

        # out 1
        self.out1_cbl = self._make_cbl(512, 256, 1) # cbl全称是conv-bn-leakyrelu
                                                    # 该层是将输入的通道数从512变为256
                                                    # 该层的作用是将输入的特征图进行下采样
        self.out1 = self._make_embedding([256, 512], 512 + 256)
        # embedding，嵌入层，用于将输入的特征图进行下采样，同时将输入的特征图与下采样后的特征图进行融合
        # 其中[256, 512]是嵌入层的卷积核的个数，512+256是输入的通道数
        # 512+256是因为下采样后的特征图的通道数是256，输入的特征图的通道数是512，所以需要将两者进行融合
        # 融合的方式是将两者的特征图进行拼接，然后再进行卷积操作，最终输出的通道数是512

        # out 2
        self.out2_cbl = self._make_cbl(256, 128, 1)
        # 该层是将输入的通道数从256变为128
        self.out2 = self._make_embedding([128, 256], 256 + 128)
        # embedding，嵌入层，用于将输入的特征图进行下采样，同时将输入的特征图与下采样后的特征图进行融合
        # 其中[128, 256]是嵌入层的卷积核的个数，256+128是输入的通道数
        # 256+128是因为下采样后的特征图的通道数是128，输入的特征图的通道数是256，所以需要将两者进行融合
        # 融合的方式是将两者的特征图进行拼接，然后再进行卷积操作，最终输出的通道数是256

        # upsample
        self.upsample = nn.Upsample(scale_factor=2, mode="nearest")
        # 上采样层，用于将输入的特征图进行上采样，上采样的方式是将输入的特征图进行放大，放大的倍数是2倍
        # scale_factor=2，表示放大的倍数是2倍
        # mode="nearest"，表示放大的方式是最近邻插值

    def _make_cbl(self, _in, _out, ks):
        return BaseConv(_in, _out, ks, stride=1, act="lrelu")

    def _make_embedding(self, filters_list, in_filters):
        m = nn.Sequential(
            *[
                self._make_cbl(in_filters, filters_list[0], 1),
                self._make_cbl(filters_list[0], filters_list[1], 3),
                self._make_cbl(filters_list[1], filters_list[0], 1),
                self._make_cbl(filters_list[0], filters_list[1], 3),
                self._make_cbl(filters_list[1], filters_list[0], 1),
            ]
        )
        return m

    def load_pretrained_model(self, filename="./weights/darknet53.mix.pth"):
        # load pretrained model的作用是加载预训练模型
        # 其中mix是指将darknet53的预训练模型与yolov5的预训练模型进行融合
        with open(filename, "rb") as f:
            state_dict = torch.load(f, map_location="cpu")
            # map_location="cpu"，表示将模型加载到cpu上
        print("loading pretrained weights...")
        self.backbone.load_state_dict(state_dict)
        # backbone是Darknet53，用于提取特征图，load_state_dict是用于加载预训练模型的函数

    def forward(self, inputs):
        """
        Args:
            inputs (Tensor): input image.

        Returns:
            Tuple[Tensor]: FPN output features..
        """
        #  backbone
        out_features = self.backbone(inputs)
        x2, x1, x0 = [out_features[f] for f in self.in_features]
        # 将dark3, dark4, dark5三个层级的特征图进行提取，对应的是yolo的三个分支

        #  yolo branch 1
        x1_in = self.out1_cbl(x0)
        # 将输入的特征图进行下采样，下采样的方式是将输入的特征图进行卷积操作，卷积核的个数是256
        x1_in = self.upsample(x1_in)
        # 将下采样后的特征图进行上采样，上采样的方式是将输入的特征图进行放大，放大的倍数是2倍
        x1_in = torch.cat([x1_in, x1], 1)
        # 将下采样后的特征图与输入的特征图进行融合，融合的方式是将两者的特征图进行拼接
        out_dark4 = self.out1(x1_in)
        # 将融合后的特征图进行卷积操作，卷积核的个数是[256, 512]，最终输出的通道数是512

        #  yolo branch 2
        x2_in = self.out2_cbl(out_dark4)
        # 将输入的特征图进行下采样，下采样的方式是将输入的特征图进行卷积操作，卷积核的个数是128
        x2_in = self.upsample(x2_in)
        # 将下采样后的特征图进行上采样，上采样的方式是将输入的特征图进行放大，放大的倍数是2倍
        x2_in = torch.cat([x2_in, x2], 1)
        # 将下采样后的特征图与输入的特征图进行融合，融合的方式是将两者的特征图进行拼接
        out_dark3 = self.out2(x2_in)
        # 将融合后的特征图进行卷积操作，卷积核的个数是[128, 256]，最终输出的通道数是256

        outputs = (out_dark3, out_dark4, x0)
        # 将dark3, dark4, dark5三个层级的特征图进行融合，对应的是yolo的三个分支
        return outputs
# YOLOFPN的数据流程图
# W*H*3 -> backbone -> dark3, dark4, dark5 -> out1_cbl -> out1 -> out2_cbl -> out2 -> upsample -> out_dark4 -> upsample -> out_dark3
# 该模型的输入是W*H*3的图片，经过backbone提取特征图，然后将dark3, dark4, dark5三个层级的特征图进行提取，对应的是yolo的三个分支
# 然后将dark3, dark4, dark5三个层级的特征图进行融合，对应的是yolo的三个分支