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