what.models.detection.ssd.utils.box_utils API documentation

# what.models.detection.ssd.utils.box_utils

View Source

      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: View Source

    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): View Source

     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): View Source

    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: View Source

    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): View Source

    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): View Source

    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): View Source

    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): View Source

    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): View Source

    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): View Source

    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): View Source

    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): View Source

    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.