[人工智能]基于pytorch的Faster RCNN中RPN实现

RegionProposalNetwork

在Faster RCNN中第一阶段是由RegionProposalNetwork生成anchors，并通过筛选得到proposal。代码中详细注释了每一部分的过程。

import torch
import torchvision
from torch import nn, Tensor
from torch.nn import functional as F
import math
from typing import Dict


def smooth_l1_loss(input, target, beta: float = 1. / 9, size_average: bool = True):
    """
    very similar to the smooth_l1_loss from pytorch, but with
    the extra beta parameter
    """
    n = torch.abs(input - target)
    # cond = n < beta. lt: 代表小于操作
    cond = torch.lt(n, beta)
    loss = torch.where(cond, 0.5 * n ** 2 / beta, n - 0.5 * beta)
    if size_average:
        return loss.mean()
    return loss.sum()


def nms(boxes, scores, iou_threshold):
    # type: (Tensor, Tensor, float) -> Tensor
    """
    Performs non-maximum suppression (NMS) on the boxes according
    to their intersection-over-union (IoU).

    NMS iteratively removes lower scoring boxes which have an
    IoU greater than iou_threshold with another (higher scoring)
    box.

    Parameters
    ----------
    boxes : Tensor[N, 4])
        boxes to perform NMS on. They
        are expected to be in (x1, y1, x2, y2) format
    scores : Tensor[N]
        scores for each one of the boxes
    iou_threshold : float
        discards all overlapping
        boxes with IoU < iou_threshold

    Returns
    -------
    keep : Tensor
        int64 tensor with the indices
        of the elements that have been kept
        by NMS, sorted in decreasing order of scores
    """
    return torchvision.ops.nms(boxes, scores, iou_threshold)


class RPNHead(nn.Module):
    # 计算预测目标概率与bbox regression参数
    def __init__(self, in_channels, num_anchors):
        super(RPNHead, self).__init__()

        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
        # 计算预测的类别分数（这里的类别指前景或者背景）
        self.cls_logits = nn.Conv2d(in_channels, num_anchors, kernel_size=1, stride=1)
        # 计算预测的目标bbox regression参数
        self.bbox_pred = nn.Conv2d(in_channels, num_anchors * 4, kernel_size=1, stride=1)
        # 初始化
        for layer in self.children():
            if isinstance(layer, nn.Conv2d):
                torch.nn.init.normal_(layer.weight, std=0.01)
                torch.nn.init.constant_(layer.bias, 0)

    def forward(self, x):
        # x.shape: 不同特征图的输出维度 C=256 [B, 256, H, W]
        logits = []
        bbox_reg = []
        for i, feature in enumerate(x):
            output_33 = F.relu(self.conv(feature))
            logits.append(self.cls_logits(output_33))
            bbox_reg.append(self.bbox_pred(output_33))
        return logits, bbox_reg


