Minor API and typo cleanups.

PiperOrigin-RevId: 874530870

Author: mjanusz

ffn/input/volume.py

@@ -153,7 +153,7 @@ def _update_config_for_augmentation(
         )
 
       if vol_cfg.filter_shape != (1, 1, 1) and vol_cfg.filter_shape is not None:
-        vol_cfg.filer_shape = tuple(
+        vol_cfg.filter_shape = tuple(
             augmentation.input_size_for_rotated_output(
                 vol_cfg.filter_shape, voxel_size
             )
@@ -174,7 +174,7 @@ def _postprocess_augmented_data(
       shape = vol_cfg.load_shape
 
       def _update_array(x, name=name, shape=shape):
-        setattr(x, name, mask.crop(x[name], (0, 0, 0), shape))
+        x[name] = mask.crop(x[name], (0, 0, 0), shape)
         return x
 
       ds = ds.map(

ffn/training/augmentation.py

@@ -369,7 +369,7 @@ def random_contrast_brightness_adjustment(
   Returns:
     tf.Tensor: Adjusted Tensor.
   """
-  adjust_tensor = tf.identity(input_tensor)
+  adjust_tensor = input_tensor
   if contrast_factor_range:
     min_contrast, max_contrast = contrast_factor_range
     contrast_factor = tf.random.uniform([], min_contrast, max_contrast)
@@ -383,7 +383,7 @@ def random_contrast_brightness_adjustment(
   if apply_adjustment_to == 'foreground':
     adjust_tensor = tf.where(seg_tensor > 0, adjust_tensor, input_tensor)
   elif apply_adjustment_to == 'background':
-    adjust_tensor = tf.where(seg_tensor <= 0, input_tensor, adjust_tensor)
+    adjust_tensor = tf.where(seg_tensor <= 0, adjust_tensor, input_tensor)
   return adjust_tensor
 
 
@@ -451,15 +451,15 @@ def __init__(
 
     if self.reflectable_axes.size > 0:
       self.reflect_decisions = (
-          tf.random_uniform([len(self.reflectable_axes)], seed=reflection_seed)
+          tf.random.uniform([len(self.reflectable_axes)], seed=reflection_seed)
           > 0.5
       )
       self.reflected_axes = tf.boolean_mask(
           self.reflectable_axes, self.reflect_decisions
       )
 
     if self.permutable_axes.size > 0:
-      self.permutation = tf.random_shuffle(
+      self.permutation = tf.random.shuffle(
           self.permutable_axes, seed=permutation_seed
       )
       # full_permutation must be a list rather than an np.array of int32 because

ffn/training/inputs.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Tensorflow Python ops and utilities for generating network inputs."""
 
+import functools
 import random
 import re
 from typing import Any, Callable, Optional, Sequence
@@ -553,7 +554,7 @@ def soften_labels(bool_labels, softness=0.05, scope='soften_labels'):
     Tensor with same shape as bool_labels with dtype `float32` and values 0.05
     for False and 0.95 for True.
   """
-  with tf.op_scope([bool_labels, softness], scope):
+  with tf.name_scope(scope):
     label_shape = tf.shape(bool_labels, name='label_shape')
     return tf.where(bool_labels,
                     tf.fill(label_shape, 1.0 - softness, name='soft_true'),
@@ -703,6 +704,35 @@ def sample(
   return sampled_dataset
 
 
+@functools.lru_cache(maxsize=None)
+def _parse_bounding_boxes(
+    volinfo_map_string: str, use_bboxes: bool = True
+) -> dict[bytes, list[bounding_box.BoundingBox]]:
+  boxes_by_volname = {}
+  for mapping in volinfo_map_string.split(','):
+    k, volinfo_path = mapping.split(':')
+    k = k.encode('utf-8')
+    assert k not in boxes_by_volname
+
+    if volinfo_path.endswith('metadata.json'):
+      f = open(volinfo_path, 'r')
+      meta = metadata.VolumeMetadata.from_json(f.read())
+      if use_bboxes:
+        bboxes = meta.bounding_boxes
+      else:
+        bboxes = [
+            bounding_box.BoundingBox(
+                (0, 0, 0),
+                (meta.volume_size.x, meta.volume_size.y, meta.volume_size.z),
+            )
+        ]
+      boxes_by_volname[k] = bboxes
+
+  if not boxes_by_volname:
+    raise ValueError('boxes_by_volname is empty.')
+  return boxes_by_volname
+
+
 def coordinates_in_bounds(
     coordinates: tf.Tensor,
     volname: tf.Tensor,
@@ -734,28 +764,7 @@ def coordinates_in_bounds(
     coordinates or an empty constant of shape `[0, 3]`, which can then be
     passed to batching (e.g. see tests).
   """
-  boxes_by_volname = {}
-  for mapping in volinfo_map_string.split(','):
-    k, volinfo_path = mapping.split(':')
-    k = k.encode('utf-8')
-    assert k not in boxes_by_volname
-
-    if volinfo_path.endswith('metadata.json'):
-      f = open(volinfo_path, 'r')
-      meta = metadata.VolumeMetadata.from_json(f.read())
-      if use_bboxes:
-        bboxes = meta.bounding_boxes
-      else:
-        bboxes = [
-            bounding_box.BoundingBox(
-                (0, 0, 0),
-                (meta.volume_size.x, meta.volume_size.y, meta.volume_size.z),
-            )
-        ]
-      boxes_by_volname[k] = bboxes
-
-  if not boxes_by_volname:
-    raise ValueError('boxes_by_volname is empty.')
+  boxes_by_volname = _parse_bounding_boxes(volinfo_map_string, use_bboxes)
 
   def _in_bounds_fn(coordinates, volname):
     boxes = boxes_by_volname[volname[0]]
@@ -767,8 +776,8 @@ def _in_bounds_fn(coordinates, volname):
     return False
 
   with tf.name_scope(name, values=[coordinates, volname]) as scope:
-    assert coordinates.shape_as_list() == [1, 3]
-    assert volname.shape_as_list() == [1]
+    assert coordinates.shape.as_list() == [1, 3]
+    assert volname.shape.as_list() == [1]
     in_bounds = tf.py_func(
         _in_bounds_fn,
         [coordinates, volname],

ffn/training/mask.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-"""Utilites for dealing with 2d and 3d object masks."""
+"""Utilities for dealing with 2d and 3d object masks."""
 
 from typing import Optional, Sequence
 

ffn/training/model.py

@@ -122,13 +122,13 @@ def set_up_sigmoid_pixelwise_loss(self, logits):
     pixel_loss *= self.loss_weights
     self.loss = tf.reduce_mean(pixel_loss)
     tf.summary.scalar('pixel_loss', self.loss)
-    self.loss = tf.verify_tensor_all_finite(self.loss, 'Invalid loss detected')
+    self.loss = tf.debugging.check_numerics(self.loss, 'Invalid loss detected')
 
   def set_up_optimizer(self, loss=None, max_gradient_entry_mag=0.7):
     """Sets up the training op for the model."""
     if loss is None:
       loss = self.loss
-    tf.summary.scalar('optimizer_loss', self.loss)
+    tf.summary.scalar('optimizer_loss', loss)
 
     opt = optimizer.optimizer_from_flags()
     self.opt = opt

ffn/training/tracker.py

@@ -361,7 +361,7 @@ def get_summaries(self) -> list[tf.Summary.Value]:
       for i, summary in enumerate(images):
         summary.tag += '/%d' % i
 
-    total_moves = sum(self.moves.tf_value)
+    total_moves = max(sum(self.moves.tf_value), 1)
     move_summaries = []
     for mt in MoveType:
       move_summaries.append(
@@ -377,14 +377,14 @@ def get_summaries(self) -> list[tf.Summary.Value]:
                 tag='fov/masked_voxel_fraction',
                 simple_value=(
                     self.fov_stats.tf_value[FovStat.MASKED_VOXELS]
-                    / self.fov_stats.tf_value[FovStat.TOTAL_VOXELS]
+                    / max(self.fov_stats.tf_value[FovStat.TOTAL_VOXELS], 1)
                 ),
             ),
             tf.Summary.Value(
                 tag='fov/average_weight',
                 simple_value=(
                     self.fov_stats.tf_value[FovStat.WEIGHTS_SUM]
-                    / self.fov_stats.tf_value[FovStat.TOTAL_VOXELS]
+                    / max(self.fov_stats.tf_value[FovStat.TOTAL_VOXELS], 1)
                 ),
             ),
             tf.Summary.Value(
@@ -418,7 +418,7 @@ def get_summaries(self) -> list[tf.Summary.Value]:
     )
 
     for r, r_moves in self.moves_by_r.items():
-      total_moves = sum(r_moves.tf_value)
+      total_moves = max(sum(r_moves.tf_value), 1)
       summaries.extend([
           tf.Summary.Value(
               tag='moves/r=%d/correct' % r,