Add ReconcileAndFilterFlows processor.

PiperOrigin-RevId: 673524211

Author: timblakely

processor/flow.py

@@ -15,16 +15,25 @@
 """Flow field estimation from SOFIMA."""
 
 import dataclasses
-from typing import Optional
-
+from typing import Any, Sequence
+from connectomics.common import beam_utils
 from connectomics.common import bounding_box
 from connectomics.common import utils
-from connectomics.volume import mask
+from connectomics.volume import base
+from connectomics.volume import mask as mask_lib
+from connectomics.volume import metadata
 from connectomics.volume import subvolume
 from connectomics.volume import subvolume_processor
 import dataclasses_json
 import numpy as np
+from scipy import interpolate
 from sofima import flow_field
+from sofima import flow_utils
+from sofima import map_utils
+
+
+Subvolume = subvolume.Subvolume
+SubvolumeOrMany = Subvolume | list[Subvolume]
 
 
 class EstimateFlow(subvolume_processor.SubvolumeProcessor):
@@ -77,9 +86,9 @@ class EstimateFlowConfig(utils.NPDataClassJsonMixin):
     stride: int
     z_stride: int = 1
     fixed_current: bool = False
-    mask_configs: Optional[mask.MaskConfigs] = None
+    mask_configs: mask_lib.MaskConfigs | None = None
     mask_only_for_patch_selection: bool = False
-    selection_mask_configs: Optional[mask.MaskConfigs] = None
+    selection_mask_configs: mask_lib.MaskConfigs | None = None
     batch_size: int = 1024
 
   _config: EstimateFlowConfig
@@ -143,7 +152,7 @@ def process(self, subvol: subvolume.Subvolume) -> subvolume.Subvolume:
     if config.mask_configs is not None:
       # TODO(blakely): Remove the unused lambda here and below when the external
       # paths support DecoratorSpecs.
-      initial_mask = mask.build_mask(
+      initial_mask = mask_lib.build_mask(
           config.mask_configs, subvol.bbox, lambda x: x
       )
 
@@ -160,7 +169,7 @@ def process(self, subvol: subvolume.Subvolume) -> subvolume.Subvolume:
           subvol.bbox.size - xy * config.patch_size + xy * config.stride
       ) / scale
       sel_box = bounding_box.BoundingBox(sel_start, sel_size)
