def forward(self, locations, box_cls, box_regression, centerness, proposal_embed, proposal_margin, pixel_embed, image_sizes, targets): """ Arguments: anchors: list[list[BoxList]] box_cls: list[tensor] box_regression: list[tensor] image_sizes: list[(h, w)] Returns: boxlists (list[BoxList]): the post-processed anchors, after applying box decoding and NMS """ sampled_boxes = [] for i, (l, o, b, c) in enumerate(zip(locations, box_cls, box_regression, centerness)): em = proposal_embed[i] mar = proposal_margin[i] if self.fix_margin: mar = torch.ones_like(mar) * self.init_margin sampled_boxes.append( self.forward_for_single_feature_map( l, o, b, c, em, mar, image_sizes, i ) ) boxlists = list(zip(*sampled_boxes)) boxlists = [cat_boxlist(boxlist) for boxlist in boxlists] boxlists = self.select_over_all_levels(boxlists) # resize pixel embedding for higher resolution N, dim, m_h, m_w = pixel_embed.shape o_h = m_h * self.mask_scale_factor o_w = m_w * self.mask_scale_factor pixel_embed = interpolate(pixel_embed, size=(o_h, o_w), mode='bilinear', align_corners=False) boxlists = self.forward_for_mask(boxlists, pixel_embed) return boxlists
def forward(self, anchors, objectness, box_regression, targets=None): """ Arguments: anchors: list[list[BoxList]] objectness: list[tensor] box_regression: list[tensor] Returns: boxlists (list[BoxList]): the post-processed anchors, after applying box decoding and NMS """ sampled_boxes = [] num_levels = len(objectness) anchors = list(zip(*anchors)) for a, o, b in zip(anchors, objectness, box_regression): sampled_boxes.append(self.forward_for_single_feature_map(a, o, b)) boxlists = list(zip(*sampled_boxes)) boxlists = [cat_boxlist(boxlist) for boxlist in boxlists] if num_levels > 1: boxlists = self.select_over_all_levels(boxlists) # append ground-truth bboxes to proposals if self.training and targets is not None: boxlists = self.add_gt_proposals(boxlists, targets) return boxlists
def forward(self, locations, box_cls, box_regression, centerness, box_mask, image_sizes): """ Arguments: anchors: list[list[BoxList]] box_cls: list[tensor] box_regression: list[tensor] box_mask: list[tensor] image_sizes: list[(h, w)] Returns: boxlists (list[BoxList]): the post-processed anchors, after applying box decoding and NMS """ sampled_boxes = [] for _, (l, o, b, c, m) in enumerate( zip(locations, box_cls, box_regression, centerness, box_mask)): sampled_boxes.append( self.forward_for_single_feature_map(l, o, b, c, m, image_sizes)) boxlists = list(zip(*sampled_boxes)) boxlists = [cat_boxlist(boxlist) for boxlist in boxlists] boxlists = self.select_over_all_levels(boxlists) return boxlists
def forward(self, locations, box_cls, box_regression, centerness, image_sizes, benchmark=False, timers=None): """ Arguments: anchors: list[list[BoxList]] box_cls: list[tensor] box_regression: list[tensor] image_sizes: list[(h, w)] Returns: boxlists (list[BoxList]): the post-processed anchors, after applying box decoding and NMS """ if benchmark and timers is not None: torch.cuda.synchronize() timers[2].tic() sampled_boxes = [] for _, (l, o, b, c) in enumerate(zip(locations, box_cls, box_regression, centerness)): sampled_boxes.append( self.forward_for_single_feature_map( l, o, b, c, image_sizes ) ) if benchmark and timers is not None: torch.cuda.synchronize() timers[2].toc() timers[3].tic() boxlists = list(zip(*sampled_boxes)) boxlists = [cat_boxlist(boxlist) for boxlist in boxlists] if not self.bbox_aug_enabled: boxlists = self.select_over_all_levels(boxlists) if benchmark and timers is not None: torch.cuda.synchronize() timers[3].toc() return boxlists
def select_over_all_levels(self, boxlists): num_images = len(boxlists) results = [] for i in range(num_images): scores = boxlists[i].get_field("scores") labels = boxlists[i].get_field("labels") diagonals = boxlists[i].get_field("diagonals") boxes = boxlists[i].bbox boxlist = boxlists[i] result = [] # skip the background for j in tqdm(range(1, self.num_classes)): inds = (labels == j).nonzero().view(-1) scores_j = scores[inds] boxes_j = boxes[inds, :].view(-1, 4) diagonals_j = diagonals[inds, :].view(-1, 16) boxlist_for_class = BoxList(boxes_j, boxlist.size, mode="xyxy") boxlist_for_class.add_field("scores", scores_j) boxlist_for_class.add_field("diagonals", diagonals_j) if scores_j.size()[0] != 0: boxlist_for_class = diagonal_nms(boxlist_for_class, self.nms_thresh, self.nms_topk, score_field="scores") num_labels = len(boxlist_for_class) boxlist_for_class.add_field( "labels", torch.full((num_labels, ), j, dtype=torch.int64, device=scores.device)) result.append(boxlist_for_class) result = cat_boxlist(result) # new_result = all_class_nms(result, self.nms_thresh) number_of_detections = len(result) print('Number of detections in this image {}'.format( number_of_detections)) # Limit to max_per_image detections **over all classes** if number_of_detections > self.fpn_post_nms_top_n > 0: cls_scores = result.get_field("scores") image_thresh, _ = torch.kthvalue( cls_scores.cpu(), number_of_detections - self.fpn_post_nms_top_n + 1) # image_thresh = torch.as_tensor(0.5) keep = cls_scores >= image_thresh.item() keep = torch.nonzero(keep).squeeze(1) result = result[keep] result = all_class_nms(result, self.nms_thresh) results.append(result) return results
def cat_boxlist_with_keypoints(boxlists): assert all(boxlist.has_field("keypoints") for boxlist in boxlists) kp = [boxlist.get_field("keypoints").keypoints for boxlist in boxlists] kp = cat(kp, 0) fields = boxlists[0].get_fields() fields = [field for field in fields if field != "keypoints"] boxlists = [boxlist.copy_with_fields(fields) for boxlist in boxlists] boxlists = cat_boxlist(boxlists) boxlists.add_field("keypoints", kp) return boxlists
def __call__(self, box_regression, centerness, targets, anchors): labels, reg_targets = self.prepare_targets(targets, anchors) N = len(labels) box_regression_flatten = self.concat_box_prediction_layers( box_regression) centerness_flatten = [ ct.permute(0, 2, 3, 1).reshape(N, -1, 1) for ct in centerness ] centerness_flatten = torch.cat(centerness_flatten, dim=1).reshape(-1) labels_flatten = torch.cat(labels, dim=0) reg_targets_flatten = torch.cat(reg_targets, dim=0) anchors_flatten = torch.cat([ cat_boxlist(anchors_per_image).bbox for anchors_per_image in anchors ], dim=0) pos_inds = torch.nonzero(labels_flatten > 0).squeeze(1) num_gpus = get_num_gpus() total_num_pos = reduce_sum(pos_inds.new_tensor([pos_inds.numel() ])).item() num_pos_avg_per_gpu = max(total_num_pos / float(num_gpus), 1.0) box_regression_flatten = box_regression_flatten[pos_inds] reg_targets_flatten = reg_targets_flatten[pos_inds] anchors_flatten = anchors_flatten[pos_inds] centerness_flatten = centerness_flatten[pos_inds] centerness_targets = self.compute_centerness_targets( reg_targets_flatten, anchors_flatten) sum_centerness_targets_avg_per_gpu = reduce_sum( centerness_targets.sum()).item() / float(num_gpus) if pos_inds.numel() > 0: reg_loss = self.GIoULoss( box_regression_flatten, reg_targets_flatten, anchors_flatten, weight=centerness_targets) / sum_centerness_targets_avg_per_gpu centerness_loss = self.centerness_loss_func( centerness_flatten, centerness_targets) / num_pos_avg_per_gpu else: reg_loss = box_regression_flatten.sum() centerness_loss = centerness_flatten.sum() return reg_loss * self.cfg.MODEL.ATSS.REG_LOSS_WEIGHT, centerness_loss
def select_over_all_levels(self, boxlists): num_images = len(boxlists) results = [] for i in range(num_images): scores = boxlists[i].get_field("scores") labels = boxlists[i].get_field("labels") masks = boxlists[i].get_field("mask") boxes = boxlists[i].bbox boxlist = boxlists[i] result = [] # skip the background for j in range(1, self.num_classes): inds = (labels == j).nonzero().view(-1) scores_j = scores[inds] boxes_j = boxes[inds, :].view(-1, 4) boxlist_for_class = BoxList(boxes_j, boxlist.size, mode="xyxy") boxlist_for_class.add_field("scores", scores_j) boxlist_for_class, keep = boxlist_nms_with_keep( boxlist_for_class, self.nms_thresh, score_field="scores") num_labels = len(boxlist_for_class) boxlist_for_class.add_field( "labels", torch.full((num_labels, ), j, dtype=torch.int64, device=scores.device)) boxlist_for_class.add_field("mask", masks[inds[keep], :, :, :]) result.append(boxlist_for_class) result = cat_boxlist(result) number_of_detections = len(result) # Limit to max_per_image detections **over all classes** if number_of_detections > self.fpn_post_nms_top_n > 0: cls_scores = result.get_field("scores") image_thresh, _ = torch.kthvalue( cls_scores.cpu(), number_of_detections - self.fpn_post_nms_top_n + 1) keep = cls_scores >= image_thresh.item() keep = torch.nonzero(keep).squeeze(1) result = result[keep] results.append(result) return results
def __call__(self, anchors, objectness, box_regression, targets): """ Arguments: anchors (list[BoxList]) objectness (list[Tensor]) box_regression (list[Tensor]) targets (list[BoxList]) Returns: objectness_loss (Tensor) box_loss (Tensor """ anchors = [ cat_boxlist(anchors_per_image) for anchors_per_image in anchors ] labels, regression_targets = self.prepare_targets(anchors, targets) sampled_pos_inds, sampled_neg_inds = self.fg_bg_sampler(labels) 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, box_regression = \ concat_box_prediction_layers(objectness, box_regression) objectness = objectness.squeeze() labels = torch.cat(labels, dim=0) regression_targets = torch.cat(regression_targets, dim=0) box_loss = smooth_l1_loss( box_regression[sampled_pos_inds], regression_targets[sampled_pos_inds], beta=1.0 / 9, size_average=False, ) / (sampled_inds.numel()) # box_loss = self.box_reg_loss_func(box_regression[sampled_pos_inds], # regression_targets[sampled_pos_inds]) / (sampled_inds.numel()) objectness_loss = F.binary_cross_entropy_with_logits( objectness[sampled_inds], labels[sampled_inds]) return objectness_loss, box_loss
def merge_result_from_multi_scales(boxlists, nms_type='nms', vote_thresh=0.65): num_images = len(boxlists) results = [] for i in range(num_images): scores = boxlists[i].get_field("scores") labels = boxlists[i].get_field("labels") boxes = boxlists[i].bbox boxlist = boxlists[i] result = [] # skip the background for j in range(1, cfg.MODEL.RETINANET.NUM_CLASSES): inds = (labels == j).nonzero().view(-1) scores_j = scores[inds] boxes_j = boxes[inds, :].view(-1, 4) boxlist_for_class = BoxList(boxes_j, boxlist.size, mode="xyxy") boxlist_for_class.add_field("scores", scores_j) boxlist_for_class = boxlist_nms(boxlist_for_class, cfg.MODEL.ATSS.NMS_TH, score_field="scores", nms_type=nms_type, vote_thresh=vote_thresh) num_labels = len(boxlist_for_class) boxlist_for_class.add_field( "labels", torch.full((num_labels, ), j, dtype=torch.int64, device=scores.device)) result.append(boxlist_for_class) result = cat_boxlist(result) number_of_detections = len(result) # Limit to max_per_image detections **over all classes** if number_of_detections > cfg.MODEL.ATSS.PRE_NMS_TOP_N > 0: cls_scores = result.get_field("scores") image_thresh, _ = torch.kthvalue( cls_scores.cpu(), number_of_detections - cfg.MODEL.ATSS.PRE_NMS_TOP_N + 1) keep = cls_scores >= image_thresh.item() keep = torch.nonzero(keep).squeeze(1) result = result[keep] results.append(result) return results
def filter_results(self, boxlist, num_classes): """Returns bounding-box detection results by thresholding on scores and applying non-maximum suppression (NMS). """ # unwrap the boxlist to avoid additional overhead. # if we had multi-class NMS, we could perform this directly on the boxlist boxes = boxlist.bbox.reshape(-1, num_classes * 4) scores = boxlist.get_field("scores").reshape(-1, num_classes) device = scores.device result = [] # Apply threshold on detection probabilities and apply NMS # Skip j = 0, because it's the background class inds_all = scores > self.score_thresh for j in range(1, num_classes): inds = inds_all[:, j].nonzero().squeeze(1) scores_j = scores[inds, j] boxes_j = boxes[inds, j * 4:(j + 1) * 4] boxlist_for_class = BoxList(boxes_j, boxlist.size, mode="xyxy") boxlist_for_class.add_field("scores", scores_j) boxlist_for_class = boxlist_nms(boxlist_for_class, self.nms) num_labels = len(boxlist_for_class) boxlist_for_class.add_field( "labels", torch.full((num_labels, ), j, dtype=torch.int64, device=device)) result.append(boxlist_for_class) result = cat_boxlist(result) number_of_detections = len(result) # Limit to max_per_image detections **over all classes** if number_of_detections > self.detections_per_img > 0: cls_scores = result.get_field("scores") image_thresh, _ = torch.kthvalue( cls_scores.cpu(), number_of_detections - self.detections_per_img + 1) keep = cls_scores >= image_thresh.item() keep = torch.nonzero(keep).squeeze(1) result = result[keep] return result
def __call__(self, anchors, box_cls, box_regression, targets): """ Arguments: anchors (list[BoxList]) box_cls (list[Tensor]) box_regression (list[Tensor]) targets (list[BoxList]) Returns: retinanet_cls_loss (Tensor) retinanet_regression_loss (Tensor """ anchors = [ cat_boxlist(anchors_per_image) for anchors_per_image in anchors ] labels, regression_targets = self.prepare_targets(anchors, targets) N = len(labels) box_cls, box_regression = \ concat_box_prediction_layers(box_cls, box_regression) labels = torch.cat(labels, dim=0) regression_targets = torch.cat(regression_targets, dim=0) pos_inds = torch.nonzero(labels > 0).squeeze(1) retinanet_regression_loss = smooth_l1_loss( box_regression[pos_inds], regression_targets[pos_inds], beta=self.bbox_reg_beta, size_average=False, ) / (max(1, pos_inds.numel() * self.regress_norm)) labels = labels.int() retinanet_cls_loss = self.box_cls_loss_func( box_cls, labels) / (pos_inds.numel() + N) return retinanet_cls_loss, retinanet_regression_loss
def add_gt_proposals(self, proposals, targets): """ Arguments: proposals: list[BoxList] targets: list[BoxList] """ # Get the device we're operating on device = proposals[0].bbox.device gt_boxes = [target.copy_with_fields([]) for target in targets] # later cat of bbox requires all fields to be present for all bbox # so we need to add a dummy for objectness that's missing for gt_box in gt_boxes: gt_box.add_field("objectness", torch.ones(len(gt_box), device=device)) proposals = [ cat_boxlist((proposal, gt_box)) for proposal, gt_box in zip(proposals, gt_boxes) ] return proposals
def __call__(self, anchors, objectness, box_regression, targets): """ Arguments: anchors (list[BoxList]) objectness (list[Tensor]) box_regression (list[Tensor]) targets (list[BoxList]) Returns: objectness_loss (Tensor) box_loss (Tensor """ print_log = (self.num_call % 100) == 0 self.num_call += 1 if print_log: all_anchor_sizes_each_pyramid = [[len(a) for a in anchors_per_image] for anchors_per_image in anchors] anchor_boxes = [cat_boxlist(anchors_per_image) for anchors_per_image in anchors] all_num_anchor_per_level = [[len(a) for a in anchors_per_image] for anchors_per_image in anchors] labels, regression_targets = self.prepare_targets( anchor_boxes, targets, all_num_anchor_per_level) if print_log: with torch.no_grad(): all_kind_to_num = get_kind_to_num_info(labels, all_anchor_sizes_each_pyramid) for kind_to_num in all_kind_to_num: for k in kind_to_num: kind_to_num[k] /= 1. * len(targets) from qd.qd_common import print_table print_table(all_kind_to_num) #if self.all_kind_to_num is None: #self.all_kind_to_num = all_kind_to_num #else: #for kind_to_num, self_kind_to_num in zip(all_kind_to_num, self.all_kind_to_num): #for kind, num in kind_to_num.items(): #self_kind_to_num[kind] += num #print_table(self.all_kind_to_num) sampled_pos_inds, sampled_neg_inds = self.fg_bg_sampler(labels) sampled_pos_inds = torch.nonzero(torch.cat(sampled_pos_inds, dim=0), as_tuple=False).squeeze(1) sampled_neg_inds = torch.nonzero(torch.cat(sampled_neg_inds, dim=0), as_tuple=False).squeeze(1) sampled_inds = torch.cat([sampled_pos_inds, sampled_neg_inds], dim=0) objectness, box_regression = \ concat_box_prediction_layers(objectness, box_regression) objectness = objectness.squeeze() labels = torch.cat(labels, dim=0) regression_targets = torch.cat(regression_targets, dim=0) box_loss = smooth_l1_loss( box_regression[sampled_pos_inds], regression_targets[sampled_pos_inds], beta=1.0 / 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 prepare_targets(self, targets, anchors): cls_labels = [] reg_targets = [] for im_i in range(len(targets)): targets_per_im = targets[im_i] assert targets_per_im.mode == "xyxy" bboxes_per_im = targets_per_im.bbox labels_per_im = targets_per_im.get_field("labels") anchors_per_im = cat_boxlist(anchors[im_i]) num_gt = bboxes_per_im.shape[0] if self.cfg.MODEL.ATSS.POSITIVE_TYPE == 'SSC': object_sizes_of_interest = [[-1, 64], [64, 128], [128, 256], [256, 512], [512, INF]] area_per_im = targets_per_im.area() expanded_object_sizes_of_interest = [] points = [] for l, anchors_per_level in enumerate(anchors[im_i]): anchors_per_level = anchors_per_level.bbox anchors_cx_per_level = (anchors_per_level[:, 2] + anchors_per_level[:, 0]) / 2.0 anchors_cy_per_level = (anchors_per_level[:, 3] + anchors_per_level[:, 1]) / 2.0 points_per_level = torch.stack((anchors_cx_per_level, anchors_cy_per_level), dim=1) points.append(points_per_level) object_sizes_of_interest_per_level = \ points_per_level.new_tensor(object_sizes_of_interest[l]) expanded_object_sizes_of_interest.append( object_sizes_of_interest_per_level[None].expand(len(points_per_level), -1) ) expanded_object_sizes_of_interest = torch.cat(expanded_object_sizes_of_interest, dim=0) points = torch.cat(points, dim=0) xs, ys = points[:, 0], points[:, 1] l = xs[:, None] - bboxes_per_im[:, 0][None] t = ys[:, None] - bboxes_per_im[:, 1][None] r = bboxes_per_im[:, 2][None] - xs[:, None] b = bboxes_per_im[:, 3][None] - ys[:, None] reg_targets_per_im = torch.stack([l, t, r, b], dim=2) is_in_boxes = reg_targets_per_im.min(dim=2)[0] > 0.01 max_reg_targets_per_im = reg_targets_per_im.max(dim=2)[0] is_cared_in_the_level = \ (max_reg_targets_per_im >= expanded_object_sizes_of_interest[:, [0]]) & \ (max_reg_targets_per_im <= expanded_object_sizes_of_interest[:, [1]]) locations_to_gt_area = area_per_im[None].repeat(len(points), 1) locations_to_gt_area[is_in_boxes == 0] = INF locations_to_gt_area[is_cared_in_the_level == 0] = INF locations_to_min_area, locations_to_gt_inds = locations_to_gt_area.min(dim=1) cls_labels_per_im = labels_per_im[locations_to_gt_inds] cls_labels_per_im[locations_to_min_area == INF] = 0 matched_gts = bboxes_per_im[locations_to_gt_inds] elif self.cfg.MODEL.ATSS.POSITIVE_TYPE == 'ATSS': num_anchors_per_loc = len(self.cfg.MODEL.ATSS.ASPECT_RATIOS) * self.cfg.MODEL.ATSS.SCALES_PER_OCTAVE num_anchors_per_level = [len(anchors_per_level.bbox) for anchors_per_level in anchors[im_i]] ious = boxlist_iou(anchors_per_im, targets_per_im) gt_cx = (bboxes_per_im[:, 2] + bboxes_per_im[:, 0]) / 2.0 gt_cy = (bboxes_per_im[:, 3] + bboxes_per_im[:, 1]) / 2.0 gt_points = torch.stack((gt_cx, gt_cy), dim=1) anchors_cx_per_im = (anchors_per_im.bbox[:, 2] + anchors_per_im.bbox[:, 0]) / 2.0 anchors_cy_per_im = (anchors_per_im.bbox[:, 3] + anchors_per_im.bbox[:, 1]) / 2.0 anchor_points = torch.stack((anchors_cx_per_im, anchors_cy_per_im), dim=1) distances = (anchor_points[:, None, :] - gt_points[None, :, :]).pow(2).sum(-1).sqrt() # Selecting candidates based on the center distance between anchor box and object candidate_idxs = [] star_idx = 0 for level, anchors_per_level in enumerate(anchors[im_i]): end_idx = star_idx + num_anchors_per_level[level] distances_per_level = distances[star_idx:end_idx, :] topk = min(self.cfg.MODEL.ATSS.TOPK * num_anchors_per_loc, num_anchors_per_level[level]) _, topk_idxs_per_level = distances_per_level.topk(topk, dim=0, largest=False) candidate_idxs.append(topk_idxs_per_level + star_idx) star_idx = end_idx candidate_idxs = torch.cat(candidate_idxs, dim=0) # Using the sum of mean and standard deviation as the IoU threshold to select final positive samples candidate_ious = ious[candidate_idxs, torch.arange(num_gt)] iou_mean_per_gt = candidate_ious.mean(0) iou_std_per_gt = candidate_ious.std(0) iou_thresh_per_gt = iou_mean_per_gt + iou_std_per_gt is_pos = candidate_ious >= iou_thresh_per_gt[None, :] # Limiting the final positive samples’ center to object anchor_num = anchors_cx_per_im.shape[0] for ng in range(num_gt): candidate_idxs[:, ng] += ng * anchor_num e_anchors_cx = anchors_cx_per_im.view(1, -1).expand(num_gt, anchor_num).contiguous().view(-1) e_anchors_cy = anchors_cy_per_im.view(1, -1).expand(num_gt, anchor_num).contiguous().view(-1) candidate_idxs = candidate_idxs.view(-1) l = e_anchors_cx[candidate_idxs].view(-1, num_gt) - bboxes_per_im[:, 0] t = e_anchors_cy[candidate_idxs].view(-1, num_gt) - bboxes_per_im[:, 1] r = bboxes_per_im[:, 2] - e_anchors_cx[candidate_idxs].view(-1, num_gt) b = bboxes_per_im[:, 3] - e_anchors_cy[candidate_idxs].view(-1, num_gt) is_in_gts = torch.stack([l, t, r, b], dim=1).min(dim=1)[0] > 0.01 is_pos = is_pos & is_in_gts # if an anchor box is assigned to multiple gts, the one with the highest IoU will be selected. ious_inf = torch.full_like(ious, -INF).t().contiguous().view(-1) index = candidate_idxs.view(-1)[is_pos.view(-1)] ious_inf[index] = ious.t().contiguous().view(-1)[index] ious_inf = ious_inf.view(num_gt, -1).t() anchors_to_gt_values, anchors_to_gt_indexs = ious_inf.max(dim=1) cls_labels_per_im = labels_per_im[anchors_to_gt_indexs] cls_labels_per_im[anchors_to_gt_values == -INF] = 0 matched_gts = bboxes_per_im[anchors_to_gt_indexs] elif self.cfg.MODEL.ATSS.POSITIVE_TYPE == 'TOPK': gt_cx = (bboxes_per_im[:, 2] + bboxes_per_im[:, 0]) / 2.0 gt_cy = (bboxes_per_im[:, 3] + bboxes_per_im[:, 1]) / 2.0 gt_points = torch.stack((gt_cx, gt_cy), dim=1) anchors_cx_per_im = (anchors_per_im.bbox[:, 2] + anchors_per_im.bbox[:, 0]) / 2.0 anchors_cy_per_im = (anchors_per_im.bbox[:, 3] + anchors_per_im.bbox[:, 1]) / 2.0 anchor_points = torch.stack((anchors_cx_per_im, anchors_cy_per_im), dim=1) distances = (anchor_points[:, None, :] - gt_points[None, :, :]).pow(2).sum(-1).sqrt() distances = distances / distances.max() / 1000 ious = boxlist_iou(anchors_per_im, targets_per_im) is_pos = ious * False for ng in range(num_gt): _, topk_idxs = (ious[:, ng] - distances[:, ng]).topk(self.cfg.MODEL.ATSS.TOPK, dim=0) l = anchors_cx_per_im[topk_idxs] - bboxes_per_im[ng, 0] t = anchors_cy_per_im[topk_idxs] - bboxes_per_im[ng, 1] r = bboxes_per_im[ng, 2] - anchors_cx_per_im[topk_idxs] b = bboxes_per_im[ng, 3] - anchors_cy_per_im[topk_idxs] is_in_gt = torch.stack([l, t, r, b], dim=1).min(dim=1)[0] > 0.01 is_pos[topk_idxs[is_in_gt == 1], ng] = True ious[is_pos == 0] = -INF anchors_to_gt_values, anchors_to_gt_indexs = ious.max(dim=1) cls_labels_per_im = labels_per_im[anchors_to_gt_indexs] cls_labels_per_im[anchors_to_gt_values == -INF] = 0 matched_gts = bboxes_per_im[anchors_to_gt_indexs] elif self.cfg.MODEL.ATSS.POSITIVE_TYPE == 'IoU': match_quality_matrix = boxlist_iou(targets_per_im, anchors_per_im) matched_idxs = self.matcher(match_quality_matrix) targets_per_im = targets_per_im.copy_with_fields(['labels']) matched_targets = targets_per_im[matched_idxs.clamp(min=0)] cls_labels_per_im = matched_targets.get_field("labels") cls_labels_per_im = cls_labels_per_im.to(dtype=torch.float32) # Background (negative examples) bg_indices = matched_idxs == Matcher.BELOW_LOW_THRESHOLD cls_labels_per_im[bg_indices] = 0 # discard indices that are between thresholds inds_to_discard = matched_idxs == Matcher.BETWEEN_THRESHOLDS cls_labels_per_im[inds_to_discard] = -1 matched_gts = matched_targets.bbox # Limiting positive samples’ center to object # in order to filter out poor positives and use the centerness branch pos_idxs = torch.nonzero(cls_labels_per_im > 0).squeeze(1) pos_anchors_cx = (anchors_per_im.bbox[pos_idxs, 2] + anchors_per_im.bbox[pos_idxs, 0]) / 2.0 pos_anchors_cy = (anchors_per_im.bbox[pos_idxs, 3] + anchors_per_im.bbox[pos_idxs, 1]) / 2.0 l = pos_anchors_cx - matched_gts[pos_idxs, 0] t = pos_anchors_cy - matched_gts[pos_idxs, 1] r = matched_gts[pos_idxs, 2] - pos_anchors_cx b = matched_gts[pos_idxs, 3] - pos_anchors_cy is_in_gts = torch.stack([l, t, r, b], dim=1).min(dim=1)[0] > 0.01 cls_labels_per_im[pos_idxs[is_in_gts == 0]] = -1 else: raise NotImplementedError reg_targets_per_im = self.box_coder.encode(matched_gts, anchors_per_im.bbox) cls_labels.append(cls_labels_per_im) reg_targets.append(reg_targets_per_im) return cls_labels, reg_targets