class AnchorsGenerator(nn.Module):
    # sizes: anchor的尺寸 aspect_ratios：anchor采用的不同的比例
    def __init__(self, sizes=(128, 256, 512), aspect_ratios=(0.5, 1.0, 2.0)):
        super(AnchorsGenerator, self).__init__()

        self.sizes = sizes
        self.aspect_ratios = aspect_ratios
        self.cell_anchors = None
        self._cache = {} # 在原图上生成的所有anchor的信息 存入到cache中

    def generate_anchors(self, scale, aspect_ratios, dtype=torch.float32, device='cpu'):
        # scale: tuple -> tensor
        scale = torch.as_tensor(scale, dtype=dtype, device=device)
        # aspect_ratios: tuple -> tensor
        aspect_ratios = torch.as_tensor(aspect_ratios, dtype=dtype, device=device)

        h_ratios = torch.sqrt(aspect_ratios)
        w_ratios = 1.0 / h_ratios

        # 每一个高、宽比例分别乘上对应预测特征图的尺度 得到每一个anchor对应的高和宽
        # ResNet而言 对应不同的预测特征图 scale仅有一个尺度 ws shape: [3]
        ws = (w_ratios[:, None] * scale[None, :]).view(-1)
        hs = (h_ratios[:, None] * scale[None, :]).view(-1)

        # 生成的anchors模板都是以（0, 0）为中心, 对应左上角坐标和右下角坐标. dim=1 shape-> [3, 4] 三个比例的anchor，4左上角坐标和右下角坐标
        base_anchor = torch.stack([-ws, -hs, ws, hs], dim=1) / 2
        # round 四舍五入取整
        return base_anchor.round()

    def set_cell_anchors(self, dtype, device):
        # 初始化时设置为None
        if self.cell_anchors is not None:
            cell_anchors = self.cell_anchors
            assert cell_anchors is not None
            # suppose that all anchors have the same device
            # which is a valid assumption in the current state of the codebase
            if cell_anchors[0].device == device:
                return

        # 根据提供的sizes和aspect_ratios生成anchors模板
        # 不同预测特征图所对应的sizes不同 遍历不同的预测特征图生成与之对应的anchor模板
        # 返回一个list key代表不同的预测特征图 value代表anchor模板
        cell_anchors = [
            self.generate_anchors(sizes, aspect_ratios, dtype, device)
            for sizes, aspect_ratios in zip(self.sizes, self.aspect_ratios)]

        self.cell_anchors = cell_anchors

    def cached_grid_anchors(self, grid_sizes, strides):
        key = str(grid_sizes) + str(strides)
        if key in self._cache:
            return self._cache[key]
        # 得到所有预测特征图映射回原图上所生成的anchors
        anchors = self.grid_anchors(grid_sizes, strides)
        self._cache[key] = anchors

        return anchors

    def grid_anchors(self, grid_sizes, strides):
        anchors = []
        # anchor模板
        cell_anchors = self.cell_anchors

        # 遍历每个预测特征图的grid_sizes strides和cell_anchors
        # size 预测特征图的高度和宽度
        # stride 预测特征层上的每一个stride对应原图上的高度和宽度
        # cell_anchors anchor模板
        for size, stride, base_anchors in zip(grid_sizes, strides, cell_anchors):
            grid_height, grid_width = size
            stride_height, stride_width = stride
            device = base_anchors.device
            # 生成对应原图上的x坐标
            shifts_x = torch.arange(0, grid_width, dtype=torch.float32, device=device) * stride_width
            # 生成对应原图上的x坐标
            shifts_y = torch.arange(0, grid_height, dtype=torch.float32, device=device) * stride_height

            # 得到每一个点映射回原图上的坐标
            shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x)
            shift_x = shift_x.reshape(-1) # 展平 -> 向量
            shift_y = shift_y.reshape(-1)
            # shifts: shape: 当前预测特征图上[grid_width*grid_height, 4]
            shifts = torch.stack([shift_x, shift_y, shift_x, shift_y], dim=1)
            # shift这个偏移量加上anchor模板 就可以平移到原图上的每个位置上
            # shape: [grid_width*grid_height, 3, 4] 3: 当前预测特征图上每个位置生成的3个anchor 4: 对应的左上右下坐标
            shifts_anchor = shifts.view(-1, 1, 4) + base_anchors.view(1, -1, 4)
            # 当前预测特征图生成的所有anchor存入anchors中
            # shape: [grid_width*grid_height*3, 4]
            anchors.append(shifts_anchor.reshape(-1, 4)) # 当前预测特征图生成的所有anchor存入anchors中

        return anchors

    def num_anchors_per_location(self):
        # 计算每个预测特征层上每个滑动窗口的预测目标数
        return [len(s) * len(a) for s, a in zip(self.sizes, self.aspect_ratios)]

    def forward(self, image_list, feature_maps):
        # 获取backbone中每个预测特征图的尺寸(h, w)
        grid_sizes = list([feature_map.shape[-2:] for feature_map in feature_maps])

        # 获取打包成batch的输入图像的height和width
        # image_list: 1 图像缩放后的尺寸(image_sizes) 2 预处理过程中将图像打包成一个个batch 每个batch对应一个tensor(tensors)
        image_size = image_list.tensors.shape[-2:]

        # 获取变量类型和设备类型
        dtype, device = feature_maps[0].dtype, feature_maps[0].device

        # 计算特征层上的每一步对应原始图像上的尺度
        # 图像大小 / 预测特征图大小 = 特征图上的每一步对应原图上像素的大小
        strides = [[torch.tensor(image_size[0] // g[0], dtype=torch.int64, device=device),
                    torch.tensor(image_size[1] // g[1], dtype=torch.int64, device=device)] for g in grid_sizes]

        # 根据提供的sizes和aspect_ratios生成anchors模板
        self.set_cell_anchors(dtype, device)

        # 将anchor模板应用到原图中
        # grid_sizes 预测特征图的高度、宽度信息
        # strides：预测特征图中每个scale对应原图上的尺度
        # 得到不同预测特征图在原图上生成的anchors的坐标信息
        anchors_over_all_feature_maps = self.cached_grid_anchors(grid_sizes, strides)

        # 遍历batch中的每一张图片
        anchors = []
        for i, (image_height, image_width) in enumerate(image_list.image_sizes):
            anchors_in_image = []
            # 遍历所有预测特征图映射回原图中anchor的坐标信息
            for anchors_per_feature_map in anchors_over_all_feature_maps:
                anchors_in_image.append(anchors_per_feature_map)
            anchors.append(anchors_in_image)
        # cat: 将一张图像中的所有预测特征图上的anchors拼接到一起
        anchors = [torch.cat(anchors_per_image) for anchors_per_image in anchors]
        self._cache.clear()

        return anchors


def box_area(boxes):
    """
    Computes the area of a set of bounding boxes, which are specified by its
    (x1, y1, x2, y2) coordinates.

    Arguments:
        boxes (Tensor[N, 4]): boxes for which the area will be computed. They
            are expected to be in (x1, y1, x2, y2) format

    Returns:
        area (Tensor[N]): area for each box
    """
    return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])


def box_iou(boxes1, boxes2):
    # 计算传入的两组boxes的IOU值
    area1 = box_area(boxes1)
    area2 = box_area(boxes2)

    # [gt_box_num, 1, 2] - [anchor_num, 2] -> shape: [gt_box_num, anchor_num, 2]
    left_top = torch.max(boxes1[:, None, :2], boxes2[:, :2])
    # [gt_box_num, 1, 2] - [anchor_num, 2] -> shape: [gt_box_num, anchor_num, 2]
    right_bottom = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])
    # clamp: 限制最小的输出 shape: [gt_box_num, anchor_num, 2] 2: 相交区域的w和h
    wh = (right_bottom - left_top).clamp(min=0)
    # wh[:, :, 0] 将每一个相交区域的宽取出 shape -> [gt_box_num, anchor_num]
    # wh[:, :, 1] 将每一个相交区域的高取出 shape -> [gt_box_num, anchor_num]
    # inter shape: [gt_box_num, anchor_num] 对应每个gtbox与生成的所有anchor之间相交的面积
    inter = wh[:, :, 0] * wh[:, :, 1]
    # iou shape: [gt_box_num, anchor_num]
    iou = inter / (area1[:, None] + area2 - inter)

    return iou


