Updates to EstimateFlow stages.

PiperOrigin-RevId: 673530607

connectomics/common/beam_utils.py

@@ -31,6 +31,11 @@ def timer_counter(namespace, name):
       int((time.time() - start) * 1e3))
 
 
+def counter(namespace, name):
+  """Returns a counter."""
+  return beam.metrics.Metrics.counter(namespace, name)
+
+
 class MustFollow(beam.PTransform):
   """Pass-through input, but enforces deferred processing."""
 

connectomics/common/utils.py

@@ -16,6 +16,7 @@
 
 import contextlib
 import itertools
+import re
 import time
 
 from absl import logging
@@ -92,3 +93,18 @@ def report_time(name):
   finally:
     duration = time.time() - start
     logging.info('time[%s] = %.6f', name, duration)
+
+
+_PASCAL_TO_KEBAB_RE = re.compile(
+    r"""
+    (?<=[a-z])(?=[A-Z])
+    |
+    (?<=[A-Z])(?=[A-Z][a-z])
+    """,
+    re.X,
+)
+
+
+def pascal_to_kebab(name: str) -> str:
+  """Converts a PascalCase name to kebab-case."""
+  return _PASCAL_TO_KEBAB_RE.sub('-', name).lower()

connectomics/volume/subvolume_processor.py

@@ -25,6 +25,7 @@
 from connectomics.common import bounding_box
 from connectomics.common import counters
 from connectomics.common import file
+from connectomics.common import utils
 from connectomics.volume import descriptor
 from connectomics.volume import subvolume
 import dataclasses_json
@@ -131,6 +132,12 @@ class SubvolumeProcessor:
   # only uses the type and geometry of the array for further processing.
   ignores_input_data = False
 
+  # Namespace to use for counters. Overridable by subclasses; defaults to
+  # kebab-case of the class name.
+  @property
+  def namespace(self) -> str:
+    return utils.pascal_to_kebab(type(self).__name__)
+
   def output_type(self, input_type: Union[np.uint8, np.uint64, np.float32]):
     """Returns Numpy output type of self.process for given input_type.