快速上手！MMTracking 食用指南 之 VID 篇（附 AAAI2021 论文解读 ！）

python ./演示demo/演示demo_vid.py \
    ./configs/vid/temporal_roi_align/selsa_troialign_faster_rcnn_r101_dc5_7e_imagenetvid.py \
    --input ${VIDEO_FILE} \
    --checkpoint checkpoints/selsa_troialign_faster_rcnn_r101_dc5_7e_imagenetvid_20210822_111621-22cb96b9.pth \
    --output ${OUTPUT} \
    --show

./tools/dist_test.sh configs/vid/temporal_roi_align/selsa_troialign_faster_rcnn_r101_dc5_7e_imagenetvid.py 8 \
    --checkpoint checkpoints/selsa_troialign_faster_rcnn_r101_dc5_7e_imagenetvid_20210822_111621-22cb96b9.pth \
    --out results.pkl \
    --eval bbox

./tools/dist_train.sh ./configs/temporal_roi_align/selsa_troialign_faster_rcnn_r101_dc5_7e_imagenetvid.py 8 \
    --work-dir ./work_dirs/

bbox_roi_extractor=dict(
    type='TemporalRoIAlign',
    num_most_similar_points=2,
    num_temporal_attention_blocks=4,
    roi_layer=dict(
        type='RoIAlign', output_size=7, sampling_ratio=2),
    out_channels=512,
    featmap_strides=[16])

from mmdet.models.roi_heads.roi_extractors import SingleRoIExtractor
class TemporalRoIAlign(SingleRoIExtractor):
    """Temporal RoI Align 组件.
    本组件在以下论文中提出
    "Temporal ROI Align for Video Object Recognition".
    `TRoI Align <https://arxiv.org/abs/2109.03495>`_.
    Args:
        num_most_similar_points (int): 代表寻找的最大的 K 个点。 默认为 2
        num_temporal_attention_blocks (int):代表使用多少个 attention block 来聚合
            特征，如果该值小于等于 0，则使用平均池化(Pooling)(Average-Pooling)来聚合特征。默认为 4
    """
    def __init__(self,
                 num_most_similar_points=2,
                 num_temporal_attention_blocks=4,
                 *args,
                 **kwargs):
        super(TemporalRoIAlign, self).__init__(*args, **kwargs)
        self.num_most_similar_points = num_most_similar_points
        self.num_temporal_attention_blocks = num_temporal_attention_blocks
    def temporal_attentional_feature_aggregation(self, x, ref_x):
        """聚合 RoI features `x` 和 Most Similar RoI features `ref_x`.
        聚合主要包含 3 步：
        1. 通过一个小的 embed network 提取特征
        2. 使用 multi-head attention 计算 `x` 和 `ref_x` 之间的权重
        3. 使用归一化(Normalization)之后的权重聚合 `x` 和 `ref_x`
        Args:
            x (Tensor): 大小为 [1, roi_n, C, roi_h, roi_w]。 roi_n, roi_h 和
                roi_w 分别代表当前帧 proposals 的数量， RoI Features 的高和宽。
            ref_x (Tensor): 大小为 [img_n, roi_n, C, roi_h, roi_w]。 img_n 是其
                它帧的数量。
        Returns:
            Tensor: 聚合之后的 Temporal RoI Features，大小为
                [roi_n, C, roi_h, roi_w]。
        """
    def most_similar_roi_align(self, roi_feats, ref_feats):
        """基于 RoI Features `roi_feats` 来从其它帧的特征图 `ref_feats` 提取
        Most Similar RoI features。
        提取过程主要包含 3 步：
        1. 计算 `roi_feats` 和 `ref_feats`之间的 cosine 相似度图
        2. 在相似图中挑去 Top K 个点
        3. 将 Top K 个点投射到 `ref_feats` 中提取 Most Similar RoI features
        Args:
            roi_feats (Tensor): 大小为 [roi_n, C, roi_h, roi_w]。
                roi_n, roi_h 和 roi_w 分别代表当前帧 proposals 的数量
                RoI Features 的高和宽。
            ref_feats (Tensor): 大小为 [img_n, C, img_h, img_w]。
                img_n, img_h 和 img_w 分别达标其它帧的数量，其它帧特征图的高和宽
        Returns:
            Tensor: 从其它帧特征图提取出的 Most Similar RoI features，
                大小为 [img_n, roi_n, C, roi_h, roi_w]。
        """
    def forward(self, feats, rois, roi_scale_factor=None, ref_feats=None):
        """Forward function."""
        # 第一步
        roi_feats = super().forward(feats, rois, roi_scale_factor)
        if ref_feats is None:
            # 当没有其它帧的特征图时，直接返回 RoI Features
            return roi_feats
        else:
            # 第二步
            # 我们只用其它帧特征图的最后一层特征图来计算 Most Similar RoI Align
            ref_roi_feats = self.most_similar_roi_align(
                roi_feats, ref_feats[-1])
            roi_feats = roi_feats.unsqueeze(0)
            if self.num_temporal_attention_blocks > 0:
                # 第三步
                temporal_roi_feats = \
                    self.temporal_attentional_feature_aggregation(
                        roi_feats, ref_roi_feats)
            else:
                temporal_roi_feats = torch.cat((roi_feats, ref_roi_feats),
                                               dim=0)
                temporal_roi_feats = temporal_roi_feats.mean(dim=0)
            return temporal_roi_feats