def permute_and_flatten(layer, N, A, C, H, W):
    # 调整tensor顺序 并进行展平操作
    # layer: 预测特征层上预测的目标概率或bboxes regression参数
    # N: batch_size
    # A: anchors的数量
    # C: classes_num or 4(bbox coordinate) RPN中C=1
    # H: height
    # W: width
    # C: class_num RPN只需计算前景还是背景 C=1 shape: [batch_size, anchor_num, C, H, W]
    layer = layer.view(N, -1, C, H, W)
    # shape: -> [batch_size, H, W, anchor_num, C]
    layer = layer.permute(0, 3, 4, 1, 2)
    # shape: -> [batch_size, H * W * anchor_num, C]
    # view和reshape功能是一样的，先展平所有元素在按照给定shape排列
    # view函数只能用于内存中连续存储的tensor，permute等操作会使tensor在内存中变得不再连续，此时就不能再调用view函数
    # reshape则不需要依赖目标tensor是否在内存中是连续的
    layer = layer.reshape(N, -1, C)

    return layer


def concat_box_pred_layers(box_cls, box_regression):

    # 存储预测目标分类数的参数
    box_cls_flattened = []
    # 存储预测bbox回归的参数
    box_regression_flattened = []

    # 逐层遍历每个预测特征图
    for box_cls_per_level, box_regression_per_level in zip(box_cls, box_regression):
        # N: batch_size A: anchor的个数 C: 分类数（在RPN中C=1，只区分目标和背景）H: 特征矩阵的高度 W: 特征矩阵的宽度
        N, AxC, H, W = box_cls_per_level.shape
        Ax4 = box_regression_per_level.shape[1]
        # anchor的个数
        A = Ax4 // 4
        # 分类的个数
        C = AxC // A

        # 对参数进行展平处理 shape: [batch_size, anchor_num*H*W， C=1]
        box_cls_per_level = permute_and_flatten(box_cls_per_level, N, A, C, H, W)
        box_cls_flattened.append(box_cls_per_level)

        # shape: [batch_size, anchor_num*H*W， C=4]
        box_regression_per_level = permute_and_flatten(box_regression_per_level, N, A, 4, H, W)
        box_regression_flattened.append(box_regression_per_level)
    # torch.cat(box_cls_flattened, dim=1) shape: [batch, anchor_num*H*W, C=1] 将不同预测特征图中生成的anchor数量加到一起
    # flatten(0, -2) shape: [Batch*anchor_num*H*W, 1] 其中参数表示从0维度开始，到-2维度停止 来进行展平操作。
    box_cls = torch.cat(box_cls_flattened, dim=1).flatten(0, -2)
    # shape: [Batch*anchor_num*H*W, 4]
    box_regression = torch.cat(box_regression_flattened, dim=1).reshape(-1, 4)

    return box_cls, box_regression


def clip_boxes_to_image(boxes, size):
    # 裁剪预测的boxes信息，将越界的坐标调整到图片边界上
    boxes_x = boxes[..., 0::2] # xmin xmax
    boxes_y = boxes[..., 1::2] # ymin ymax
    # size 对应当前图像的高宽信息
    height, width = size
    # clamp 限制信息。 将x坐标范围限制在0-width之间
    boxes_x = boxes_x.clamp(min=0, max=width)
    # 将y坐标范围限制在0 - height之间
    boxes_y = boxes_y.clamp(min=0, max=height)

    # 得到裁剪之后的boxes
    clipped_boxes = torch.cat((boxes_x, boxes_y), dim=1)
    return clipped_boxes


def remove_small_boxes(boxes, min_size):
    # proposal的宽高信息
    ws, hs = boxes[:, 2] - boxes[:, 0], boxes[:, 3] - boxes[:, 1]
    # 当满足宽，高都大于给定阈值时为True
    keep = (ws >= min_size) & (hs >= min_size)
    # 获取keep中为True的索引
    keep = keep.nonzero().squeeze(1)
    return keep


def batched_nms(boxes, scores, level_idxs, iou_threshold):
    # 判断proposal的数量是否为0
    if boxes.numel() == 0:
        return torch.empty((0, ), dtype=torch.int64, device=boxes.device)

    # 获取所有proposal中坐标最大的数值
    max_coordinate = boxes.max()
    # to: 使level_idxs的dtype和device与boxes保持一致
    offset = level_idxs.to(boxes) * (max_coordinate + 1)
    # 生成的proposal是在不同的预测特征图上生成的，加一个很大的offset，是为了将不同预测特征图的proposal分开，保证每一个预测特征图的proposal与其他层的proposal不会相交。
    # nms对不同预测特征图的proposal进行处理
    boxes_offset = boxes + offset[:, None]
    keep = nms(boxes_offset, scores, iou_threshold)

    return keep