-      sel_mask = mask.build_mask(
+      sel_mask = mask_lib.build_mask(
           config.selection_mask_configs, sel_box, lambda x: x
       )
 
@@ -259,3 +268,238 @@ def expected_output_box(
         + self._config.stride
     ) // self._config.stride
     return bounding_box.BoundingBox(start, size)
+
+
+# TODO(blakely): Remove references to volinfos in favor of metadata
+class ReconcileAndFilterFlows(subvolume_processor.SubvolumeProcessor):
+  """Filters 4-channel or 3-channel flow volumes.
+
+  The input flow volume(s) (generated by EstimateFlow) are filtered to
+  only retain 'valid' entries fulfilling local consistency and estimation
+  confidence criteria. If additional (lower-resolution) flow estimates
+  are provided via 'flow_volinfos', they are used to fill any flow
+  entries considered 'invalid' after filtering the higher resolution
+  results.
+  """
+
+  crop_at_borders = False
+
+  @dataclasses_json.dataclass_json
+  @dataclasses.dataclass(eq=True)
+  class ReconcileFlowsConfig(utils.NPDataClassJsonMixin):
+    """Configuration for ReconcileAndFilterFlows.
+
+    Attributes:
+      flow_volinfos: List or comma-separated string of volinfo paths, sorted in
+        ascending order of voxel size; a path can optionally be followed by
+        ':scale', which defines a divisor to apply to the corresponding flow
+        field. If the divisor is not specified, its value is inferred from the
+        pixel size ratio between the given flow field and the first flow field
+        on the list.
+      mask_configs: MaskConfigs proto in text format; masked voxels will be set
+        to nan (in both channels)
+      min_peak_ratio: See flow_utils.clean_flow.
+      min_peak_sharpness: See flow_utils.clean_flow.
+      max_magnitude: See flow_utils.clean_flow.
+      max_deviation: See flow_utils.clean_flow.
+      max_gradient: See flow_utils.clean_flow.
+      min_patch_size: See flow_utils.clean_flow.
+      multi_section: If generating a multi-section volume, the value of the 3rd
+        channel to initialize the output flow with
+      base_delta_z: If generating a multi-section volume, the value of the 3rd
+        channel to initialize the output flow with
+    """
+
+    flow_volinfos: Sequence[str] | str | None = None
+    mask_configs: str | mask_lib.MaskConfigs | None = None
+    min_peak_ratio: float = 1.6
+    min_peak_sharpness: float = 1.6
+    max_magnitude: float = 40
+    max_deviation: float = 10
+    max_gradient: float = 40
+    min_patch_size: int = 400
+    multi_section: bool = False
+    base_delta_z: int = 0
+
+  _config: ReconcileFlowsConfig
+
+  def __init__(
+      self,
+      config: ReconcileFlowsConfig,
+      input_volinfo_or_metadata: str | metadata.VolumeMetadata | None = None,
+  ):
+    """Constructor.
+
+    Args:
+      config: Parameters for ReconcileAndFilterFlows
+      input_volinfo_or_metadata: input volume with a voxel size equal or smaller
+        than the first volume in the flow_volinfos list
+    """
+    self._config = config
+
+    self._scales = [None]
+    self._metadata: list[metadata.VolumeMetadata] = []
+    if input_volinfo_or_metadata is not None:
+      self._metadata.append(self._get_metadata(input_volinfo_or_metadata))
+    if isinstance(config.flow_volinfos, str):
+      config.flow_volinfos = config.flow_volinfos.split(',')
+
+    for path in config.flow_volinfos:
+      path, _, scale = path.partition(':')
+      if scale:
+        scale = float(scale)
+      else:
+        scale = None
+
+      self._scales.append(scale)
+      self._metadata.append(self._get_metadata(path))
+
+    # Ensure that the volumes are correctly sorted.
+    for a, b in zip(self._metadata, self._metadata[1:]):
+      assert a.pixel_size.x <= b.pixel_size.x
+      assert a.pixel_size.y <= b.pixel_size.y
+      assert a.pixel_size.x / b.pixel_size.x == a.pixel_size.y / b.pixel_size.y
+      assert a.pixel_size.z == b.pixel_size.z
+
+    if config.mask_configs is not None:
+      if isinstance(config.mask_configs, str):
+        config.mask_configs = self._get_mask_configs(config.mask_configs)
+
+  def _open_volume(self, path: str) -> base.Volume:
+    """Returns a CZYX-shaped ndarray-like object."""
+    raise NotImplementedError(
+        'This function needs to be defined in a subclass.'
+    )
+
+  def _get_metadata(self, path) -> metadata.VolumeMetadata:
+    raise NotImplementedError(
+        'This function needs to be defined in a subclass.'
+    )
+
+  def _get_mask_configs(self, mask_configs: str) -> mask_lib.MaskConfigs:
+    raise NotImplementedError(
+        'This function needs to be defined in a subclass.'
+    )
+
+  def _build_mask(
+      self,
+      mask_configs: mask_lib.MaskConfigs,
+      box: bounding_box.BoundingBoxBase,
+  ) -> Any:
+    """Returns a CZYX-shaped ndarray-like object."""
+    raise NotImplementedError(
+        'This function needs to be defined in a subclass.'
+    )
+
+  def num_channels(self, input_channels=0):
+    del input_channels
+    return 2 if not self._config.multi_section else 3
+
+  def process(self, subvol: Subvolume) -> SubvolumeOrMany:
+    box = subvol.bbox
+    if self._config.mask_configs is not None:
+      mask = self._build_mask(self._config.mask_configs, box)
+    else:
+      mask = None
+
+    # Points in image space at which the base (highest resolution) flow
+    # is defined. Pixel values are assumed to correspond to the middle
+    # point of the pixel.
+    qy, qx = np.mgrid[: box.size[1], : box.size[0]]
+    qx = qx + box.start[0]
+    qy = qy + box.start[1]
+
+    flows = []
+    volumes = [self._open_volume(v) for v in self._metadata]
+
+    for i, (vol, mag_scale) in enumerate(zip(volumes, self._scales)):
+      if i > 0:
+        scale = self._metadata[0].pixel_size.x / self._metadata[i].pixel_size.x
+        assert scale <= 1.0
+        read_box = box.scale((scale, scale, 1))
+        if scale < 1:
+          read_box = read_box.adjusted_by(
+              start=-self._context[0], end=self._context[1]
+          )
+        read_box = vol.clip_box_to_volume(read_box)
+        assert read_box is not None
+      else:
+        scale = 1
+        read_box = box
+
+      with beam_utils.timer_counter(
+          'reconcile-flows', 'time-volstore-load-%d' % i
+      ):
+        flow = vol[read_box.to_slice4d()]
+
+      with beam_utils.timer_counter('reconcile-flows', 'time-clean-%d' % i):
+        flow = flow_utils.clean_flow(
+            flow,
+            self._config.min_peak_ratio,
+            self._config.min_peak_sharpness,
+            self._config.max_magnitude,
+            self._config.max_deviation,
+        )
+
+      if i == 0 or scale == 1:
+        if self._config.multi_section and flow.shape[0] != 3:
+          shape = np.array(flow.shape)
+          shape[0] = 3
+          nflow = np.full(shape, np.nan, dtype=flow.dtype)
+          nflow[:2, ...] = flow[:2, ...]
+          nflow[2, ...][np.isfinite(nflow[0, ...])] = self._config.base_delta_z
+          flow = nflow
+
+        flows.append(flow)
+        continue
+
+      # Upsample flow to the base resolution.
+      hires_flow = np.zeros_like(flows[0])
+
+      oy, ox = np.ogrid[: read_box.size[1], : read_box.size[0]]
+      ox = ox + read_box.start[0]
+      oy = oy + read_box.start[1]
+      ox = (ox / scale).ravel()
+      oy = (oy / scale).ravel()
+
+      if mag_scale is None:
+        mag_scale = scale
+
+      with beam_utils.timer_counter('reconcile-flows', 'time-upsample-%d' % i):
+        for z in range(flow.shape[1]):
+          rgi = interpolate.RegularGridInterpolator(
+              (oy, ox), flow[0, z, ...], method='nearest', bounds_error=False
+          )
+          invalid_mask = np.isnan(rgi((qy, qx)))
+
+          # We want to upsample the spatial components of the flow with
+          # at least linear interpolation. Doing so with RegularGridInterpolator
+          # in the presence of invalid entries (NaN) will cause the invalid
+          # regions to grow beyond what 'nearest' upsampling would generate.
+          # To avoid this, we use a resampling scheme with interpolation and
+          # mask out invalid entries as if the field was resampled in
+          # the 'nearest' interpolation mode.
+          resampled = map_utils.resample_map(
+              flow[:2, z : z + 1, ...], read_box, box, 1 / scale, 1  #
+          )
+          hires_flow[:2, z : z + 1, ...] = resampled / mag_scale
+          hires_flow[0, z, ...][invalid_mask] = np.nan
+          hires_flow[1, z, ...][invalid_mask] = np.nan
+
+          for c in range(2, self.num_channels()):
+            rgi = interpolate.RegularGridInterpolator(
+                (oy, ox), flow[c, z, ...], method='nearest', bounds_error=False
+            )
+            hires_flow[c, z, ...] = rgi((qy, qx)).astype(np.float32)
+
+      if mask is not None:
+        flow_utils.apply_mask(hires_flow, mask)
+      flows.append(hires_flow)
+
+    ret = flow_utils.reconcile_flows(
+        flows,
+        self._config.max_gradient,
+        self._config.max_deviation,
+        self._config.min_patch_size,
+    )
+    return self.crop_box_and_data(box, ret)