yolo-world 源码解析（五）（4）

# 版权声明
# 导入数学库
import math
# 导入类型提示相关库
from typing import List, Optional, Tuple, Union, Sequence
# 导入 PyTorch 库
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torch.nn.modules.batchnorm import _BatchNorm
# 导入 mmcv 库中的模块
from mmcv.cnn import ConvModule
from mmengine.config import ConfigDict
from mmengine.dist import get_dist_info
from mmengine.structures import InstanceData
from mmdet.structures import SampleList
from mmdet.utils import (ConfigType, OptConfigType, OptInstanceList,
                         OptMultiConfig, InstanceList)
from mmdet.models.utils import multi_apply, unpack_gt_instances
from mmyolo.models.dense_heads import YOLOv8HeadModule
from mmyolo.models.utils import gt_instances_preprocess
from mmyolo.registry import MODELS, TASK_UTILS
from mmyolo.models.dense_heads.yolov5_ins_head import (
    ProtoModule, YOLOv5InsHead
)
# 导入自定义的模块
from .yolo_world_head import ContrastiveHead, BNContrastiveHead
# 注册 YOLOWorldSegHeadModule 类为模型
@MODELS.register_module()
class YOLOWorldSegHeadModule(YOLOv8HeadModule):
    # 初始化方法
    def __init__(self,
                 *args,
                 embed_dims: int,
                 proto_channels: int,
                 mask_channels: int,
                 freeze_bbox: bool = False,
                 use_bn_head: bool = False,
                 **kwargs) -> None:
        # 初始化属性
        self.freeze_bbox = freeze_bbox
        self.embed_dims = embed_dims
        self.proto_channels = proto_channels
        self.mask_channels = mask_channels
        self.use_bn_head = use_bn_head
        # 调用父类的初始化方法
        super().__init__(*args, **kwargs)
    def init_weights(self, prior_prob=0.01):
        """初始化PPYOLOE头部的权重和偏置。"""
        # 调用父类的初始化权重方法
        super().init_weights()
        # 遍历分类预测、分类对比和特征图步长，分别初始化偏置
        for cls_pred, cls_contrast, stride in zip(self.cls_preds,
                                                  self.cls_contrasts,
                                                  self.featmap_strides):
            cls_pred[-1].bias.data[:] = 0.0  # 重置偏置
            # 如果分类对比具有偏置属性，则初始化为特定值
            if hasattr(cls_contrast, 'bias'):
                nn.init.constant_(
                    cls_contrast.bias.data,
                    math.log(5 / self.num_classes / (640 / stride)**2))
    def head_norm_eval(self):
        # 遍历分类预测模块，将所有批归一化层设置为评估模式
        for m in self.cls_preds:
            for q in m.modules():
                if isinstance(q, _BatchNorm):
                    q.eval()
        # 遍历回归预测模块，将所有批归一化层设置为评估模式
        for m in self.reg_preds:
            for q in m.modules():
                if isinstance(q, _BatchNorm):
                    q.eval()
    def train(self, mode: bool = True):
        """将模型转换为训练模式，同时保持归一化层冻结。"""
        # 调用父类的训练方法
        super().train(mode)
        # 如果冻结边界框，则调用头部归一化评估方法
        if self.freeze_bbox:
            self.head_norm_eval()
    def forward(self, img_feats: Tuple[Tensor],
                txt_feats: Tensor) -> Tuple[List]:
        """从上游网络前向传播特征。"""
        # 断言图像特征的长度等于级别数
        assert len(img_feats) == self.num_levels
        # 将文本特征复制多份以匹配级别数
        txt_feats = [txt_feats for _ in range(self.num_levels)]
        # 生成掩码原型
        mask_protos = self.proto_pred(img_feats[0])
        # 多路并行处理，获取分类logit、边界框预测、边界框距离预测和系数预测
        cls_logit, bbox_preds, bbox_dist_preds, coeff_preds = multi_apply(
            self.forward_single, img_feats, txt_feats, self.cls_preds,
            self.reg_preds, self.cls_contrasts, self.seg_preds)
        # 如果处于训练模式，则返回所有预测结果和掩码原型
        if self.training:
            return cls_logit, bbox_preds, bbox_dist_preds, coeff_preds, mask_protos
        # 否则，返回分类logit、边界框预测、系数预测和掩码原型
        else:
            return cls_logit, bbox_preds, None, coeff_preds, mask_protos
    def forward_single(self, img_feat: Tensor, txt_feat: Tensor,
                       cls_pred: nn.ModuleList, reg_pred: nn.ModuleList,
                       cls_contrast: nn.ModuleList,
                       seg_pred: nn.ModuleList) -> Tuple:
        """Forward feature of a single scale level."""
        # 获取输入特征的形状信息
        b, _, h, w = img_feat.shape
        # 使用分类预测模型对图像特征进行预测
        cls_embed = cls_pred(img_feat)
        # 使用对比损失模型对分类嵌入进行预测
        cls_logit = cls_contrast(cls_embed, txt_feat)
        # 使用回归预测模型对图像特征进行预测
        bbox_dist_preds = reg_pred(img_feat)
        # 使用分割预测模型对图像特征进行预测
        coeff_pred = seg_pred(img_feat)
        # 如果回归最大值大于1
        if self.reg_max > 1:
            # 重塑回归预测结果的形状
            bbox_dist_preds = bbox_dist_preds.reshape(
                [-1, 4, self.reg_max, h * w]).permute(0, 3, 1, 2)
            # TODO: get_flops脚本无法处理矩阵乘法的情况，稍后需要修复
            # 计算边界框预测结果
            bbox_preds = bbox_dist_preds.softmax(3).matmul(
                self.proj.view([-1, 1])).squeeze(-1)
            bbox_preds = bbox_preds.transpose(1, 2).reshape(b, -1, h, w)
        else:
            bbox_preds = bbox_dist_preds
        # 如果处于训练模式
        if self.training:
            return cls_logit, bbox_preds, bbox_dist_preds, coeff_pred
        else:
            return cls_logit, bbox_preds, None, coeff_pred