def encode_boxes(reference_boxes, anchors, weights):
    wx = weights[0]
    wy = weights[1]
    ww = weights[2]
    wh = weights[3]

    # shape: [anchor_nums, 1]
    # 对应anchors的坐标 xmin ymin xmax ymax
    anchors_x1 = anchors[:, 0].unsqueeze(1)
    anchors_y1 = anchors[:, 1].unsqueeze(1)
    anchors_x2 = anchors[:, 2].unsqueeze(1)
    anchors_y2 = anchors[:, 3].unsqueeze(1)
    # 每个anchors对应的gtbox坐标 xmin ymin xmax ymax
    reference_boxes_x1 = reference_boxes[:, 0].unsqueeze(1)
    reference_boxes_y1 = reference_boxes[:, 1].unsqueeze(1)
    reference_boxes_x2 = reference_boxes[:, 2].unsqueeze(1)
    reference_boxes_y2 = reference_boxes[:, 3].unsqueeze(1)
    # anchors的宽度
    ex_width = anchors_x2 - anchors_x1
    # anchors的高度
    ex_height = anchors_y2 - anchors_y1
    # anchor的中心坐标x
    ex_ctr_x = anchors_x1 + 0.5 * ex_width
    # anchor的中心坐标y
    ex_ctr_y = anchors_y1 + 0.5 * ex_height

    # 每一个anchor对应gtbox的w h center_x center_y
    gt_widths = reference_boxes_x2 - reference_boxes_x1
    gt_heights = reference_boxes_y2 - reference_boxes_y1
    gt_ctr_x = reference_boxes_x1 + 0.5 * gt_widths
    gt_ctr_y = reference_boxes_y1 + 0.5 * gt_heights

    # 计算gtbox对应anchors的回归参数
    # ti*: 第i个anchor对应gtbox的回归参数.
    # tx* = (gtbox的x坐标 - anchor的x的坐标) / anchor的宽
    # ty* = (gtbox的y坐标 - anchor的y的坐标) / anchor的高
    # 得到gtbox相对anchor的xy坐标偏移量
    target_dx = wx * (gt_ctr_x - ex_ctr_x) / ex_width
    target_dy = wy * (gt_ctr_y - ex_ctr_y) / ex_height
    # tw* = log(gtbox的w / anchor的w)
    # th* = log(gtbox的h / anchor的h)
    target_dw = ww * torch.log(gt_widths / ex_width)
    target_dh = wh * torch.log(gt_heights / ex_height)
    # xywh -> xmin ymin xmax ymax
    pred_boxes_xmin = target_dx - torch.tensor(0.5, dtype=target_dx.dtype, device=target_dx.device) * target_dw
    pred_boxes_ymin = target_dy - torch.tensor(0.5, dtype=target_dy.dtype, device=target_dy.device) * target_dh
    pred_boxes_xmax = target_dx + torch.tensor(0.5, dtype=target_dx.dtype, device=target_dx.device) * target_dw
    pred_boxes_ymax = target_dy + torch.tensor(0.5, dtype=target_dy.dtype, device=target_dy.device) * target_dh
    # shape: [anchors, 4] 4为xywh的偏移量
    # targets = torch.cat((target_dx, target_dy, target_dw, target_dh), dim=1)
    # shape: [anchors, 4] 4为xyxy的偏移量
    targets = torch.cat((pred_boxes_xmin, pred_boxes_ymin, pred_boxes_xmax, pred_boxes_ymax), dim=1)

    return targets


