From f989bccd5b1c5bc7e3ca1a81243e2d4d25e641dc Mon Sep 17 00:00:00 2001
From: Tim Blakely <blakely@google.com>
Date: Wed, 6 Nov 2024 12:29:49 -0800
Subject: [PATCH] Move `downsample_area` into public connectomics repository.

PiperOrigin-RevId: 693823309
---
 connectomics/common/geom_utils.py | 54 +++++++++++++++++++++++++++++++
 1 file changed, 54 insertions(+)

diff --git a/connectomics/common/geom_utils.py b/connectomics/common/geom_utils.py
index 8f9e520..aadd755 100644
--- a/connectomics/common/geom_utils.py
+++ b/connectomics/common/geom_utils.py
@@ -18,6 +18,8 @@
 
 import numbers
 from typing import Optional, Sequence
+
+from connectomics.common import bounding_box
 import numpy as np
 
 
@@ -100,3 +102,55 @@ def point_query_integral_image(summed_volume_table: np.ndarray,
   else:
     raise NotImplementedError(
         'Only 2 and 3-dimensional integral images are supported.')
+
+
+def downsample_area(
+    svt: np.ndarray,
+    box: bounding_box.BoundingBox,
+    scale: np.ndarray,
+    dtype: np.dtype,
+    mask_svt: Optional[np.ndarray] = None,
+) -> tuple[bounding_box.BoundingBoxBase, np.ndarray]:
+  """Downsamples data by area-averaging.
+
+  Args:
+    svt: [z, y, x] summed-volume table for the data to downsample
+    box: bounding box from which the source data was originates
+    scale: xyz downsampling factors
+    dtype: data type for the output array
+    mask_svt: [z, y, x] summed-volume table for the mask
+
+  Returns:
+    bounding box for the downsampled subvolume, downsampled subvolume as a
+    [1, z', y', x']-shaped array
+  """
+  # The offset needs to be such that the first input voxel to the svt is
+  # evenly divisible by scale.
+  off = box.start - box.start // scale * scale
+  off[off > 0] = (scale - off)[off > 0]
+
+  scale_zyx = scale[::-1]
+  scale_vol = np.prod(scale)
+
+  # The value computed in ret below corresponds to an output cropped by
+  # scale - 1 on each side prior to striding, but we need the output to
+  # be cropped by the context expected by the processor.
+  ret = query_integral_image(
+      svt[off[2] :, off[1] :, off[0] :], diam=scale_zyx, stride=scale_zyx
+  )
+  if mask_svt is None:
+    ret = ret / scale_vol
+  else:
+    missing = query_integral_image(
+        mask_svt[off[2] :, off[1] :, off[0] :], diam=scale_zyx, stride=scale_zyx
+    )
+    norm = np.clip(scale_vol - missing, 1, None)
+    ret = ret / norm
+    ret[missing == scale_vol] = np.nan
+
+  ret = np.round(ret).astype(dtype)
+  out_box = bounding_box.BoundingBox(
+      start=(box.start + off) // scale, size=ret.shape[::-1]
+  )
+  ret = ret[np.newaxis, ...]
+  return out_box, ret