Skip to content

Instantly share code, notes, and snippets.

@sukhitashvili

sukhitashvili/non_maximum_suppression.py

Last active September 5, 2023 11:59
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save sukhitashvili/745b7528f0418172e84d4ddb23f666b9 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save sukhitashvili/745b7528f0418172e84d4ddb23f666b9 to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. sukhitashvili revised this gist Sep 5, 2023. 1 changed file with 2 additions and 0 deletions.
    2 changes: 2 additions & 0 deletions non_maximum_suppression.py
    View file
    Original file line number Diff line number Diff line change
    @@ -20,6 +20,8 @@ def non_max_supression(predictions: np.ndarray, conf_thr: float, iou_thr: float)
    Does NMS on predictions.
    Args:
    predictions: [[x1, y1, x2, y2, conf, cls], ...]
    conf_thr: float
    iou_thr: float
    Returns:
    predictions: [[x1, y1, x2, y2, conf, cls], ...]
  2. sukhitashvili revised this gist Sep 5, 2023. 1 changed file with 33 additions and 26 deletions.
    59 changes: 33 additions & 26 deletions non_maximum_suppression.py
    View file
    Original file line number Diff line number Diff line change
    @@ -1,37 +1,44 @@
    import torch

    def get_iou(box1: torch.Tensor, box2: torch.Tensor):
    def get_iou(pred1: torch.Tensor, pred2: torch.Tensor):
    """
    box - tensor of shape [5]
    prediction: [x1, y1, x2, y2, conf, cls]
    """
    ar1 = abs((box1[3] - box1[1]) * (box1[4] - box1[2]))
    ar2 = abs((box2[3] - box2[1]) * (box2[4] - box2[2]))
    x1 = torch.max(box1[1], box2[1])
    y1 = torch.max(box1[2], box2[2])
    x2 = torch.min(box1[3], box2[3])
    y2 = torch.min(box1[4], box2[4])
    inter = (x2-x1).clamp(0) * (y2-y1).clamp(0) # clamp to zero if they don't intersect
    iou = inter/(ar1+ar2-inter)
    ar1 = abs((pred1[2] - pred1[0]) * (pred1[3] - pred1[1]))
    ar2 = abs((pred2[2] - pred2[0]) * (pred2[3] - pred2[1]))
    x1 = torch.max(pred1[0], pred2[0])
    y1 = torch.max(pred1[1], pred2[1])
    x2 = torch.min(pred1[2], pred2[2])
    y2 = torch.min(pred1[3], pred2[3])
    inter = (x2 - x1).clamp(0) * (y2 - y1).clamp(0) # clamp to zero if they don't intersect
    iou = inter / (ar1 + ar2 - inter)
    return iou


    def non_max_supression(boxes, iou_val=0.5, prob_thr: 0.5, max_boxes=10000):
    def non_max_supression(predictions: np.ndarray, conf_thr: float, iou_thr: float) -> np.ndarray:
    """
    boxes - [(class_id, score, x1, y1, x2, y2), ...]
    Does NMS on predictions.
    Args:
    predictions: [[x1, y1, x2, y2, conf, cls], ...]
    Returns:
    predictions: [[x1, y1, x2, y2, conf, cls], ...]
    """
    if max_boxes>0:
    boxes = boxes[:max_boxes]
    boxes = [box for box in boxes if box[1] >= prob_thr]
    boxes = sorted(boxes, key = lambda x:x[1], reverse=True)
    rboxes = []
    while len(boxes)>0:
    cur_box = boxes.pop(0)
    res_predictions = [pred for pred in predictions if pred[4] >= conf_thr]
    res_predictions = sorted(res_predictions, key=lambda x: x[4], reverse=True)

    if len(res_predictions) == 0:
    return predictions

    results = []
    while len(res_predictions) > 0:
    cur_pred = res_predictions.pop(0)
    # recreate boxes
    boxes =[
    box for box in boxes
    if (get_iou(torch.tensor(box), torch.tensor(cur_box)).item() < iou_val) or (box[0] != cur_box[0])
    res_predictions = [
    pred for pred in res_predictions
    if
    (get_iou(torch.tensor(pred), torch.tensor(cur_pred)).item() < iou_thr) or (pred[4] != cur_pred[4])
    ]
    rboxes.append(cur_box)
    return rboxes
    results.append(cur_pred)
    results = np.array(results)
    return results
  3. sukhitashvili revised this gist Sep 18, 2022. 1 changed file with 1 addition and 1 deletion.
    2 changes: 1 addition & 1 deletion non_maximum_suppression.py
    View file
    Original file line number Diff line number Diff line change
    @@ -27,7 +27,7 @@ def non_max_supression(boxes, iou_val=0.5, prob_thr: 0.5, max_boxes=10000):

    rboxes = []
    while len(boxes)>0:
    cur_box = boxes.pop(0)cur_box
    cur_box = boxes.pop(0)
    # recreate boxes
    boxes =[
    box for box in boxes
  4. sukhitashvili revised this gist Sep 18, 2022. 1 changed file with 11 additions and 5 deletions.
    16 changes: 11 additions & 5 deletions non_maximum_suppression.py
    View file
    Original file line number Diff line number Diff line change
    @@ -1,6 +1,9 @@
    import torch

    def get_iou(box1: torch.Tensor, box2: torch.Tensor):
    """
    box - tensor of shape [5]
    """
    ar1 = abs((box1[3] - box1[1]) * (box1[4] - box1[2]))
    ar2 = abs((box2[3] - box2[1]) * (box2[4] - box2[2]))
    x1 = torch.max(box1[1], box2[1])
    @@ -12,20 +15,23 @@ def get_iou(box1: torch.Tensor, box2: torch.Tensor):
    return iou


    def non_max_supression(boxes, iou_val=0.1, max_boxes=10000):
    def non_max_supression(boxes, iou_val=0.5, prob_thr: 0.5, max_boxes=10000):
    """
    boxes - [(score, x1, y1, x2, y2), ...]
    boxes - [(class_id, score, x1, y1, x2, y2), ...]
    """
    boxes = sorted(boxes, key = lambda x:x[0], reverse=True)
    if max_boxes>0:
    boxes = boxes[:max_boxes]

    boxes = [box for box in boxes if box[1] >= prob_thr]
    boxes = sorted(boxes, key = lambda x:x[1], reverse=True)

    rboxes = []
    while len(boxes)>0:
    cur_box = boxes.pop(0)
    cur_box = boxes.pop(0)cur_box
    # recreate boxes
    boxes =[
    box for box in boxes
    if get_iou(torch.tensor(box), torch.tensor(cur_box)).item() < iou_val
    if (get_iou(torch.tensor(box), torch.tensor(cur_box)).item() < iou_val) or (box[0] != cur_box[0])
    ]
    rboxes.append(cur_box)
    return rboxes
  5. sukhitashvili revised this gist Sep 17, 2022. 1 changed file with 11 additions and 9 deletions.
    20 changes: 11 additions & 9 deletions non_maximum_suppression.py
    View file
    Original file line number Diff line number Diff line change
    @@ -1,11 +1,13 @@
    def get_iou(box1, box2):
    ar1 = (box1[3] - box1[1]) * (box1[4] - box1[2])
    ar2 = (box2[3] - box2[1]) * (box2[4] - box2[2])
    x1 = max(box1[1], box2[1])
    y1 = max(box1[2], box2[2])
    x2 = min(box1[3], box2[3])
    y2 = min(box1[4], box2[4])
    inter = (x2-x1)*(y2-y1)
    import torch

    def get_iou(box1: torch.Tensor, box2: torch.Tensor):
    ar1 = abs((box1[3] - box1[1]) * (box1[4] - box1[2]))
    ar2 = abs((box2[3] - box2[1]) * (box2[4] - box2[2]))
    x1 = torch.max(box1[1], box2[1])
    y1 = torch.max(box1[2], box2[2])
    x2 = torch.min(box1[3], box2[3])
    y2 = torch.min(box1[4], box2[4])
    inter = (x2-x1).clamp(0) * (y2-y1).clamp(0) # clamp to zero if they don't intersect
    iou = inter/(ar1+ar2-inter)
    return iou

    @@ -23,7 +25,7 @@ def non_max_supression(boxes, iou_val=0.1, max_boxes=10000):
    # recreate boxes
    boxes =[
    box for box in boxes
    if get_iou(box, cur_box)<iou_val
    if get_iou(torch.tensor(box), torch.tensor(cur_box)).item() < iou_val
    ]
    rboxes.append(cur_box)
    return rboxes
  6. sukhitashvili created this gist Jun 23, 2022.
    29 changes: 29 additions & 0 deletions non_maximum_suppression.py
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,29 @@
    def get_iou(box1, box2):
    ar1 = (box1[3] - box1[1]) * (box1[4] - box1[2])
    ar2 = (box2[3] - box2[1]) * (box2[4] - box2[2])
    x1 = max(box1[1], box2[1])
    y1 = max(box1[2], box2[2])
    x2 = min(box1[3], box2[3])
    y2 = min(box1[4], box2[4])
    inter = (x2-x1)*(y2-y1)
    iou = inter/(ar1+ar2-inter)
    return iou


    def non_max_supression(boxes, iou_val=0.1, max_boxes=10000):
    """
    boxes - [(score, x1, y1, x2, y2), ...]
    """
    boxes = sorted(boxes, key = lambda x:x[0], reverse=True)
    if max_boxes>0:
    boxes = boxes[:max_boxes]
    rboxes = []
    while len(boxes)>0:
    cur_box = boxes.pop(0)
    # recreate boxes
    boxes =[
    box for box in boxes
    if get_iou(box, cur_box)<iou_val
    ]
    rboxes.append(cur_box)
    return rboxes