class box_Coder(object):

    def __init__(self, weights, bbox_xform_clip=math.log(1000. / 16)):

        self.weights = weights
        self.bbox_xform_clip = bbox_xform_clip

    def encode(self, reference_boxes, anchores):
        # 结合anchors和与之对应的gt计算回归参数
        # reference_boxes: 每个anchor对应的gtbox坐标 proposals: 对应anchor的坐标
        # 统计每张图像的anchors个数
        boxes_per_image = [len(b) for b in reference_boxes]
        # shape: [batch_size*per_anchor_num, 4]
        # 每个anchor对应的gtbox坐标进行拼接
        reference_boxes = torch.cat(reference_boxes, dim=0)
        # shape: [batch_size*per_anchor_num, 4]
        # 每个anchors进行拼接
        anchors = torch.cat(anchores, dim=0)

        targets = self.encode_single(reference_boxes, anchors)
        # shape: [per_anchor_num, 4]
        return targets.split(boxes_per_image, 0)

    def encode_single(self, reference_boxes, anchors):
        dtype = reference_boxes.dtype
        device = reference_boxes.device
        weights = torch.as_tensor(self.weights, dtype=dtype, device=device) # weights 1 1 1 1
        # shape: [anchors, 4] 4: center_x, center_y, w, h的偏移量
        # targets = encode_boxes(reference_boxes, anchors, weights)
        # shape: [anchors, 4] 4: xmin, ymin, xmax, ymax的偏移量
        targets = encode_boxes(reference_boxes, anchors, weights)

        return targets

    def decode_single(self, rel_codes, boxes):
        # rel_codes: 预测的bbox的回归参数
        # boxes: anchor模板的坐标
        boxes = boxes.to(rel_codes.dtype)

        # boxes(cell_anchors) shape: [batchsize*anchor_num, 4] 4: xmin, ymin, xmax, ymax
        width = boxes[:, 2] - boxes[:, 0] # shape: [batchsize*anchor_num]
        height = boxes[:, 3] - boxes[:, 1]
        center_x = boxes[:, 0] + 0.5 * width
        center_y = boxes[:, 1] + 0.5 * height

        wx, wy, ww, wh = self.weights
        # 0::4 从0开始以4为间隔进行采样 但可以多保留一个维度 dx shape: [batchsize*anchor_num, 1]
        # rel_codes shape: [batchsize*anchor_num, 4]
        dx = rel_codes[:, 0::4] / wx # 预测anchor的中心坐标回归参数x
        dy = rel_codes[:, 1::4] / wy # 预测anchor的中心坐标回归参数y
        dw = rel_codes[:, 2::4] / ww # 预测anchor的宽度回归参数w
        dh = rel_codes[:, 3::4] / wh # 预测anchor的高度回归参数h

        # clamp: 限制数值的上下限
        # 预防指数爆炸
        dw = torch.clamp(dw, max=self.bbox_xform_clip)
        dh = torch.clamp(dh, max=self.bbox_xform_clip)

        # 将预测值应用于anchor模板中
        # tx: 预测得到的中心坐标x的回归参数 wa: anchor的宽度 xa: anchor的中心x坐标
        # tx = (x - xa) / wa        ty = (y - ya) / ha
        # tw = log(w / wa)          th = log(h / ha)
        # x y w h 为预测值xywh
        pred_center_x = dx * width[:, None] + center_x[:, None]
        pred_center_y = dy * height[:, None] + center_y[:, None]
        pred_w = torch.exp(dw) * width[:, None]
        pred_h = torch.exp(dh) * height[:, None]

        # [x, y, w, h] -> [xmin, ymin, xmax. ymax]
        pred_boxes_xmin = pred_center_x - torch.tensor(0.5, dtype=pred_center_x.dtype, device=pred_w.device) * pred_w
        pred_boxes_ymin = pred_center_y - torch.tensor(0.5, dtype=pred_center_x.dtype, device=pred_w.device) * pred_h
        pred_boxes_xmax = pred_center_x + torch.tensor(0.5, dtype=pred_center_x.dtype, device=pred_w.device) * pred_w
        pred_boxes_ymax = pred_center_y + torch.tensor(0.5, dtype=pred_center_x.dtype, device=pred_w.device) * pred_h

        # stack: dim=2 首先新增一个维度 flatten()展平 ->  shape: [batchsize*anchor_num, 4]
        # stack: 会新增维度 cat: 在指定维度上进行拼接
        pred_boxes = torch.stack((pred_boxes_xmin, pred_boxes_ymin, pred_boxes_xmax, pred_boxes_ymax), dim=2).flatten(1)
        # pred_boxes = torch.cat((pred_boxes_xmin, pred_boxes_ymin, pred_boxes_xmax, pred_boxes_ymax), dim=1)

        return pred_boxes

    def decode(self, rel_codes, boxes):
        # rel_codes: 预测的bbox的回归参数
        # boxes: anchor的坐标
        # 将一个batch中所有anchor的坐标信息拼接在一起
        concat_boxes = torch.cat(boxes, dim=0)

        # 获取一个batch中anchor的总数
        box_sum = concat_boxes.shape[0]

        # 将预测的bbox的回归参数应用到anchor上
        pred_boxes = self.decode_single(rel_codes, concat_boxes)
        # shape: [batchsize*anchor_num, 1, 4]
        pred_boxes = pred_boxes.reshape(box_sum, -1, 4)

        return pred_boxes


class Matcher(object):
    # 计算anchors与每个gtboxes匹配的iou最大值，并记录索引，
    def __init__(self,  high_threshold, low_threshold, allow_low_quality_matches=False):

        self.BELOW_LOW_THRESHOLD = -1
        self.high_threshold = high_threshold  # 0.7
        self.low_threshold = low_threshold  # 0.3
        self.allow_low_quality_matches = allow_low_quality_matches

    def __call__(self, match_quality_matrix):
        # matched_vals: 每一个anchor对应每个GTbox上最大的iou的值, 即每列的最大值。
        # matches: 每列最大值的索引
        # 获得anchors与GTbox最大的iou值
        matched_value, matches_idx = match_quality_matrix.max(dim=0)

        if self.allow_low_quality_matches:
            all_matches = matches_idx.clone()
        else:
            all_matches = None

        # 计算iou小于low_threshold. 小于的位置为True, 不小于的位置为False
        below_low_threshold = matched_value < self.low_threshold

        # 计算iou大于low_threshold且小于high_threshold之间的索引
        between_threshold = (matched_value >= self.low_threshold) & (matched_value < self.high_threshold)

        # 将matches_idx中小于low_threshold位置的值，设置为-1
        matches_idx[below_low_threshold] = -1

        # iou在[low_threshold, high_threshold]之间的matches_idx索引值设为-2
        matches_idx[between_threshold] = -2

        # 对于每一个GTbox而言，与之匹配IOU值最大的anchor，也将其设置为正样本。
        # 相当于对于每一个GTbox都有一个与之匹配的anchor，将其设置为正样本，不局限于iou大于0.7时才为正样本。
        if self.allow_low_quality_matches:
            self.set_low_quality_matches_(matches_idx, all_matches, match_quality_matrix)

        return matches_idx

    def set_low_quality_matches_(self, matches_idx, all_matches, match_quality_matrix):
        # 对于每个GTbox，寻找与其iou值最大的anchor. 对于每一行而言，每一行为一个GTbox
        highest_quality_gt_value, _ = match_quality_matrix.max(dim=1)
        # 在相同位置上 数值相等的位置为True. nonzero: 寻找为True的位置. 返回每一个非0元素的坐标
        gt_anchor_matches_highest_coordiate = torch.nonzero(match_quality_matrix == highest_quality_gt_value[:, None])
        gt_anchor_matches_highest_coordiate_update = gt_anchor_matches_highest_coordiate[:, 1]
        # 保留该anchor匹配gt最大iou的索引，即使iou低于设定的阈值
        matches_idx[gt_anchor_matches_highest_coordiate_update] = all_matches[gt_anchor_matches_highest_coordiate_update]


