def NMSIndices(self, bboxes, scores, max_output_size, nms_iou_threshold=0.3, score_threshold=0.01): """Apply NMS to a series of 3d bounding boxes in 7-DOF format. Args: bboxes: A [num_boxes, 7] floating point Tensor of bounding boxes in [x, y, z, dx, dy, dz, phi] format. scores: A [num_boxes] floating point Tensor containing box scores. max_output_size: Maximum number of boxes to predict per input. nms_iou_threshold: IoU threshold to use when determining whether two boxes overlap for purposes of suppression. score_threshold: The score threshold passed to NMS that allows NMS to quickly ignore irrelevant boxes. Returns: The NMS indices and the mask of the padded indices. """ bboxes = py_utils.HasShape(bboxes, [-1, 7]) # Extract x, y, w, h, then convert to extrema. # # Note that we drop the rotation angle because we don't have an NMS # operation that takes rotation into account. bboxes_2d = tf.stack( [bboxes[:, 0], bboxes[:, 1], bboxes[:, 3], bboxes[:, 4]], axis=-1) bboxes_extrema = geometry.XYWHToBBoxes(bboxes_2d) # Compute NMS with padding; we use the padded version so this function can # be used in a map_fn. This function returns the scalar number of boxes # for each example. # # We use an IoU threshold of 0.3 since our anchor boxes have rotations # that make the default IoU threshold of 0.5 possibly too high. nms_index_padded, num_valid = tf.image.non_max_suppression_padded( bboxes_extrema, scores, iou_threshold=nms_iou_threshold, max_output_size=max_output_size, score_threshold=score_threshold, pad_to_max_output_size=True) # Return the mask of valid indices instead of just a scalar number. mask = tf.concat( [tf.ones([num_valid]), tf.zeros([max_output_size - num_valid])], axis=0) nms_index_padded = tf.where(mask > 0, nms_index_padded, tf.zeros_like(nms_index_padded)) return nms_index_padded, mask
def testXYWHBBoxesCentroid(self): with self.session(): xywh = np.tile( np.array([[[10, 20, 8, 6], [0, 0, 0, 0], [-10, -20, 2.4, 3.6]]]), (3, 2, 1)) bboxes = np.tile( np.array([[[17, 6, 23, 14], [0, 0, 0, 0], [-21.8, -11.2, -18.2, -8.8]]]), (3, 2, 1)) centroid = np.tile(np.array([[[10, 20], [0, 0], [-10, -20]]]), (3, 2, 1)) print('shapes = ', xywh.shape, bboxes.shape, centroid.shape) self.assertAllClose(geometry.XYWHToBBoxes(xywh).eval(), bboxes) self.assertAllClose(geometry.BBoxesToXYWH(bboxes).eval(), xywh) self.assertAllClose(geometry.BBoxesCentroid(bboxes).eval(), centroid)