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([])
class SSDBoxSizes(builtins.tuple):

SSDBoxSizes(min, max)

SSDBoxSizes(min, max)

Create new instance of SSDBoxSizes(min, max)

min

Alias for field number 0

max

Alias for field number 1

Inherited Members
builtins.tuple
index
count
class SSDSpec(builtins.tuple):

SSDSpec(feature_map_size, shrinkage, box_sizes, aspect_ratios)

SSDSpec(feature_map_size, shrinkage, box_sizes, aspect_ratios)

Create new instance of SSDSpec(feature_map_size, shrinkage, box_sizes, aspect_ratios)

feature_map_size

Alias for field number 0

shrinkage

Alias for field number 1

box_sizes

Alias for field number 2

aspect_ratios

Alias for field number 3

Inherited Members
builtins.tuple
index
count
def generate_ssd_priors( specs: List[what.models.detection.ssd.utils.box_utils.SSDSpec], image_size, clamp=True) -> torch.Tensor:
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.

def convert_locations_to_boxes(locations, priors, center_variance, size_variance):
 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.

def convert_boxes_to_locations( center_form_boxes, center_form_priors, center_variance, size_variance):
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)
def area_of(left_top, right_bottom) -> torch.Tensor:
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.

def iou_of(boxes0, boxes1, eps=1e-05):
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.

def assign_priors(gt_boxes, gt_labels, corner_form_priors, iou_threshold):
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.

def hard_negative_mining(loss, labels, neg_pos_ratio):
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.

def center_form_to_corner_form(locations):
209def center_form_to_corner_form(locations):
210    return torch.cat([locations[..., :2] - locations[..., 2:]/2,
211                     locations[..., :2] + locations[..., 2:]/2], locations.dim() - 1) 
def corner_form_to_center_form(boxes):
214def corner_form_to_center_form(boxes):
215    return torch.cat([
216        (boxes[..., :2] + boxes[..., 2:]) / 2,
217         boxes[..., 2:] - boxes[..., :2]
218    ], boxes.dim() - 1)
def hard_nms(box_scores, iou_threshold, top_k=-1, candidate_size=200):
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

def nms( box_scores, nms_method=None, score_threshold=None, iou_threshold=None, sigma=0.5, top_k=-1, candidate_size=200):
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)
def soft_nms(box_scores, score_threshold, sigma=0.5, top_k=-1):
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.