class BalancedPositiveNegativeSampler(object):

    def __init__(self, batch_size_per_image, positive_fraction):

        self.batch_size_per_image = batch_size_per_image
        self.positive_fraction = positive_fraction

    def __call__(self, matched_idxs):
        # 记录正负样本的索引
        pos_idx = []
        neg_idx = []
        # 遍历每张图像的matched_idxs 划分为正、负和丢弃样本的labels
        for matched_idxs_per_image in matched_idxs:
            # >= 1 positive sample, nonzero返回非零元素的索引(正样本对应的索引)
            # torch.nonzero(matched_idxs_per_image >= 1) shape: [正样本对应的索引, 1]
            # squeeze(1): 降维 -> shape: [正样本对应的索引]
            positive = torch.nonzero(matched_idxs_per_image >= 1).squeeze(1)
            # = 0的为负样本
            negative = torch.nonzero(matched_idxs_per_image == 0).squeeze(1)
            # 指定正样本数量
            num_pos = int(self.batch_size_per_image * self.positive_fraction)
            # 若正样本数量不足，直接使用所有正样本
            num_pos = min(positive.numel(), num_pos)
            # 负样本数量
            num_neg = self.batch_size_per_image - num_pos
            # 如果负样本数量不够就直接采用所有负样本 numel: 元素个数
            num_neg = min(negative.numel(), num_neg)

            # 随机选择指定数量的正负样本 返回的是 索引信息
            perm1 = torch.randperm(positive.numel(), device=positive.device)[:num_pos]
            perm2 = torch.randperm(negative.numel(), device=negative.device)[:num_neg]

            pos_idx_per_image = positive[perm1]
            neg_idx_per_image = negative[perm2]

            # 创建两个全0的模板 用于存储正样本和负样本
            pos_idx_per_image_mask = torch.zeros_like(
                matched_idxs_per_image, dtype=torch.uint8
            )
            neg_idx_per_image_mask = torch.zeros_like(
                matched_idxs_per_image, dtype=torch.uint8
            )
            # 对于正样本模板 对应正样本的索引位置处的值 全部设置为1
            pos_idx_per_image_mask[pos_idx_per_image] = 1
            # 对于负样本模板 对应负样本的索引位置处的值 全部设置为1
            neg_idx_per_image_mask[neg_idx_per_image] = 1

            pos_idx.append(pos_idx_per_image_mask)
            neg_idx.append(neg_idx_per_image_mask)

        return pos_idx, neg_idx


