From 8c7fcb6881193181f0c30c9ce5ca141f8f46744a Mon Sep 17 00:00:00 2001 From: Michal Januszewski Date: Wed, 9 Sep 2026 04:52:23 -0700 Subject: [PATCH] Accept 1d and 2d coord tensors in _filter_coordinates_by_bbox. PiperOrigin-RevId: 978474804 --- ffn/input/volume.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/ffn/input/volume.py b/ffn/input/volume.py index e8bd4d9..c6d0539 100644 --- a/ffn/input/volume.py +++ b/ffn/input/volume.py @@ -280,7 +280,7 @@ def sample_coordinates( 'coord': [1, 3] XYZ int64 array 'volname': [1] string array """ - if config.sampling.bounding_boxes: + if config.sampling.bounding_boxes: boxes_cfg = [] volume_names = [] # Compile boxes and volumes strings. @@ -317,7 +317,10 @@ def _filter_coordinates_by_bbox( ) -> tf.Tensor: ret = tf.numpy_function( lambda c, v: _coord_in_bboxes_np(c, v, bboxes), - [item['coord'][0], tf.reshape(item['volname'], [-1])[0]], + [ + tf.reshape(item['coord'], [-1, 3])[0], + tf.reshape(item['volname'], [-1])[0], + ], tf.bool, ) ret.set_shape([])