Add more unit tests.

PiperOrigin-RevId: 874047196

Author: mjanusz

ffn/input/load_data_test.py

@@ -0,0 +1,20 @@
+# Copyright 2024 Google Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from absl.testing import absltest
+
+
+if __name__ == '__main__':
+  absltest.main()

ffn/input/volume.py

@@ -18,12 +18,14 @@
 import copy
 import dataclasses
 import functools as ft
+import hashlib
 from typing import Any, Callable, Sequence, TypeVar
 
 from absl import logging
 import array_record
 from connectomics.common import array
 from connectomics.common import bounding_box
+from connectomics.common import box_generator
 from connectomics.common import io_utils
 from ffn.input import segmentation
 from ffn.training import augmentation
@@ -193,11 +195,14 @@ def _load_data(
   ret = dict(ex)
   dtype_remap = {tf.uint64: tf.int64}
 
+  ret['coord'] = tf.reshape(ex['coord'], [1, 3])
+  ret['volname'] = tf.reshape(ex['volname'], [1])
+
   for name, vol in config.volumes.items():
     if vol.oob_mask:
       ret[name] = inputs.make_oob_mask(
-          ex['coord'],
-          ex['volname'],
+          ret['coord'],
+          ret['volname'],
           shape=vol.load_shape,
           volinfo_map_string=get_path_str(vol.paths),
       )
@@ -312,7 +317,7 @@ def _filter_coordinates_by_bbox(
 ) -> tf.Tensor:
   ret = tf.numpy_function(
       lambda c, v: _coord_in_bboxes_np(c, v, bboxes),
-      [item['coord'][0], item['volname'][0]],
+      [item['coord'][0], tf.reshape(item['volname'], [-1])[0]],
       tf.bool,
   )
   ret.set_shape([])

ffn/input/volume_test.py

@@ -0,0 +1,20 @@
+# Copyright 2024 Google Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from absl.testing import absltest
+
+
+if __name__ == '__main__':
+  absltest.main()

ffn/training/inputs.py

@@ -785,8 +785,10 @@ def filter_oob(
     use_bboxes: bool = True
 ) -> tf.Tensor:
   radius = np.floor_divide(patch_size, 2)
+  coord = tf.reshape(item['coord'], [1, 3])
+  volname = tf.reshape(item['volname'], [1])
   return coordinates_in_bounds(
-      item['coord'], item['volname'], radius, volinfo_map_string,
+      coord, volname, radius, volinfo_map_string,
       use_bboxes
   )
 

ffn/training/inputs_test.py

@@ -0,0 +1,97 @@
+# Copyright 2026 Google Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import functools as ft
+import os
+
+
+from absl.testing import absltest
+from connectomics.common import bounding_box
+from connectomics.common import tuples
+from connectomics.volume import metadata
+from ffn.training import inputs
+import numpy as np
+import tensorflow.compat.v1 as tf
+
+tf.disable_eager_execution()
+
+
+class _FilterOobTestBase:
+
+  _volinfo_map_string: str
+
+  def test_filter_oob_in_bounds(self):
+    coord = tf.constant([[50, 50, 50]], dtype=tf.int64)
+    volname = tf.constant(['testvol'], dtype=tf.string)
+    item = {'coord': coord, 'volname': volname}
+    result = inputs.filter_oob(
+        item, self._volinfo_map_string, patch_size=[10, 10, 10]
+    )
+    with tf.Session() as sess:
+      self.assertTrue(sess.run(result))
+
+  def test_filter_oob_out_of_bounds(self):
+    coord = tf.constant([[0, 0, 0]], dtype=tf.int64)
+    volname = tf.constant(['testvol'], dtype=tf.string)
+    item = {'coord': coord, 'volname': volname}
+    result = inputs.filter_oob(
+        item, self._volinfo_map_string, patch_size=[10, 10, 10]
+    )
+    with tf.Session() as sess:
+      self.assertFalse(sess.run(result))
+
+  def test_filter_oob_in_dataset_filter(self):
+    ds = tf.data.Dataset.from_tensors({
+        'coord': tf.constant([[50, 50, 50]], dtype=tf.int64),
+        'volname': tf.constant(['testvol'], dtype=tf.string),
+    })
+    ds = ds.filter(
+        ft.partial(
+            inputs.filter_oob,
+            volinfo_map_string=self._volinfo_map_string,
+            patch_size=[10, 10, 10],
+        )
+    )
+    iterator = tf.data.make_one_shot_iterator(ds)
+    item = iterator.get_next()
+    with tf.Session() as sess:
+      result = sess.run(item)
+      np.testing.assert_array_equal(result['coord'], [[50, 50, 50]])
+
+
+class FilterOobMetadataJsonTest(_FilterOobTestBase, absltest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    tf.reset_default_graph()
+
+    self._tmpdir = self.create_tempdir().full_path
+    meta = metadata.VolumeMetadata(
+        path='none',
+        volume_size=tuples.XYZ(100, 100, 100),
+        pixel_size=tuples.XYZ(8, 8, 30),
+        bounding_boxes=[
+            bounding_box.BoundingBox(start=(0, 0, 0), size=(100, 100, 100))
+        ],
+    )
+    self._metadata_path = os.path.join(self._tmpdir, 'metadata.json')
+    with open(self._metadata_path, 'w') as f:
+      f.write(meta.to_json())
+
+    self._volinfo_map_string = f'testvol:{self._metadata_path}'
+
+
+if __name__ == '__main__':
+  absltest.main()