class RegionProposalNetwork(nn.Module):
    # batch_size_per_image：RPN计算损失时采用正负样本的总个数 positive_fraction：正样本在总样本中的比例
    def __init__(self, anchor_generate, rpn_head, fg_iou_thresh, bg_iou_thresh, batch_size_per_image,
                 positive_fraction, pre_nms_top_n, post_nms_top_n, nms_thresh):
        super(RegionProposalNetwork, self).__init__()

        self.anchor_generator = anchor_generate
        self.head = rpn_head
        # 平衡系数
        self.box_coder = box_Coder(weights=(1.0, 1.0, 1.0, 1.0))

        # train
        self.box_similarity = box_iou
        # fg_iou_thresh: 当iou大于fg_iou_thresh(0.7)，设置为正样本. bg_iou_thresh: 当iou小于bg_iou_thresh(0.3)，设置为负样本.
        self.proposal_matcher = Matcher(fg_iou_thresh, bg_iou_thresh, allow_low_quality_matches=True)

        # batch_size_per_image: 计算损失时，选择的正负样本的总个数
        # positive_fraction: 正样本占全部样本的比例
        self.fg_bg_sampler = BalancedPositiveNegativeSampler(
            batch_size_per_image, positive_fraction  # 256, 0.5
        )

        # test
        self._pre_nms_top_n = pre_nms_top_n
        self._post_nms_top_n = post_nms_top_n
        self.nms_thresh = nms_thresh
        # 过滤proposal时使用
        self.min_size = 1e-3

    def pre_nms_top_n(self):
        # 判断训练还是预测. 训练 pre_nms_top_n = 2000 预测: pre_nms_top_n = 1000
        if self.training:
            return self._pre_nms_top_n['training']
        return self._pre_nms_top_n['testing']

    def post_nms_top_n(self):
        if self.training:
            return self._post_nms_top_n['training']
        return self._post_nms_top_n['testing']

    def _get_top_n_idx(self, objectness, num_anchors_per_level):

        # objectness: 一个batch中每张图像的预测目标概率信息
        # num_anchors_per_level: 每个预测特征图上预测anchor的数量
        # 记录每个预测特征图上预测目标概率前pre_nms_top_n的索引信息
        r = []
        offset = 0
        # 遍历每个预测特征图上的预测目标概率信息
        # split: 在指定维度上 信息按多长进行分割。 对应每一个预测特征图上anchor的个数
        for ob in objectness.split(num_anchors_per_level, 1):
            # ob shape: [batch_size, 对应预测特征图上anchor的数量]
            num_anchors = ob.shape[1]
            # 针对每一层取前topn个anchors，得到proposal.
            pre_nms_top_n = min(self.pre_nms_top_n(), num_anchors)
            # topk. 默认从大到小排序。 return：1、排序后的数值 2、排序后的index
            # 对每一层预测特征图进行排序
            _, top_n_idx = ob.topk(pre_nms_top_n, dim=1)
            # offset. 目的是获得将所有anchor(所有预测特征图中的anchor)合并到一起后的索引， 得到proposal。
            r.append(top_n_idx + offset)
            offset += num_anchors
        # 将每一层的前topk的索引进行拼接
        # r中每一层预测特征图的shape: [batch_size, topk]
        return torch.cat(r, dim=1)

    def assign_targets_to_anchors(self, anchors, targets):
        # 对于一个batch中，每一张图像中生成的anchors，为每一个anchors匹配所对应的标签以及GTboxes。

        # 记录anchors匹配的标签以及gtboxes所对应的坐标
        labels = []
        matched_gt_boxes = []

        # 遍历每张图像的anchors和target
        for anchors_per_image, targets_per_image in zip(anchors, targets):
            # targets为dict，分别对应boxes, labels, image_id, area, iscrowd信息。
            # RPN只区别目标是前景和背景 因此只提取坐标信息，不需要具体的类别信息。
            gt_boxes = targets_per_image['boxes']
            # numel(): 获取所有元素的个数
            # 如果图像中没有检测目标(gtbox)
            if gt_boxes.numel() == 0:
                device = anchors_per_image.device
                matched_gt_boxes_per_image = torch.zeros(anchors_per_image.shape, dtype=torch.float32, device=device)
                labels_per_image = torch.zeros((anchors_per_image.shape[0],), dtype=torch.float32, device=device)

            else:
                # 计算GTboxes与anchors的IOU
                # return shape: [gt_num, anchor_num]
                match_quality_matrix = box_iou(gt_boxes, anchors_per_image)
                # 通过计算得到的iou值，为每个anchor分配匹配到的gtbox的索引。
                # shape: [anchor_num]. 由-1 -2 anchor的索引组成. -1为负样本 -2为丢弃的样本 >=0为正样本
                matched_idxs = self.proposal_matcher(match_quality_matrix)
                # 获得每一个anchor得到的gtbox的坐标
                matched_gt_boxes_per_image = gt_boxes[matched_idxs.clamp(min=0)]
                # 记录所有anchors匹配后的标签(所有正样本的位置) 正样本对应索引位置上的值为1
                labels_per_image = matched_idxs >= 0
                labels_per_image = labels_per_image.to(dtype=torch.float32)
                # 记录所有负样本的索引 负样本对应索引位置上的值设为0
                bg_indices = matched_idxs == -1
                labels_per_image[bg_indices] = 0.0
                # 记录丢弃样本的索引 丢弃样本对应索引位置上的值设为-1
                between_indices = matched_idxs == -2
                labels_per_image[between_indices] = -1.0

            labels.append(labels_per_image)
            matched_gt_boxes.append(matched_gt_boxes_per_image)

        return labels, matched_gt_boxes

    def filter_proposals(self, proposals, objectness, image_shapes, num_anchors_per_level):
        # proposal shape: [Batch, anchor_num, 4]
        # batch size
        num_images = proposals.shape[0]
        device = proposals.device

        # detach 对于fast_rcnn部分，proposal是输入参数，丢弃计算梯度信息，只保留数值信息，requires_grad=False.
        objectness = objectness.detach()
        # shape: [batch_size, anchor_num]
        objectness = objectness.view(num_images, -1)
        # levels 记录不同预测特征层上的anchors索引信息。 为了知道对于每一个anchor其属于哪一个预测特征图上的
        # idx: 对应预测特征图的索引 n: 该预测特征图上anchor的个数
        # torch.full生成维度为n，其中用idx值填充
        levels = [torch.full((n,), idx, dtype=torch.int64, device=device) for idx, n in enumerate(num_anchors_per_level)]
        # [anchor_num]
        levels = torch.cat(levels, dim=0)
        # reshape shape: [1, anchor_num]
        # expand_as: [batch, anchor_num]
        levels = levels.reshape(1, -1).expand_as(objectness)

        # 获取每张预测特征图上预测概率排在pre_nms_top_n的anchor的索引值，得到proposals.
        # top_n_idx shape: [batch_size, proposal_num]
        top_n_idx = self._get_top_n_idx(objectness, num_anchors_per_level)

        # batch_size的个数
        image_range = torch.arange(num_images, device=device)
        # shape: [batch_size, 1]
        batch_idx = image_range[:, None]

        # 根据每个预测特征图中排在前pre_nms_top_n个anchor的索引值，得到其相应的概率信息，从而得到proposal。
        # 在objectness中，每一张图像中对应的topn的索引，得到其概率信息。
        # objectness shape: [batch_size, proposal_num]
        objectness = objectness[batch_idx, top_n_idx]
        # 存储anchor属于某个预测特征层的信息
        levels = levels[batch_idx, top_n_idx]
        # 预测概率排在pre_nms_top_n个anchor的索引值，得到相应的bbox坐标信息，得到proposal
        proposals = proposals[batch_idx, top_n_idx]

        final_boxes = []
        final_scores = []

        # 遍历每张图像上的相关预测信息
        for boxes, scores, level, img_shape in zip(proposals, objectness, levels, image_shapes):
            # 调整预测的boxes坐标，得到最终的proposal坐标。目的是将越界的坐标调整到图像边界上，限制在图像的内部。
            boxes = clip_boxes_to_image(boxes, img_shape)
            # 移除proposal中的小boxes 根据self.min_size
            # 返回boxes满足宽，高都大于self.min_size的索引
            keep = remove_small_boxes(boxes, self.min_size)
            boxes, scores, level = boxes[keep], scores[keep], level[keep]
            # 经过nms后保留下来的proposal.
            # return: 执行完nms后，按照目标的类别scores，从大到小排序输出的索引信息。
            keep = batched_nms(boxes, scores, level, self.nms_thresh)
            # 只获取post_nms_top_n个proposal
            keep = keep[: self.post_nms_top_n()]
            boxes, scores = boxes[keep], scores[keep]

            final_boxes.append(boxes)
            final_scores.append(scores)

        return final_boxes, final_scores

    def compute_loss(self, objectness, pred_bbox_deltas, labels, regression_targets):
        # 按照给定的batch_size_per_image, positive_fraction选择正负样本
        # labels: 完成对正负样本的划分。
        # fg_bg_sampler: 不是所有的正负样本都拿来训练，选择在计算损失时所使用到的正负样本。
        sampled_pos_inds, sampled_neg_inds = self.fg_bg_sampler(labels)
        # 将一个batch中的所有正负样本分别拼接在一起，并获取非零位置的索引
        sampled_pos_inds = torch.nonzero(torch.cat(sampled_pos_inds, dim=0)).squeeze(1)
        sampled_neg_inds = torch.nonzero(torch.cat(sampled_neg_inds, dim=0)).squeeze(1)

        # 将所有正负样本索引拼接在一起
        sampled_inds = torch.cat([sampled_pos_inds, sampled_neg_inds], dim=0)
        objectness = objectness.flatten()

        labels = torch.cat(labels, dim=0)
        regression_targets = torch.cat(regression_targets, dim=0)

        # 计算边界框回归损失 只需要计算正样本的损失
        box_loss = smooth_l1_loss(pred_bbox_deltas[sampled_pos_inds], regression_targets[sampled_pos_inds], beta=1 / 9, size_average=False) / (sampled_inds.numel())
        # 计算目标预测概率损失
        objectness_loss = F.binary_cross_entropy_with_logits(
            objectness[sampled_inds], labels[sampled_inds]
        )

        return objectness_loss, box_loss

    def forward(self, image_list, features, targets=None):
        # 首先提取所有预测特征图的特征矩阵。features为字典类型，只提取value(特征矩阵)，不要key.
        # features是所有预测特征层组成的OrderedDict
        features = list(features.values())

        # 将预测特征图输入到RPNHead中
        # 计算每个预测特征层上的预测目标概率和bboxes regression参数
        # objectness和pred_bbox_deltas都是list
        # objectness: 根据预测特征层预测每一个anchor属于前景还是背景。 shape: [B, 3(每个预测特征图中anchor的个数), H, W]
        # pred_bbox_deltas: 根据预测特征层输出坐标偏移量 shape: [B, 3*4, H, W]
        objectness, pred_bbox_deltas = self.head(features)

        # 得到一个batch中对应每一张图像生成的所有anchors信息
        anchors = self.anchor_generator(image_list, features)

        # batch_size
        num_images = len(anchors)

        # 计算每个预测特征图上anchor的数量 shape: [C, H, W]
        num_anchors_per_level_shape_tensor = [o[0].shape for o in objectness]
        # 每一个预测特征层上所生成anchor的个数
        num_anchors_per_level = [s[0] * s[1] * s[2] for s in num_anchors_per_level_shape_tensor]

        # 调整box_cls和box_reg两个list中的每个预测特征层的预测信息的tensor排列顺序和shape进行调整
        # objectness shape -> [Batch*anchor_num*H*W, 1]
        # pred_bbox_deltas shape -> [Batch*anchor_num*H*W, 4]
        objectness, pred_bbox_deltas = concat_box_pred_layers(objectness, pred_bbox_deltas)

        # 将预测得到的bbox参数应用到anchor中，得到proposal
        # detach(): 不计算梯度
        proposals = self.box_coder.decode(pred_bbox_deltas.detach(), anchors)
        # shape: [Batch*anchor_num, 4] -> [Batch, anchor_num, 4]
        proposals = proposals.view(num_images, -1, 4)

        # 过滤proposal. 筛出小面积的proposal框, nms处理，根据预测概率取前post_nms_top_n个目标
        # proposal shape: [Batch, anchor_num, 4]
        boxes, scores = self.filter_proposals(proposals, objectness, image_list.image_sizes, num_anchors_per_level)

        # 在训练期间，计算loss
        losses = {}
        if self.training:
            # 计算每个anchor最匹配的GTbox，并将anchor分类成 前景、背景、丢弃 用1，0，-1表示.
            labels, matched_gt_boxes = self.assign_targets_to_anchors(anchors, targets)
            # 计算gtbox坐标相对于anchor坐标之间的回归参数
            regression_targets = self.box_coder.encode(matched_gt_boxes, anchors)

            # objectness: 预测的目标分数 pred_bbox_deltas: 预测的坐标偏移量
            # labels: 真实的label(正样本1 负样本0 丢弃样本-1) regression_targets: gtbox坐标相对于anchor坐标之间的回归参数
            loss_objectness, loss_rpn_box_reg = self.compute_loss(objectness, pred_bbox_deltas, labels, regression_targets)

            losses = {
                'loss_objectness': loss_objectness,
                'loss_rpn_box_reg': loss_rpn_box_reg
            }
        return boxes, losses