From 87c34f7938b078aa472d104bfc558da20cd04eaa Mon Sep 17 00:00:00 2001 From: Ranbir Aulakh Date: Fri, 11 Oct 2024 19:00:21 +0000 Subject: [PATCH] test: new tests cases for imported ensemble_boxes_nms class --- .../model_runner/common/ensemble_boxes_nms.py | 75 ++++++----- .../common/test_ensemble_boxes_nms.py | 118 ++++++++++++++++++ 2 files changed, 161 insertions(+), 32 deletions(-) create mode 100644 test/aws/osml/model_runner/common/test_ensemble_boxes_nms.py diff --git a/src/aws/osml/model_runner/common/ensemble_boxes_nms.py b/src/aws/osml/model_runner/common/ensemble_boxes_nms.py index 909023ca..5d8d631c 100644 --- a/src/aws/osml/model_runner/common/ensemble_boxes_nms.py +++ b/src/aws/osml/model_runner/common/ensemble_boxes_nms.py @@ -1,16 +1,20 @@ -# Copyright 2024 Amazon.com, Inc. or its affiliates. +# Copyright 2024 Amazon.com, Inc. or its affiliates. """ -This file contains code for performing non-maximum suppression (NMS) on bounding boxes obtained from multiple object detection models. The code was originally taken from the following GitHub repository: +Non-Maximum Suppression (NMS) and Soft-NMS implementation for bounding boxes. -https://github.com/ZFTurbo/Weighted-Boxes-Fusion/blob/master/ensemble_boxes/ensemble_boxes_nms.py +This module implements standard NMS, linear Soft-NMS, and Gaussian Soft-NMS +for bounding boxes with normalized coordinates. It supports weighted scores +and multiple labels. -The original author is 'ZFTurbo' (https://kaggle.com/zfturbo). +Original implementation inspired by: +- https://github.com/ZFTurbo/Weighted-Boxes-Fusion -This implementation provides functions for standard NMS, linear soft-NMS, and gaussian soft-NMS. It also includes a method for preparing the boxes, scores, and labels before applying NMS. +Author: ZFTurbo (https://kaggle.com/zfturbo) +Refactored for internal use in OSML. """ -from typing import List, Optional, Tuple +from typing import Any, List, Optional, Tuple import numpy as np from numba import jit @@ -20,9 +24,9 @@ def prepare_boxes(boxes: np.ndarray, scores: np.ndarray, labels: np.ndarray) -> """ Prepare boxes by correcting invalid coordinates and removing boxes with zero area. - :param boxes: Array of shape (N, 4) with box coordinates in the format [x1, y1, x2, y2], where all values are normalized [0, 1]. - :param scores: Array of shape (N,) with confidence scores for each box. - :param labels: Array of shape (N,) with labels for each box. + :param boxes: Array of shape (N, 4) with box coordinates, [x1, y1, x2, y2], where all values are normalized [0, 1]. + :param scores: Array of shape (N, ) with confidence scores for each box. + :param labels: Array of shape (N, ) with labels for each box. :return: Tuple containing the filtered and corrected boxes, scores, and labels. """ @@ -58,14 +62,14 @@ def prepare_boxes(boxes: np.ndarray, scores: np.ndarray, labels: np.ndarray) -> return result_boxes, scores, labels -def cpu_soft_nms_float(dets: np.ndarray, sc: np.ndarray, Nt: float, sigma: float, thresh: float, method: int) -> np.ndarray: +def cpu_soft_nms_float(dets: np.ndarray, sc: np.ndarray, nt: float, sigma: float, thresh: float, method: int) -> np.ndarray: """ Based on: https://github.com/DocF/Soft-NMS/blob/master/soft_nms.py It's different from original soft-NMS because we have float coordinates on range [0; 1] :param dets: boxes format [x1, y1, x2, y2] :param sc: scores for boxes - :param Nt: required iou + :param nt: required iou :param sigma: Sigma value for Gaussian soft-NMS. :param thresh: Score threshold to filter boxes. :param method: 1 - linear soft-NMS, 2 - gaussian soft-NMS, 3 - standard NMS @@ -74,8 +78,8 @@ def cpu_soft_nms_float(dets: np.ndarray, sc: np.ndarray, Nt: float, sigma: float """ # indexes concatenate boxes with the last column - N = dets.shape[0] - indexes = np.array([np.arange(N)]) + n = dets.shape[0] + indexes = np.array([np.arange(n)]) dets = np.concatenate((dets, indexes.T), axis=1) # the order of boxes coordinate is [y1, x1, y2, x2] @@ -86,15 +90,15 @@ def cpu_soft_nms_float(dets: np.ndarray, sc: np.ndarray, Nt: float, sigma: float scores = sc areas = (x2 - x1) * (y2 - y1) - for i in range(N): + for i in range(n): # intermediate parameters for later parameters exchange - tBD = dets[i, :].copy() + tbd = dets[i, :].copy() tscore = scores[i].copy() tarea = areas[i].copy() pos = i + 1 # - if i != N - 1: + if i != n - 1: maxscore = np.max(scores[pos:], axis=0) maxpos = np.argmax(scores[pos:], axis=0) else: @@ -102,16 +106,16 @@ def cpu_soft_nms_float(dets: np.ndarray, sc: np.ndarray, Nt: float, sigma: float maxpos = 0 if tscore < maxscore: dets[i, :] = dets[maxpos + i + 1, :] - dets[maxpos + i + 1, :] = tBD - tBD = dets[i, :] + dets[maxpos + i + 1, :] = tbd + # tbd = dets[i, :] scores[i] = scores[maxpos + i + 1] scores[maxpos + i + 1] = tscore - tscore = scores[i] + # tscore = scores[i] areas[i] = areas[maxpos + i + 1] areas[maxpos + i + 1] = tarea - tarea = areas[i] + # tarea = areas[i] # IoU calculate xx1 = np.maximum(dets[i, 1], dets[pos:, 1]) @@ -127,12 +131,12 @@ def cpu_soft_nms_float(dets: np.ndarray, sc: np.ndarray, Nt: float, sigma: float # Three methods: 1.linear 2.gaussian 3.original NMS if method == 1: # linear weight = np.ones(ovr.shape) - weight[ovr > Nt] = weight[ovr > Nt] - ovr[ovr > Nt] + weight[ovr > nt] = weight[ovr > nt] - ovr[ovr > nt] elif method == 2: # gaussian weight = np.exp(-(ovr * ovr) / sigma) else: # original NMS weight = np.ones(ovr.shape) - weight[ovr > Nt] = 0 + weight[ovr > nt] = 0 scores[pos:] = weight * scores[pos:] @@ -143,7 +147,7 @@ def cpu_soft_nms_float(dets: np.ndarray, sc: np.ndarray, Nt: float, sigma: float @jit(nopython=True) -def nms_float_fast(dets: np.ndarray, scores: np.ndarray, thresh: float) -> np.ndarray: +def nms_fast(dets: np.ndarray, scores: np.ndarray, thresh: float) -> list[np.ndarray[Any, Any]]: """ It's different from original nms because we have float coordinates on range [0; 1] @@ -193,7 +197,8 @@ def nms_method( """ Perform NMS on a list of boxes, scores, and labels from multiple models. - :param boxes: list of boxes predictions from each model, each box is 4 numbers. It has 3 dimensions (models_number, model_preds, 4). Order of boxes: x1, y1, x2, y2. We expect float normalized coordinates [0; 1]. + :param boxes: list of boxes predictions from each model, each box is 4 numbers. It has 3 dimensions + (models_number, model_preds, 4). Order of boxes: x1, y1, x2, y2. We expect float normalized coordinates [0; 1]. :param scores: list of scores for each model. :param labels: list of labels for each model. :param method: 1 - linear soft-NMS, 2 - gaussian soft-NMS, 3 - standard NMS. @@ -205,6 +210,10 @@ def nms_method( :return: tuple of (boxes, scores, labels) after NMS. """ + # Validate input lengths + if not (len(boxes) == len(scores) == len(labels)): + raise ValueError(f"Input lengths must match: boxes={len(boxes)}, scores={len(scores)}, labels={len(labels)}") + # If weights are specified if weights is not None: if len(boxes) != len(weights): @@ -246,19 +255,19 @@ def nms_method( final_boxes = [] final_scores = [] final_labels = [] - for l in unique_labels: - condition = labels == l + for label in unique_labels: + condition = labels == label boxes_by_label = boxes[condition] scores_by_label = scores[condition] - labels_by_label = np.array([l] * len(boxes_by_label)) + labels_by_label = np.array([label] * len(boxes_by_label)) if method != 3: keep = cpu_soft_nms_float( - boxes_by_label.copy(), scores_by_label.copy(), Nt=iou_thr, sigma=sigma, thresh=thresh, method=method + boxes_by_label.copy(), scores_by_label.copy(), nt=iou_thr, sigma=sigma, thresh=thresh, method=method ) else: # Use faster function - keep = nms_float_fast(boxes_by_label, scores_by_label, thresh=iou_thr) + keep = nms_fast(boxes_by_label, scores_by_label, thresh=iou_thr) final_boxes.append(boxes_by_label[keep]) final_scores.append(scores_by_label[keep]) @@ -280,7 +289,8 @@ def nms( """ Short call for standard NMS - :param boxes: list of boxes predictions from each model, each box is 4 numbers. It has 3 dimensions (models_number, model_preds, 4). Order of boxes: x1, y1, x2, y2. We expect float normalized coordinates [0; 1]. + :param boxes: list of boxes predictions from each model, each box is 4 numbers. It has 3 dimensions (models_number, + model_preds, 4). Order of boxes: x1, y1, x2, y2. We expect float normalized coordinates [0; 1]. :param scores: list of scores for each model. :param labels: list of labels for each model. :param iou_thr: IoU threshold value for boxes. @@ -304,7 +314,8 @@ def soft_nms( """ Perform soft-NMS on the given set of boxes for each label. - :param boxes: list of boxes predictions from each model, each box is 4 numbers. It has 3 dimensions (models_number, model_preds, 4). Order of boxes: x1, y1, x2, y2. We expect float normalized coordinates [0; 1]. + :param boxes: list of boxes predictions from each model, each box is 4 numbers. It has 3 dimensions + (models_number, model_preds, 4). Order of boxes: x1, y1, x2, y2. We expect float normalized coordinates [0; 1]. :param scores: list of scores for each model. :param labels: list of labels for each model. :param method: 1 - linear soft-NMS, 2 - gaussian soft-NMS. @@ -315,4 +326,4 @@ def soft_nms( :return: Tuple containing the final boxes, scores, and labels after soft-NMS. """ - return nms_method(boxes, scores, labels, method=method, iou_thr=iou_thr, sigma=sigma, thresh=thresh, weights=weights) + return nms_method(boxes, scores, labels, method, iou_thr, sigma, thresh, weights) diff --git a/test/aws/osml/model_runner/common/test_ensemble_boxes_nms.py b/test/aws/osml/model_runner/common/test_ensemble_boxes_nms.py new file mode 100644 index 00000000..c0ac9988 --- /dev/null +++ b/test/aws/osml/model_runner/common/test_ensemble_boxes_nms.py @@ -0,0 +1,118 @@ +# Copyright 2024 Amazon.com, Inc. or its affiliates. + +import unittest + +import numpy as np +import pytest + + +class TestNMSMethods(unittest.TestCase): + """ + Unit tests for the Non-Maximum Suppression (NMS) and Soft-NMS functions. + """ + + def setUp(self): + """ + Sets up mock bounding boxes, scores, and labels for testing. + """ + # Bounding boxes (x1, y1, x2, y2) and scores from two models + self.boxes = [ + np.array([[0.1, 0.1, 0.4, 0.4], [0.15, 0.15, 0.45, 0.45], [0.6, 0.6, 0.9, 0.9]]), # Model 1 + np.array([[0.2, 0.2, 0.5, 0.5], [0.7, 0.7, 1.0, 1.0]]), # Model 2 + ] + self.scores = [ + np.array([0.9, 0.85, 0.6]), # Scores for Model 1 + np.array([0.8, 0.7]), # Scores for Model 2 + ] + self.labels = [ + np.array([1, 1, 2]), # Labels for Model 1 + np.array([1, 2]), # Labels for Model 2 + ] + + def test_prepare_boxes(self): + """ + Test the prepare_boxes function to ensure it: + 1. Corrects invalid box coordinates. + 2. Removes boxes with zero area. + """ + from aws.osml.model_runner.common.ensemble_boxes_nms import prepare_boxes + + # Create invalid boxes with out-of-bound coordinates and zero area + invalid_boxes = np.array([[-0.1, 0.2, 1.1, 1.2], [0.5, 0.5, 0.5, 0.5]]) + invalid_scores = np.array([0.9, 0.8]) + invalid_labels = np.array([1, 1]) + + filtered_boxes, filtered_scores, filtered_labels = prepare_boxes(invalid_boxes, invalid_scores, invalid_labels) + + # Assertions + assert filtered_boxes.shape[0] == 1 + assert np.all(filtered_boxes >= 0) and np.all(filtered_boxes <= 1) + + def test_nms(self): + """ + Test the standard NMS function to ensure it suppresses overlapping boxes + based on an IoU threshold of 0.5. + """ + from aws.osml.model_runner.common.ensemble_boxes_nms import nms + + final_boxes, final_scores, final_labels = nms(self.boxes, self.scores, self.labels, 0.5) + + # Assertions + assert final_boxes.shape[0] == 4 + + def test_soft_nms(self): + """ + Test the Soft-NMS function with the linear method (method=1). + """ + from aws.osml.model_runner.common.ensemble_boxes_nms import soft_nms + + final_boxes, final_scores, final_labels = soft_nms(self.boxes, self.scores, self.labels, 1, 0.5) + + # Assertions + assert final_boxes.shape[0] == 5 + + def test_nms_fast(self): + """ + Test the optimized NMS implementation (nms_fast) for speed and correctness. + """ + from aws.osml.model_runner.common.ensemble_boxes_nms import nms_fast + + dets = np.array([[0.1, 0.1, 0.4, 0.4], [0.15, 0.15, 0.45, 0.45], [0.6, 0.6, 0.9, 0.9]]) + scores = np.array([0.9, 0.85, 0.6]) + + keep = nms_fast(dets, scores, 0.5) + + # Assertions + assert len(keep) == 2 + + def test_nms_with_weights(self): + """ + Test the NMS function with model weights applied to scores. + """ + from aws.osml.model_runner.common.ensemble_boxes_nms import nms + + weights = [0.5, 0.5] # Apply equal weights to both models + final_boxes, final_scores, final_labels = nms(self.boxes, self.scores, self.labels, 0.5, weights=weights) + + # Assertions + assert final_boxes.shape[0] == 4 + assert np.all(final_scores <= 1.0) # Scores should remain normalized + + def test_invalid_input_lengths(self): + """ + Test that NMS raises a ValueError when input lengths are mismatched. + """ + from aws.osml.model_runner.common.ensemble_boxes_nms import nms + + # Mismatched input: boxes have fewer entries than scores and labels + invalid_boxes = [np.array([[0.1, 0.1, 0.4, 0.4]])] # 1 box + invalid_scores = [np.array([0.9, 0.8])] # 2 scores + invalid_labels = [np.array([1, 2])] # 2 labels + + # Verify that a ValueError is raised with a clear message + with pytest.raises(ValueError): + nms(invalid_boxes, invalid_scores, invalid_labels, 0.5) + + +if __name__ == "__main__": + unittest.main()