def test_multiclass_nms(self): boxes = np.array( [[0.2, 0.4, 0.8, 0.8], [0.4, 0.2, 0.8, 0.8], [0.6, 0.0, 1.0, 1.0]], dtype=np.float32) mask0 = np.array([[0, 0, 0, 0, 0], [0, 0, 1, 1, 0], [0, 0, 1, 1, 0], [0, 0, 1, 1, 0], [0, 0, 0, 0, 0]], dtype=np.uint8) mask1 = np.array([[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 1, 1, 1, 0], [0, 1, 1, 1, 0], [0, 0, 0, 0, 0]], dtype=np.uint8) mask2 = np.array([[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [1, 1, 1, 1, 1], [1, 1, 1, 1, 1]], dtype=np.uint8) masks = np.stack([mask0, mask1, mask2]) box_mask_list = np_box_mask_list.BoxMaskList( box_data=boxes, mask_data=masks) scores = np.array([[-0.2, 0.1, 0.5, -0.4, 0.3], [0.7, -0.7, 0.6, 0.2, -0.9], [0.4, 0.34, -0.9, 0.2, 0.31]], dtype=np.float32) box_mask_list.add_field('scores', scores) box_mask_list_clean = np_box_mask_list_ops.multi_class_non_max_suppression( box_mask_list, score_thresh=0.25, iou_thresh=0.1, max_output_size=3) scores_clean = box_mask_list_clean.get_field('scores') classes_clean = box_mask_list_clean.get_field('classes') boxes = box_mask_list_clean.get() masks = box_mask_list_clean.get_masks() expected_scores = np.array([0.7, 0.6, 0.34, 0.31]) expected_classes = np.array([0, 2, 1, 4]) expected_boxes = np.array([[0.4, 0.2, 0.8, 0.8], [0.4, 0.2, 0.8, 0.8], [0.6, 0.0, 1.0, 1.0], [0.6, 0.0, 1.0, 1.0]], dtype=np.float32) self.assertAllClose(scores_clean, expected_scores) self.assertAllClose(classes_clean, expected_classes) self.assertAllClose(boxes, expected_boxes)
def test_multiclass_nms(self): boxes = np.array( [[0.2, 0.4, 0.8, 0.8], [0.4, 0.2, 0.8, 0.8], [0.6, 0.0, 1.0, 1.0]], dtype=np.float32) mask0 = np.array([[0, 0, 0, 0, 0], [0, 0, 1, 1, 0], [0, 0, 1, 1, 0], [0, 0, 1, 1, 0], [0, 0, 0, 0, 0]], dtype=np.uint8) mask1 = np.array([[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 1, 1, 1, 0], [0, 1, 1, 1, 0], [0, 0, 0, 0, 0]], dtype=np.uint8) mask2 = np.array([[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [1, 1, 1, 1, 1], [1, 1, 1, 1, 1]], dtype=np.uint8) masks = np.stack([mask0, mask1, mask2]) box_mask_list = np_box_mask_list.BoxMaskList( box_data=boxes, mask_data=masks) scores = np.array([[-0.2, 0.1, 0.5, -0.4, 0.3], [0.7, -0.7, 0.6, 0.2, -0.9], [0.4, 0.34, -0.9, 0.2, 0.31]], dtype=np.float32) box_mask_list.add_field('scores', scores) box_mask_list_clean = np_box_mask_list_ops.multi_class_non_max_suppression( box_mask_list, score_thresh=0.25, iou_thresh=0.1, max_output_size=3) scores_clean = box_mask_list_clean.get_field('scores') classes_clean = box_mask_list_clean.get_field('classes') boxes = box_mask_list_clean.get() masks = box_mask_list_clean.get_masks() expected_scores = np.array([0.7, 0.6, 0.34, 0.31]) expected_classes = np.array([0, 2, 1, 4]) expected_boxes = np.array([[0.4, 0.2, 0.8, 0.8], [0.4, 0.2, 0.8, 0.8], [0.6, 0.0, 1.0, 1.0], [0.6, 0.0, 1.0, 1.0]], dtype=np.float32) self.assertAllClose(scores_clean, expected_scores) self.assertAllClose(classes_clean, expected_classes) self.assertAllClose(boxes, expected_boxes)