what.models.detection.ssd.utils.box_utils
1import math 2import itertools 3import collections 4from typing import List 5 6import torch 7 8 9SSDBoxSizes = collections.namedtuple('SSDBoxSizes', ['min', 'max']) 10 11SSDSpec = collections.namedtuple('SSDSpec', ['feature_map_size', 'shrinkage', 'box_sizes', 'aspect_ratios']) 12 13 14def generate_ssd_priors(specs: List[SSDSpec], image_size, clamp=True) -> torch.Tensor: 15 """Generate SSD Prior Boxes. 16 17 It returns the center, height and width of the priors. The values are relative to the image size 18 Args: 19 specs: SSDSpecs about the shapes of sizes of prior boxes. i.e. 20 specs = [ 21 SSDSpec(38, 8, SSDBoxSizes(30, 60), [2]), 22 SSDSpec(19, 16, SSDBoxSizes(60, 111), [2, 3]), 23 SSDSpec(10, 32, SSDBoxSizes(111, 162), [2, 3]), 24 SSDSpec(5, 64, SSDBoxSizes(162, 213), [2, 3]), 25 SSDSpec(3, 100, SSDBoxSizes(213, 264), [2]), 26 SSDSpec(1, 300, SSDBoxSizes(264, 315), [2]) 27 ] 28 image_size: image size. 29 clamp: if true, clamp the values to make fall between [0.0, 1.0] 30 Returns: 31 priors (num_priors, 4): The prior boxes represented as [[center_x, center_y, w, h]]. All the values 32 are relative to the image size. 33 """ 34 priors = [] 35 for spec in specs: 36 scale = image_size / spec.shrinkage 37 for j, i in itertools.product(range(spec.feature_map_size), repeat=2): 38 x_center = (i + 0.5) / scale 39 y_center = (j + 0.5) / scale 40 41 # small sized square box 42 size = spec.box_sizes.min 43 h = w = size / image_size 44 priors.append([ 45 x_center, 46 y_center, 47 w, 48 h 49 ]) 50 51 # big sized square box 52 size = math.sqrt(spec.box_sizes.max * spec.box_sizes.min) 53 h = w = size / image_size 54 priors.append([ 55 x_center, 56 y_center, 57 w, 58 h 59 ]) 60 61 # change h/w ratio of the small sized box 62 size = spec.box_sizes.min 63 h = w = size / image_size 64 for ratio in spec.aspect_ratios: 65 ratio = math.sqrt(ratio) 66 priors.append([ 67 x_center, 68 y_center, 69 w * ratio, 70 h / ratio 71 ]) 72 priors.append([ 73 x_center, 74 y_center, 75 w / ratio, 76 h * ratio 77 ]) 78 79 priors = torch.tensor(priors) 80 if clamp: 81 torch.clamp(priors, 0.0, 1.0, out=priors) 82 return priors 83 84 85def convert_locations_to_boxes(locations, priors, center_variance, 86 size_variance): 87 """Convert regressional location results of SSD into boxes in the form of (center_x, center_y, h, w). 88 89 The conversion: 90 $$predicted\_center * center_variance = \frac {real\_center - prior\_center} {prior\_hw}$$ 91 $$exp(predicted\_hw * size_variance) = \frac {real\_hw} {prior\_hw}$$ 92 We do it in the inverse direction here. 93 Args: 94 locations (batch_size, num_priors, 4): the regression output of SSD. It will contain the outputs as well. 95 priors (num_priors, 4) or (batch_size/1, num_priors, 4): prior boxes. 96 center_variance: a float used to change the scale of center. 97 size_variance: a float used to change of scale of size. 98 Returns: 99 boxes: priors: [[center_x, center_y, h, w]]. All the values 100 are relative to the image size. 101 """ 102 # priors can have one dimension less. 103 if priors.dim() + 1 == locations.dim(): 104 priors = priors.unsqueeze(0) 105 return torch.cat([ 106 locations[..., :2] * center_variance * priors[..., 2:] + priors[..., :2], 107 torch.exp(locations[..., 2:] * size_variance) * priors[..., 2:] 108 ], dim=locations.dim() - 1) 109 110 111def convert_boxes_to_locations(center_form_boxes, center_form_priors, center_variance, size_variance): 112 # priors can have one dimension less 113 if center_form_priors.dim() + 1 == center_form_boxes.dim(): 114 center_form_priors = center_form_priors.unsqueeze(0) 115 return torch.cat([ 116 (center_form_boxes[..., :2] - center_form_priors[..., :2]) / center_form_priors[..., 2:] / center_variance, 117 torch.log(center_form_boxes[..., 2:] / center_form_priors[..., 2:]) / size_variance 118 ], dim=center_form_boxes.dim() - 1) 119 120 121def area_of(left_top, right_bottom) -> torch.Tensor: 122 """Compute the areas of rectangles given two corners. 123 124 Args: 125 left_top (N, 2): left top corner. 126 right_bottom (N, 2): right bottom corner. 127 128 Returns: 129 area (N): return the area. 130 """ 131 hw = torch.clamp(right_bottom - left_top, min=0.0) 132 return hw[..., 0] * hw[..., 1] 133 134 135def iou_of(boxes0, boxes1, eps=1e-5): 136 """Return intersection-over-union (Jaccard index) of boxes. 137 138 Args: 139 boxes0 (N, 4): ground truth boxes. 140 boxes1 (N or 1, 4): predicted boxes. 141 eps: a small number to avoid 0 as denominator. 142 Returns: 143 iou (N): IoU values. 144 """ 145 overlap_left_top = torch.max(boxes0[..., :2], boxes1[..., :2]) 146 overlap_right_bottom = torch.min(boxes0[..., 2:], boxes1[..., 2:]) 147 148 overlap_area = area_of(overlap_left_top, overlap_right_bottom) 149 area0 = area_of(boxes0[..., :2], boxes0[..., 2:]) 150 area1 = area_of(boxes1[..., :2], boxes1[..., 2:]) 151 return overlap_area / (area0 + area1 - overlap_area + eps) 152 153 154def assign_priors(gt_boxes, gt_labels, corner_form_priors, iou_threshold): 155 """Assign ground truth boxes and targets to priors. 156 157 Args: 158 gt_boxes (num_targets, 4): ground truth boxes. 159 gt_labels (num_targets): labels of targets. 160 priors (num_priors, 4): corner form priors 161 Returns: 162 boxes (num_priors, 4): real values for priors. 163 labels (num_priros): labels for priors. 164 """ 165 # size: num_priors x num_targets 166 ious = iou_of(gt_boxes.unsqueeze(0), corner_form_priors.unsqueeze(1)) 167 # size: num_priors 168 best_target_per_prior, best_target_per_prior_index = ious.max(1) 169 # size: num_targets 170 best_prior_per_target, best_prior_per_target_index = ious.max(0) 171 172 for target_index, prior_index in enumerate(best_prior_per_target_index): 173 best_target_per_prior_index[prior_index] = target_index 174 # 2.0 is used to make sure every target has a prior assigned 175 best_target_per_prior.index_fill_(0, best_prior_per_target_index, 2) 176 # size: num_priors 177 labels = gt_labels[best_target_per_prior_index] 178 labels[best_target_per_prior < iou_threshold] = 0 # the backgournd id 179 boxes = gt_boxes[best_target_per_prior_index] 180 return boxes, labels 181 182 183def hard_negative_mining(loss, labels, neg_pos_ratio): 184 """ 185 It used to suppress the presence of a large number of negative prediction. 186 It works on image level not batch level. 187 For any example/image, it keeps all the positive predictions and 188 cut the number of negative predictions to make sure the ratio 189 between the negative examples and positive examples is no more 190 the given ratio for an image. 191 192 Args: 193 loss (N, num_priors): the loss for each example. 194 labels (N, num_priors): the labels. 195 neg_pos_ratio: the ratio between the negative examples and positive examples. 196 """ 197 pos_mask = labels > 0 198 num_pos = pos_mask.long().sum(dim=1, keepdim=True) 199 num_neg = num_pos * neg_pos_ratio 200 201 loss[pos_mask] = -math.inf 202 _, indexes = loss.sort(dim=1, descending=True) 203 _, orders = indexes.sort(dim=1) 204 neg_mask = orders < num_neg 205 return pos_mask | neg_mask 206 207 208def center_form_to_corner_form(locations): 209 return torch.cat([locations[..., :2] - locations[..., 2:]/2, 210 locations[..., :2] + locations[..., 2:]/2], locations.dim() - 1) 211 212 213def corner_form_to_center_form(boxes): 214 return torch.cat([ 215 (boxes[..., :2] + boxes[..., 2:]) / 2, 216 boxes[..., 2:] - boxes[..., :2] 217 ], boxes.dim() - 1) 218 219 220def hard_nms(box_scores, iou_threshold, top_k=-1, candidate_size=200): 221 """ 222 223 Args: 224 box_scores (N, 5): boxes in corner-form and probabilities. 225 iou_threshold: intersection over union threshold. 226 top_k: keep top_k results. If k <= 0, keep all the results. 227 candidate_size: only consider the candidates with the highest scores. 228 Returns: 229 picked: a list of indexes of the kept boxes 230 """ 231 scores = box_scores[:, -1] 232 boxes = box_scores[:, :-1] 233 picked = [] 234 _, indexes = scores.sort(descending=True) 235 indexes = indexes[:candidate_size] 236 while len(indexes) > 0: 237 current = indexes[0] 238 picked.append(current.item()) 239 if 0 < top_k == len(picked) or len(indexes) == 1: 240 break 241 current_box = boxes[current, :] 242 indexes = indexes[1:] 243 rest_boxes = boxes[indexes, :] 244 iou = iou_of( 245 rest_boxes, 246 current_box.unsqueeze(0), 247 ) 248 indexes = indexes[iou <= iou_threshold] 249 250 return box_scores[picked, :] 251 252 253def nms(box_scores, nms_method=None, score_threshold=None, iou_threshold=None, 254 sigma=0.5, top_k=-1, candidate_size=200): 255 if nms_method == "soft": 256 return soft_nms(box_scores, score_threshold, sigma, top_k) 257 else: 258 return hard_nms(box_scores, iou_threshold, top_k, candidate_size=candidate_size) 259 260 261def soft_nms(box_scores, score_threshold, sigma=0.5, top_k=-1): 262 """Soft NMS implementation. 263 264 References: 265 https://arxiv.org/abs/1704.04503 266 https://github.com/facebookresearch/Detectron/blob/master/detectron/utils/cython_nms.pyx 267 268 Args: 269 box_scores (N, 5): boxes in corner-form and probabilities. 270 score_threshold: boxes with scores less than value are not considered. 271 sigma: the parameter in score re-computation. 272 scores[i] = scores[i] * exp(-(iou_i)^2 / simga) 273 top_k: keep top_k results. If k <= 0, keep all the results. 274 Returns: 275 picked_box_scores (K, 5): results of NMS. 276 """ 277 picked_box_scores = [] 278 while box_scores.size(0) > 0: 279 max_score_index = torch.argmax(box_scores[:, 4]) 280 cur_box_prob = torch.tensor(box_scores[max_score_index, :]) 281 picked_box_scores.append(cur_box_prob) 282 if len(picked_box_scores) == top_k > 0 or box_scores.size(0) == 1: 283 break 284 cur_box = cur_box_prob[:-1] 285 box_scores[max_score_index, :] = box_scores[-1, :] 286 box_scores = box_scores[:-1, :] 287 ious = iou_of(cur_box.unsqueeze(0), box_scores[:, :-1]) 288 box_scores[:, -1] = box_scores[:, -1] * torch.exp(-(ious * ious) / sigma) 289 box_scores = box_scores[box_scores[:, -1] > score_threshold, :] 290 if len(picked_box_scores) > 0: 291 return torch.stack(picked_box_scores) 292 else: 293 return torch.tensor([])
SSDBoxSizes(min, max)
Inherited Members
- builtins.tuple
- index
- count
SSDSpec(feature_map_size, shrinkage, box_sizes, aspect_ratios)
Create new instance of SSDSpec(feature_map_size, shrinkage, box_sizes, aspect_ratios)
Inherited Members
- builtins.tuple
- index
- count
15def generate_ssd_priors(specs: List[SSDSpec], image_size, clamp=True) -> torch.Tensor: 16 """Generate SSD Prior Boxes. 17 18 It returns the center, height and width of the priors. The values are relative to the image size 19 Args: 20 specs: SSDSpecs about the shapes of sizes of prior boxes. i.e. 21 specs = [ 22 SSDSpec(38, 8, SSDBoxSizes(30, 60), [2]), 23 SSDSpec(19, 16, SSDBoxSizes(60, 111), [2, 3]), 24 SSDSpec(10, 32, SSDBoxSizes(111, 162), [2, 3]), 25 SSDSpec(5, 64, SSDBoxSizes(162, 213), [2, 3]), 26 SSDSpec(3, 100, SSDBoxSizes(213, 264), [2]), 27 SSDSpec(1, 300, SSDBoxSizes(264, 315), [2]) 28 ] 29 image_size: image size. 30 clamp: if true, clamp the values to make fall between [0.0, 1.0] 31 Returns: 32 priors (num_priors, 4): The prior boxes represented as [[center_x, center_y, w, h]]. All the values 33 are relative to the image size. 34 """ 35 priors = [] 36 for spec in specs: 37 scale = image_size / spec.shrinkage 38 for j, i in itertools.product(range(spec.feature_map_size), repeat=2): 39 x_center = (i + 0.5) / scale 40 y_center = (j + 0.5) / scale 41 42 # small sized square box 43 size = spec.box_sizes.min 44 h = w = size / image_size 45 priors.append([ 46 x_center, 47 y_center, 48 w, 49 h 50 ]) 51 52 # big sized square box 53 size = math.sqrt(spec.box_sizes.max * spec.box_sizes.min) 54 h = w = size / image_size 55 priors.append([ 56 x_center, 57 y_center, 58 w, 59 h 60 ]) 61 62 # change h/w ratio of the small sized box 63 size = spec.box_sizes.min 64 h = w = size / image_size 65 for ratio in spec.aspect_ratios: 66 ratio = math.sqrt(ratio) 67 priors.append([ 68 x_center, 69 y_center, 70 w * ratio, 71 h / ratio 72 ]) 73 priors.append([ 74 x_center, 75 y_center, 76 w / ratio, 77 h * ratio 78 ]) 79 80 priors = torch.tensor(priors) 81 if clamp: 82 torch.clamp(priors, 0.0, 1.0, out=priors) 83 return priors
Generate SSD Prior Boxes.
It returns the center, height and width of the priors. The values are relative to the image size Args: specs: SSDSpecs about the shapes of sizes of prior boxes. i.e. specs = [ SSDSpec(38, 8, SSDBoxSizes(30, 60), [2]), SSDSpec(19, 16, SSDBoxSizes(60, 111), [2, 3]), SSDSpec(10, 32, SSDBoxSizes(111, 162), [2, 3]), SSDSpec(5, 64, SSDBoxSizes(162, 213), [2, 3]), SSDSpec(3, 100, SSDBoxSizes(213, 264), [2]), SSDSpec(1, 300, SSDBoxSizes(264, 315), [2]) ] image_size: image size. clamp: if true, clamp the values to make fall between [0.0, 1.0] Returns: priors (num_priors, 4): The prior boxes represented as [[center_x, center_y, w, h]]. All the values are relative to the image size.
86def convert_locations_to_boxes(locations, priors, center_variance, 87 size_variance): 88 """Convert regressional location results of SSD into boxes in the form of (center_x, center_y, h, w). 89 90 The conversion: 91 $$predicted\_center * center_variance = \frac {real\_center - prior\_center} {prior\_hw}$$ 92 $$exp(predicted\_hw * size_variance) = \frac {real\_hw} {prior\_hw}$$ 93 We do it in the inverse direction here. 94 Args: 95 locations (batch_size, num_priors, 4): the regression output of SSD. It will contain the outputs as well. 96 priors (num_priors, 4) or (batch_size/1, num_priors, 4): prior boxes. 97 center_variance: a float used to change the scale of center. 98 size_variance: a float used to change of scale of size. 99 Returns: 100 boxes: priors: [[center_x, center_y, h, w]]. All the values 101 are relative to the image size. 102 """ 103 # priors can have one dimension less. 104 if priors.dim() + 1 == locations.dim(): 105 priors = priors.unsqueeze(0) 106 return torch.cat([ 107 locations[..., :2] * center_variance * priors[..., 2:] + priors[..., :2], 108 torch.exp(locations[..., 2:] * size_variance) * priors[..., 2:] 109 ], dim=locations.dim() - 1)
Convert regressional location results of SSD into boxes in the form of (center_x, center_y, h, w).
The conversion: $$predicted_center * center_variance = rac {real_center - prior_center} {prior_hw}$$ $$exp(predicted_hw * size_variance) = rac {real_hw} {prior_hw}$$ We do it in the inverse direction here. Args: locations (batch_size, num_priors, 4): the regression output of SSD. It will contain the outputs as well. priors (num_priors, 4) or (batch_size/1, num_priors, 4): prior boxes. center_variance: a float used to change the scale of center. size_variance: a float used to change of scale of size. Returns: boxes: priors: [[center_x, center_y, h, w]]. All the values are relative to the image size.
112def convert_boxes_to_locations(center_form_boxes, center_form_priors, center_variance, size_variance): 113 # priors can have one dimension less 114 if center_form_priors.dim() + 1 == center_form_boxes.dim(): 115 center_form_priors = center_form_priors.unsqueeze(0) 116 return torch.cat([ 117 (center_form_boxes[..., :2] - center_form_priors[..., :2]) / center_form_priors[..., 2:] / center_variance, 118 torch.log(center_form_boxes[..., 2:] / center_form_priors[..., 2:]) / size_variance 119 ], dim=center_form_boxes.dim() - 1)
122def area_of(left_top, right_bottom) -> torch.Tensor: 123 """Compute the areas of rectangles given two corners. 124 125 Args: 126 left_top (N, 2): left top corner. 127 right_bottom (N, 2): right bottom corner. 128 129 Returns: 130 area (N): return the area. 131 """ 132 hw = torch.clamp(right_bottom - left_top, min=0.0) 133 return hw[..., 0] * hw[..., 1]
Compute the areas of rectangles given two corners.
Args: left_top (N, 2): left top corner. right_bottom (N, 2): right bottom corner.
Returns: area (N): return the area.
136def iou_of(boxes0, boxes1, eps=1e-5): 137 """Return intersection-over-union (Jaccard index) of boxes. 138 139 Args: 140 boxes0 (N, 4): ground truth boxes. 141 boxes1 (N or 1, 4): predicted boxes. 142 eps: a small number to avoid 0 as denominator. 143 Returns: 144 iou (N): IoU values. 145 """ 146 overlap_left_top = torch.max(boxes0[..., :2], boxes1[..., :2]) 147 overlap_right_bottom = torch.min(boxes0[..., 2:], boxes1[..., 2:]) 148 149 overlap_area = area_of(overlap_left_top, overlap_right_bottom) 150 area0 = area_of(boxes0[..., :2], boxes0[..., 2:]) 151 area1 = area_of(boxes1[..., :2], boxes1[..., 2:]) 152 return overlap_area / (area0 + area1 - overlap_area + eps)
Return intersection-over-union (Jaccard index) of boxes.
Args: boxes0 (N, 4): ground truth boxes. boxes1 (N or 1, 4): predicted boxes. eps: a small number to avoid 0 as denominator. Returns: iou (N): IoU values.
155def assign_priors(gt_boxes, gt_labels, corner_form_priors, iou_threshold): 156 """Assign ground truth boxes and targets to priors. 157 158 Args: 159 gt_boxes (num_targets, 4): ground truth boxes. 160 gt_labels (num_targets): labels of targets. 161 priors (num_priors, 4): corner form priors 162 Returns: 163 boxes (num_priors, 4): real values for priors. 164 labels (num_priros): labels for priors. 165 """ 166 # size: num_priors x num_targets 167 ious = iou_of(gt_boxes.unsqueeze(0), corner_form_priors.unsqueeze(1)) 168 # size: num_priors 169 best_target_per_prior, best_target_per_prior_index = ious.max(1) 170 # size: num_targets 171 best_prior_per_target, best_prior_per_target_index = ious.max(0) 172 173 for target_index, prior_index in enumerate(best_prior_per_target_index): 174 best_target_per_prior_index[prior_index] = target_index 175 # 2.0 is used to make sure every target has a prior assigned 176 best_target_per_prior.index_fill_(0, best_prior_per_target_index, 2) 177 # size: num_priors 178 labels = gt_labels[best_target_per_prior_index] 179 labels[best_target_per_prior < iou_threshold] = 0 # the backgournd id 180 boxes = gt_boxes[best_target_per_prior_index] 181 return boxes, labels
Assign ground truth boxes and targets to priors.
Args: gt_boxes (num_targets, 4): ground truth boxes. gt_labels (num_targets): labels of targets. priors (num_priors, 4): corner form priors Returns: boxes (num_priors, 4): real values for priors. labels (num_priros): labels for priors.
184def hard_negative_mining(loss, labels, neg_pos_ratio): 185 """ 186 It used to suppress the presence of a large number of negative prediction. 187 It works on image level not batch level. 188 For any example/image, it keeps all the positive predictions and 189 cut the number of negative predictions to make sure the ratio 190 between the negative examples and positive examples is no more 191 the given ratio for an image. 192 193 Args: 194 loss (N, num_priors): the loss for each example. 195 labels (N, num_priors): the labels. 196 neg_pos_ratio: the ratio between the negative examples and positive examples. 197 """ 198 pos_mask = labels > 0 199 num_pos = pos_mask.long().sum(dim=1, keepdim=True) 200 num_neg = num_pos * neg_pos_ratio 201 202 loss[pos_mask] = -math.inf 203 _, indexes = loss.sort(dim=1, descending=True) 204 _, orders = indexes.sort(dim=1) 205 neg_mask = orders < num_neg 206 return pos_mask | neg_mask
It used to suppress the presence of a large number of negative prediction. It works on image level not batch level. For any example/image, it keeps all the positive predictions and cut the number of negative predictions to make sure the ratio between the negative examples and positive examples is no more the given ratio for an image.
Args: loss (N, num_priors): the loss for each example. labels (N, num_priors): the labels. neg_pos_ratio: the ratio between the negative examples and positive examples.
221def hard_nms(box_scores, iou_threshold, top_k=-1, candidate_size=200): 222 """ 223 224 Args: 225 box_scores (N, 5): boxes in corner-form and probabilities. 226 iou_threshold: intersection over union threshold. 227 top_k: keep top_k results. If k <= 0, keep all the results. 228 candidate_size: only consider the candidates with the highest scores. 229 Returns: 230 picked: a list of indexes of the kept boxes 231 """ 232 scores = box_scores[:, -1] 233 boxes = box_scores[:, :-1] 234 picked = [] 235 _, indexes = scores.sort(descending=True) 236 indexes = indexes[:candidate_size] 237 while len(indexes) > 0: 238 current = indexes[0] 239 picked.append(current.item()) 240 if 0 < top_k == len(picked) or len(indexes) == 1: 241 break 242 current_box = boxes[current, :] 243 indexes = indexes[1:] 244 rest_boxes = boxes[indexes, :] 245 iou = iou_of( 246 rest_boxes, 247 current_box.unsqueeze(0), 248 ) 249 indexes = indexes[iou <= iou_threshold] 250 251 return box_scores[picked, :]
Args: box_scores (N, 5): boxes in corner-form and probabilities. iou_threshold: intersection over union threshold. top_k: keep top_k results. If k <= 0, keep all the results. candidate_size: only consider the candidates with the highest scores. Returns: picked: a list of indexes of the kept boxes
254def nms(box_scores, nms_method=None, score_threshold=None, iou_threshold=None, 255 sigma=0.5, top_k=-1, candidate_size=200): 256 if nms_method == "soft": 257 return soft_nms(box_scores, score_threshold, sigma, top_k) 258 else: 259 return hard_nms(box_scores, iou_threshold, top_k, candidate_size=candidate_size)
262def soft_nms(box_scores, score_threshold, sigma=0.5, top_k=-1): 263 """Soft NMS implementation. 264 265 References: 266 https://arxiv.org/abs/1704.04503 267 https://github.com/facebookresearch/Detectron/blob/master/detectron/utils/cython_nms.pyx 268 269 Args: 270 box_scores (N, 5): boxes in corner-form and probabilities. 271 score_threshold: boxes with scores less than value are not considered. 272 sigma: the parameter in score re-computation. 273 scores[i] = scores[i] * exp(-(iou_i)^2 / simga) 274 top_k: keep top_k results. If k <= 0, keep all the results. 275 Returns: 276 picked_box_scores (K, 5): results of NMS. 277 """ 278 picked_box_scores = [] 279 while box_scores.size(0) > 0: 280 max_score_index = torch.argmax(box_scores[:, 4]) 281 cur_box_prob = torch.tensor(box_scores[max_score_index, :]) 282 picked_box_scores.append(cur_box_prob) 283 if len(picked_box_scores) == top_k > 0 or box_scores.size(0) == 1: 284 break 285 cur_box = cur_box_prob[:-1] 286 box_scores[max_score_index, :] = box_scores[-1, :] 287 box_scores = box_scores[:-1, :] 288 ious = iou_of(cur_box.unsqueeze(0), box_scores[:, :-1]) 289 box_scores[:, -1] = box_scores[:, -1] * torch.exp(-(ious * ious) / sigma) 290 box_scores = box_scores[box_scores[:, -1] > score_threshold, :] 291 if len(picked_box_scores) > 0: 292 return torch.stack(picked_box_scores) 293 else: 294 return torch.tensor([])
Soft NMS implementation.
References: https://arxiv.org/abs/1704.04503 https://github.com/facebookresearch/Detectron/blob/master/detectron/utils/cython_nms.pyx
Args: box_scores (N, 5): boxes in corner-form and probabilities. score_threshold: boxes with scores less than value are not considered. sigma: the parameter in score re-computation. scores[i] = scores[i] * exp(-(iou_i)^2 / simga) top_k: keep top_k results. If k <= 0, keep all the results. Returns: picked_box_scores (K, 5): results of NMS.