# 注册 YOLO World Segmentation Head 类到 MODELS 模块
@MODELS.register_module()
class YOLOWorldSegHead(YOLOv5InsHead):
    # 特殊初始化函数，用于处理不同算法的特殊初始化过程
    def special_init(self):
        """Since YOLO series algorithms will inherit from YOLOv5Head, but
        different algorithms have special initialization process.
        The special_init function is designed to deal with this situation.
        """
        # 如果存在训练配置，则构建分配器
        if self.train_cfg:
            self.assigner = TASK_UTILS.build(self.train_cfg.assigner)
            # 添加常用属性以减少计算
            self.featmap_sizes_train = None
            self.num_level_priors = None
            self.flatten_priors_train = None
            self.stride_tensor = None
    """YOLO World head."""
    # 损失函数，计算前向传播和检测头特征的损失
    def loss(self, img_feats: Tuple[Tensor], txt_feats: Tensor,
             batch_data_samples: Union[list, dict]) -> dict:
        """Perform forward propagation and loss calculation of the detection
        head on the features of the upstream network."""
        # 执行前向传播并获取输出
        outs = self(img_feats, txt_feats)
        # 快速版本
        loss_inputs = outs + (batch_data_samples['bboxes_labels'],
                              batch_data_samples['masks'],
                              batch_data_samples['img_metas'])
        # 计算损失
        losses = self.loss_by_feat(*loss_inputs)
        return losses
    # 损失和预测函数
    def loss_and_predict(
        self,
        img_feats: Tuple[Tensor],
        txt_feats: Tensor,
        batch_data_samples: SampleList,
        proposal_cfg: Optional[ConfigDict] = None
    def forward(self, img_feats: Tuple[Tensor],
                txt_feats: Tensor) -> Tuple[List]:
        """Forward features from the upstream network."""
        # 从上游网络中前向传播特征
        return self.head_module(img_feats, txt_feats)
    def predict(self,
                img_feats: Tuple[Tensor],
                txt_feats: Tensor,
                batch_data_samples: SampleList,
                rescale: bool = False) -> InstanceList:
        """Perform forward propagation of the detection head and predict
        detection results on the features of the upstream network.
        """
        # 从检测头部进行前向传播，并在上游网络的特征上预测检测结果
        # 获取批量数据样本的元信息
        batch_img_metas = [
            data_samples.metainfo for data_samples in batch_data_samples
        ]
        # 获取模型输出
        outs = self(img_feats, txt_feats)
        # 根据模型输出进行预测
        predictions = self.predict_by_feat(*outs,
                                           batch_img_metas=batch_img_metas,
                                           rescale=rescale)
        return predictions
    def forward(self, img_feats: Tuple[Tensor],
                txt_feats: Tensor) -> Tuple[dict, InstanceList]:
        """Perform forward propagation of the head, then calculate loss and
        predictions from the features and data samples.
        """
        # 解包批量数据样本
        outputs = unpack_gt_instances(batch_data_samples)
        (batch_gt_instances, batch_gt_instances_ignore,
         batch_img_metas) = outputs
        # 获取模型输出
        outs = self(img_feats, txt_feats)
        # 构建损失函数输入
        loss_inputs = outs + (batch_gt_instances, batch_img_metas,
                              batch_gt_instances_ignore)
        # 计算损失
        losses = self.loss_by_feat(*loss_inputs)
        # 根据模型输出进行预测
        predictions = self.predict_by_feat(*outs,
                                           batch_img_metas=batch_img_metas,
                                           cfg=proposal_cfg)
        return losses, predictions
    # 定义一个测试函数，用于测试时进行数据增强
    def aug_test(self,
                 aug_batch_feats,
                 aug_batch_img_metas,
                 rescale=False,
                 with_ori_nms=False,
                 **kwargs):
        """Test function with test time augmentation."""
        # 抛出未实现错误，提示该函数还未被实现
        raise NotImplementedError('aug_test is not implemented yet.')

# 导入 YOLOWorldHead 和 YOLOWorldHeadModule 类
from .yolo_world_head import YOLOWorldHead, YOLOWorldHeadModule
# 导入 YOLOWorldSegHead 和 YOLOWorldSegHeadModule 类
from .yolo_world_seg_head import YOLOWorldSegHead, YOLOWorldSegHeadModule
# 定义 __all__ 列表，包含需要导出的类名
__all__ = [
    'YOLOWorldHead', 'YOLOWorldHeadModule', 'YOLOWorldSegHead',
    'YOLOWorldSegHeadModule'
]

# 导入所需的模块和类
from typing import List, Tuple, Union
from torch import Tensor
from mmdet.structures import OptSampleList, SampleList
from mmyolo.models.detectors import YOLODetector
from mmyolo.registry import MODELS
# 注册YOLOWorldDetector类到MODELS模块
@MODELS.register_module()
class YOLOWorldDetector(YOLODetector):
    """Implementation of YOLOW Series"""
    # 初始化函数，接受一些参数
    def __init__(self,
                 *args,
                 mm_neck: bool = False,
                 num_train_classes=80,
                 num_test_classes=80,
                 **kwargs) -> None:
        # 初始化类的属性
        self.mm_neck = mm_neck
        self.num_train_classes = num_train_classes
        self.num_test_classes = num_test_classes
        # 调用父类的初始化函数
        super().__init__(*args, **kwargs)
    # 计算损失函数的方法，接受输入和数据样本
    def loss(self, batch_inputs: Tensor,
             batch_data_samples: SampleList) -> Union[dict, list]:
        """Calculate losses from a batch of inputs and data samples."""
        # 设置bbox_head的类别数为训练类别数
        self.bbox_head.num_classes = self.num_train_classes
        # 提取图像特征和文本特征
        img_feats, txt_feats = self.extract_feat(batch_inputs,
                                                 batch_data_samples)
        # 计算损失
        losses = self.bbox_head.loss(img_feats, txt_feats, batch_data_samples)
        # 返回损失
        return losses
    # 预测模型的方法，接受批量输入和数据样本，返回带有后处理的结果列表
    def predict(self,
                batch_inputs: Tensor,
                batch_data_samples: SampleList,
                rescale: bool = True) -> SampleList:
        """Predict results from a batch of inputs and data samples with post-
        processing.
        """
        # 提取图像特征和文本特征
        img_feats, txt_feats = self.extract_feat(batch_inputs,
                                                 batch_data_samples)
        # 设置边界框头部的类别数为文本特征的第一个维度大小
        self.bbox_head.num_classes = txt_feats[0].shape[0]
        
        # 使用图像特征、文本特征和数据样本进行预测，返回结果列表
        results_list = self.bbox_head.predict(img_feats,
                                              txt_feats,
                                              batch_data_samples,
                                              rescale=rescale)
        # 将预测结果添加到数据样本中
        batch_data_samples = self.add_pred_to_datasample(
            batch_data_samples, results_list)
        
        # 返回更新后的数据样本
        return batch_data_samples
    # 网络前向传播过程，通常包括骨干网络、颈部和头部的前向传播，不包含任何后处理
    def _forward(
            self,
            batch_inputs: Tensor,
            batch_data_samples: OptSampleList = None) -> Tuple[List[Tensor]]:
        """Network forward process. Usually includes backbone, neck and head
        forward without any post-processing.
        """
        
        # 提取图像特征和文本特征
        img_feats, txt_feats = self.extract_feat(batch_inputs,
                                                 batch_data_samples)
        
        # 进行边界框头部的前向传播，返回结果
        results = self.bbox_head.forward(img_feats, txt_feats)
        
        # 返回结果
        return results
    # 定义一个方法用于提取特征，接受两个输入参数：batch_inputs（张量）和batch_data_samples（样本列表），返回一个元组
    def extract_feat(
            self, batch_inputs: Tensor,
            batch_data_samples: SampleList) -> Tuple[Tuple[Tensor], Tensor]:
        """Extract features."""
        # 如果batch_data_samples是字典类型，则获取其中的'texts'键对应的值
        if isinstance(batch_data_samples, dict):
            texts = batch_data_samples['texts']
        # 如果batch_data_samples是列表类型，则遍历其中的数据样本，获取每个数据样本的文本信息
        elif isinstance(batch_data_samples, list):
            texts = [data_sample.texts for data_sample in batch_data_samples]
        # 如果batch_data_samples既不是字典类型也不是列表类型，则抛出类型错误异常
        else:
            raise TypeError('batch_data_samples should be dict or list.')
        # 调用backbone模型提取图像和文本特征
        img_feats, txt_feats = self.backbone(batch_inputs, texts)
        # 如果模型包含neck部分
        if self.with_neck:
            # 如果使用多模态neck
            if self.mm_neck:
                # 将图像特征和文本特征输入到neck模块中进行处理
                img_feats = self.neck(img_feats, txt_feats)
            else:
                # 只将图像特征输入到neck模块中进行处理
                img_feats = self.neck(img_feats)
        # 返回提取的图像特征和文本特征
        return img_feats, txt_feats