def assign_targets_to_proposals(self, proposals, gt_boxes, gt_labels):
# type: (List[Tensor], List[Tensor], List[Tensor]) -> Tuple[List[Tensor], List[Tensor]]
"""
为每个proposal匹配对应的gt_box，并划分到正负样本中
Args:
proposals:
gt_boxes:
gt_labels:
Returns:
"""
matched_idxs = []
labels = []
# 遍历每张图像的proposals, gt_boxes, gt_labels信息
for proposals_in_image, gt_boxes_in_image, gt_labels_in_image in zip(
proposals, gt_boxes, gt_labels):
if gt_boxes_in_image.numel() == 0: # 该张图像中没有gt框，为背景
# background image
device = proposals_in_image.device
clamped_matched_idxs_in_image = torch.zeros(
(proposals_in_image.shape[0], ),
dtype=torch.int64,
device=device)
labels_in_image = torch.zeros((proposals_in_image.shape[0], ),
dtype=torch.int64,
device=device)
else:
# set to self.box_similarity when https://github.com/pytorch/pytorch/issues/27495 lands
# 计算proposal与每个gt_box的iou重合度
match_quality_matrix = box_ops.box_iou(gt_boxes_in_image,
proposals_in_image)
# 计算proposal与每个gt_box匹配的iou最大值，并记录索引，
# iou < low_threshold索引值为 -1， low_threshold <= iou < high_threshold索引值为 -2
matched_idxs_in_image = self.proposal_matcher(
match_quality_matrix)
# 限制最小值，防止匹配标签时出现越界的情况
# 注意-1, -2对应的gt索引会调整到0,获取的标签类别为第0个gt的类别（实际上并不是）,后续会进一步处理
clamped_matched_idxs_in_image = matched_idxs_in_image.clamp(
min=0)
# 获取proposal匹配到的gt对应标签
labels_in_image = gt_labels_in_image[
clamped_matched_idxs_in_image]
labels_in_image = labels_in_image.to(dtype=torch.int64)
# label background (below the low threshold)
# 将gt索引为-1的类别设置为0，即背景，负样本
bg_inds = matched_idxs_in_image == self.proposal_matcher.BELOW_LOW_THRESHOLD # -1
labels_in_image[bg_inds] = 0
# label ignore proposals (between low and high threshold)
# 将gt索引为-2的类别设置为-1, 即废弃样本
ignore_inds = matched_idxs_in_image == self.proposal_matcher.BETWEEN_THRESHOLDS # -2
labels_in_image[ignore_inds] = -1 # -1 is ignored by sampler
matched_idxs.append(clamped_matched_idxs_in_image)
labels.append(labels_in_image)
return matched_idxs, labels
def assign_targets_to_anchors(self, anchors, targets):
# type: (List[Tensor], List[Dict[str, Tensor]]) -> Tuple[List[Tensor], List[Tensor]]
"""
计算每个anchors最匹配的gt，并划分为正样本，背景以及废弃的样本
Args：
anchors: (List[Tensor])
targets: (List[Dict[Tensor])
Returns:
labels: 标记anchors归属类别（1, 0, -1分别对应正样本，背景，废弃的样本）
注意，在RPN中只有前景和背景，所有正样本的类别都是1，0代表背景
matched_gt_boxes：与anchors匹配的gt
"""
labels = []
matched_gt_boxes = []
# 遍历每张图像的anchors和targets
for anchors_per_image, targets_per_image in zip(anchors, targets):
# for anchors_per_image, targets_per_image in zip(anchors, [targets[i] for i in range(targets.shape[0])]):
gt_boxes = targets_per_image["boxes"]
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:
# 计算anchors与真实bbox的iou信息
# set to self.box_similarity when https://github.com/pytorch/pytorch/issues/27495 lands
match_quality_matrix = box_ops.box_iou(gt_boxes,
anchors_per_image)
# 计算每个anchors与gt匹配iou最大的索引（如果iou<0.3索引置为-1，0.3<iou<0.7索引为-2）
matched_idxs = self.proposal_matcher(match_quality_matrix)
# get the targets corresponding GT for each proposal
# NB: need to clamp the indices because we can have a single
# GT in the image, and matched_idxs can be -2, which goes
# out of bounds
matched_gt_boxes_per_image = gt_boxes[matched_idxs.clamp(
min=0)]
labels_per_image = matched_idxs >= 0
labels_per_image = labels_per_image.to(dtype=torch.float32)
# background (negative examples)
bg_indices = matched_idxs == self.proposal_matcher.BELOW_LOW_THRESHOLD # -1
labels_per_image[bg_indices] = 0.0
# discard indices that are between thresholds
inds_to_discard = matched_idxs == self.proposal_matcher.BETWEEN_THRESHOLDS # -2
labels_per_image[inds_to_discard] = -1.0
labels.append(labels_per_image)
matched_gt_boxes.append(matched_gt_boxes_per_image)
return labels, matched_gt_boxes
