Filter cochlea with segmentation table and export lower resolution

schilling40 · schilling40 · commit 65ccdd743397 · 2025-07-22T03:58:50.000+02:00
diff --git a/flamingo_tools/segmentation/postprocessing.py b/flamingo_tools/segmentation/postprocessing.py
@@ -10,6 +10,7 @@
 import pandas as pd
 
 from elf.io import open_file
+from scipy.ndimage import distance_transform_edt, binary_dilation, binary_closing
 from scipy.sparse import csr_matrix
 from scipy.spatial import distance
 from scipy.spatial import cKDTree, ConvexHull
@@ -615,3 +616,177 @@ def postprocess_ihc_seg(
     table.loc[:, "component_labels"] = comp_labels
 
     return table
+
+
+def dilate_and_trim(
+    arr_orig: np.ndarray,
+    edt: np.ndarray,
+    iterations: int = 15,
+    offset: float = 0.4,
+) -> np.ndarray:
+    """Dilate and trim original binary array according to a
+    Euclidean Distance Trasform computed for a separate target array.
+
+    Args:
+        arr_orig: Original 3D binary array
+        edt: 3D array containing Euclidean Distance transform for guiding dilation
+        iterations: Number of iterations for dilations
+        offset: Offset for regulating dilation. value should be in range(0, 0.45)
+
+    Returns:
+        Dilated binary array
+    """
+    border_coords = [(1, 0, 0), (-1, 0, 0), (0, 1, 0), (0, -1, 0), (0, 0, 1), (0, 0, -1)]
+    for _ in range(iterations):
+        arr_dilated = binary_dilation(arr_orig)
+        for x in range(arr_dilated.shape[0]):
+            for y in range(arr_dilated.shape[1]):
+                for z in range(arr_dilated.shape[2]):
+                    if arr_dilated[x, y, z] != 0:
+                        if arr_orig[x, y, z] == 0:
+                            min_dist = float('inf')
+                            for dx, dy, dz in border_coords:
+                                nx, ny, nz = x+dx, y+dy, z+dz
+                                if arr_orig[nx, ny, nz] == 1:
+                                    min_dist = min([min_dist, edt[nx, ny, nz]])
+                            if edt[x, y, z] >= min_dist - offset:
+                                arr_dilated[x, y, z] = 0
+        arr_orig = arr_dilated
+    return arr_dilated
+
+
+def filter_cochlea_volume_single(
+    table: pd.DataFrame,
+    components: Optional[List[int]] = [1],
+    scale_factor: int = 48,
+    resolution: float = 0.38,
+    dilation_iterations: int = 12,
+    padding: int = 1200,
+) -> np.ndarray:
+    """Filter cochlea volume based on a segmentation table.
+    Centroids contained in the segmentation table are used to create a down-scaled binary array.
+    The array can be dilated.
+
+    Args:
+        table: Segmentation table.
+        components: Component labels for filtering segmentation table.
+        scale_factor: Down-sampling factor for filtering.
+        resolution: Resolution of pixel in µm.
+        dilation_iterations: Iterations for dilating binary segmentation mask. A negative value omits binary closing.
+        padding: Padding in pixel to apply to guessed dimensions based on centroid coordinates.
+
+    Returns:
+        Binary 3D array of filtered cochlea.
+    """
+    # filter components
+    if components is not None:
+        table = table[table["component_labels"].isin(components)]
+
+    # identify approximate input dimensions for down-scaling
+    centroids = list(zip(table["anchor_x"] / resolution,
+                         table["anchor_y"] / resolution,
+                         table["anchor_z"] / resolution))
+
+    # padding the array allows for dilation without worrying about array borders
+    max_x = table["anchor_x"].max() / resolution + padding
+    max_y = table["anchor_y"].max() / resolution + padding
+    max_z = table["anchor_z"].max() / resolution + padding
+    ref_dimensions = (max_x, max_y, max_z)
+
+    # down-scale arrays
+    array_downscaled = downscaled_centroids(centroids, ref_dimensions=ref_dimensions,
+                                            scale_factor=scale_factor, downsample_mode="capped")
+
+    if dilation_iterations > 0:
+        array_dilated = binary_dilation(array_downscaled, np.ones((3, 3, 3)), iterations=dilation_iterations)
+        return binary_closing(array_dilated, np.ones((3, 3, 3)), iterations=1)
+
+    elif dilation_iterations == 0:
+        return binary_closing(array_downscaled, np.ones((3, 3, 3)), iterations=1)
+
+    else:
+        return array_downscaled
+
+
+def filter_cochlea_volume(
+    sgn_table: pd.DataFrame,
+    ihc_table: pd.DataFrame,
+    sgn_components: Optional[List[int]] = [1],
+    ihc_components: Optional[List[int]] = [1],
+    scale_factor: int = 48,
+    resolution: float = 0.38,
+    dilation_iterations: int = 12,
+    padding: int = 1200,
+    dilation_method = "individual",
+) -> np.ndarray:
+    """Filter cochlea volume with SGN and IHC segmentation.
+    Centroids contained in the segmentation tables are used to create down-scaled binary arrays.
+    The arrays are then dilated using guided dilation to fill the section inbetween SGNs and IHCs.
+
+    Args:
+        sgn_table: SGN segmentation table.
+        ihc_table: IHC segmentation table.
+        sgn_components: Component labels for filtering SGN segmentation table.
+        ihc_components: Component labels for filtering IHC segmentation table.
+        scale_factor: Down-sampling factor for filtering.
+        resolution: Resolution of pixel in µm.
+        dilation_iterations: Iterations for dilating binary segmentation mask.
+        padding: Padding in pixel to apply to guessed dimensions based on centroid coordinates.
+        dilation_method: Dilation style for SGN and IHC segmentation, either 'individual', 'combined' or no dilation.
+
+    Returns:
+        Binary 3D array of filtered cochlea.
+    """
+    # filter components
+    if sgn_components is not None:
+        sgn_table = sgn_table[sgn_table["component_labels"].isin(sgn_components)]
+    if ihc_components is not None:
+        ihc_table = ihc_table[ihc_table["component_labels"].isin(ihc_components)]
+
+    # identify approximate input dimensions for down-scaling
+    centroids_sgn = list(zip(sgn_table["anchor_x"] / resolution,
+                             sgn_table["anchor_y"] / resolution,
+                             sgn_table["anchor_z"] / resolution))
+    centroids_ihc = list(zip(ihc_table["anchor_x"] / resolution,
+                             ihc_table["anchor_y"] / resolution,
+                             ihc_table["anchor_z"] / resolution))
+
+    # padding the array allows for dilation without worrying about array borders
+    max_x = max([sgn_table["anchor_x"].max(), ihc_table["anchor_x"].max()]) / resolution + padding
+    max_y = max([sgn_table["anchor_y"].max(), ihc_table["anchor_y"].max()]) / resolution + padding
+    max_z = max([sgn_table["anchor_z"].max(), ihc_table["anchor_z"].max()]) / resolution + padding
+    ref_dimensions = (max_x, max_y, max_z)
+
+    # down-scale arrays
+    array_downscaled_sgn = downscaled_centroids(centroids_sgn, ref_dimensions=ref_dimensions,
+                                                scale_factor=scale_factor, downsample_mode="capped")
+
+    array_downscaled_ihc = downscaled_centroids(centroids_ihc, ref_dimensions=ref_dimensions,
+                                                scale_factor=scale_factor, downsample_mode="capped")
+
+    # dilate down-scaled SGN array in direction of IHC segmentation
+    distance_from_sgn = distance_transform_edt(~array_downscaled_sgn.astype(bool))
+    iterations = 20
+    arr_dilated = dilate_and_trim(array_downscaled_ihc.copy(), distance_from_sgn, iterations=iterations, offset=0.4)
+
+    # dilate single structures first
+    if dilation_method == "individual":
+        ihc_dilated = binary_dilation(array_downscaled_ihc, np.ones((3, 3, 3)), iterations=dilation_iterations)
+        sgn_dilated = binary_dilation(array_downscaled_sgn, np.ones((3, 3, 3)), iterations=dilation_iterations)
+        combined_dilated = arr_dilated + ihc_dilated + sgn_dilated
+        combined_dilated[combined_dilated > 0] = 1
+        combined_dilated = binary_dilation(combined_dilated, np.ones((3, 3, 3)), iterations=1)
+
+    # dilate combined structure
+    elif dilation_method == "combined":
+        # combine SGN, IHC, and region between both to form output mask
+        combined_structure = arr_dilated + array_downscaled_ihc + array_downscaled_sgn
+        combined_structure[combined_structure > 0] = 1
+        combined_dilated = binary_dilation(combined_structure, np.ones((3, 3, 3)), iterations=dilation_iterations)
+
+    # no dilation of combined structure
+    else:
+        combined_dilated = arr_dilated + ihc_dilated + sgn_dilated
+        combined_dilated[combined_dilated > 0] = 1
+
+    return combined_dilated
diff --git a/scripts/export_lower_resolution.py b/scripts/export_lower_resolution.py
@@ -1,12 +1,14 @@
 import argparse
 import os
+from typing import List, Optional
 
 import numpy as np
 import pandas as pd
 import tifffile
 import zarr
 
 from flamingo_tools.s3_utils import get_s3_path, BUCKET_NAME, SERVICE_ENDPOINT
+from flamingo_tools.segmentation.postprocessing import filter_cochlea_volume, filter_cochlea_volume_single
 # from skimage.segmentation import relabel_sequential
 
 
@@ -26,37 +28,157 @@ def filter_component(fs, segmentation, cochlea, seg_name, components):
     return segmentation
 
 
+def filter_cochlea(
+    cochlea: str,
+    filter_cochlea_channels: str,
+    sgn_components: Optional[List[int]] = None,
+    ihc_components: Optional[List[int]] = None,
+    ds_factor: int = 24,
+    dilation_iterations: int = 8,
+) -> np.ndarray:
+    """Pre-process information for filtering cochlea volume based on segmentation table.
+    Differentiates between the input of a single channel of either IHC or SGN or if both are supplied.
+    If a single channel is given, the filtered volume contains
+    a down-sampled segmentation area, which has been dilated.
+    If both IHC and SGN segmentation are supplied, a more specialized dilation
+    is applied to ensure that the connecting volume is not filtered.
+
+    Args:
+        cochlea: Name of cochlea.
+        filter_cochlea_channels: Segmentation table(s) used for filtering.
+        sgn_components: Component labels for filtering SGN segmentation table.
+        ihc_components: Component labels for filtering IHC segmentation table.
+        ds_factor: Down-sampling factor for filtering.
+        dilation_iterations: Iterations for dilating binary segmentation mask.
+
+    Returns:
+        Binary 3D array of filtered cochlea
+    """
+    # we check if the supplied channels contain an SGN and IHC channel
+    sgn_channels = [ch for ch in filter_cochlea_channels if "SGN" in ch]
+    sgn_channel = None if len(sgn_channels) == 0 else sgn_channels[0]
+
+    ihc_channels = [ch for ch in filter_cochlea_channels if "IHC" in ch]
+    ihc_channel = None if len(ihc_channels) == 0 else ihc_channels[0]
+
+    if ihc_channel is None and sgn_channel is None:
+        raise ValueError("Channels supplied for filtering cochlea volume do not contain an IHC or SGN segmentation.")
+
+    if sgn_channel is not None:
+        internal_path = os.path.join(cochlea, "tables",  sgn_channel, "default.tsv")
+        tsv_path, fs = get_s3_path(internal_path, bucket_name=BUCKET_NAME, service_endpoint=SERVICE_ENDPOINT)
+        with fs.open(tsv_path, "r") as f:
+            table_sgn = pd.read_csv(f, sep="\t")
+
+    if ihc_channel is not None:
+        internal_path = os.path.join(cochlea, "tables",  ihc_channel, "default.tsv")
+        tsv_path, fs = get_s3_path(internal_path, bucket_name=BUCKET_NAME, service_endpoint=SERVICE_ENDPOINT)
+        with fs.open(tsv_path, "r") as f:
+            table_ihc = pd.read_csv(f, sep="\t")
+
+    if sgn_channel is None:
+        # filter based in IHC segmentation
+        return filter_cochlea_volume_single(table_ihc, components=ihc_components,
+                                            scale_factor=ds_factor, dilation_iterations=dilation_iterations)
+    elif ihc_channel is None:
+        # filter based on SGN segmentation
+        return filter_cochlea_volume_single(table_sgn, components=sgn_components,
+                                            scale_factor=ds_factor, dilation_iterations=dilation_iterations)
+    else:
+        # filter based on SGN and IHC segmentation with a specialized function
+        return filter_cochlea_volume(table_sgn, table_ihc,
+                                     sgn_components=sgn_components,
+                                     ihc_components=ihc_components,
+                                     scale_factor=ds_factor,
+                                     dilation_iterations=dilation_iterations)
+
+
+def upscale_volume(
+    target_data: np.ndarray,
+    downscaled_volume: np.ndarray,
+    upscale_factor: int,
+) -> np.ndarray:
+    """Up-scale binary 3D mask to dimensions of target data.
+    After an initial up-scaling, the dimensions are cropped or zero-padded to fit the target shape.
+
+    Args:
+        target_data: Reference data for up-scaling.
+        downscaled_volume: Down-scaled binary 3D array.
+        upscale_factor: Initial factor for up-scaling binary array.
+
+    Returns:
+        Resized binary array.
+    """
+    target_shape = target_data.shape
+    upscaled_filter = np.repeat(
+        np.repeat(
+            np.repeat(downscaled_volume, upscale_factor, axis=0),
+            upscale_factor, axis=1),
+        upscale_factor, axis=2)
+    resized = np.zeros(target_shape, dtype=target_data.dtype)
+    min_x, min_y, min_z = tuple(min(upscaled_filter.shape[i], target_shape[i]) for i in range(3))
+    resized[:min_x, :min_y, :min_z] = upscaled_filter[:min_x, :min_y, :min_z]
+    return resized
+
+
 def export_lower_resolution(args):
-    output_folder = os.path.join(args.output_folder, args.cochlea, f"scale{args.scale}")
-    os.makedirs(output_folder, exist_ok=True)
-
-    input_key = f"s{args.scale}"
-    for channel in args.channels:
-        out_path = os.path.join(output_folder, f"{channel}.tif")
-        if os.path.exists(out_path):
-            continue
-
-        print("Exporting channel", channel)
-        internal_path = os.path.join(args.cochlea, "images",  "ome-zarr", f"{channel}.ome.zarr")
-        s3_store, fs = get_s3_path(internal_path, bucket_name=BUCKET_NAME, service_endpoint=SERVICE_ENDPOINT)
-        with zarr.open(s3_store, mode="r") as f:
-            data = f[input_key][:]
-        print(data.shape)
-        if args.filter_by_components is not None:
-            data = filter_component(fs, data, args.cochlea, channel, args.filter_by_components)
-        if args.binarize:
-            data = (data > 0).astype("uint16")
-        tifffile.imwrite(out_path, data, bigtiff=True, compression="zlib")
+    # calculate single filter mask for all lower resolutions
+    if args.filter_cochlea_channels is not None:
+            ds_factor = 48
+            filter_volume = filter_cochlea(args.cochlea, args.filter_cochlea_channels,
+                                           sgn_components=args.filter_sgn_components,
+                                           ihc_components=args.filter_ihc_components,
+                                           dilation_iterations=args.filter_dilation_iterations, ds_factor=ds_factor)
+            filter_volume = np.transpose(filter_volume, (2,1,0))
+
+    # iterate through exporting lower resolutions
+    for scale in args.scale:
+        if args.filter_cochlea_channels is not None:
+            output_folder = os.path.join(args.output_folder, args.cochlea,
+                                        f"scale{scale}_dilation{args.filter_dilation_iterations}")
+        else:
+            output_folder = os.path.join(args.output_folder, args.cochlea, f"scale{scale}")
+        os.makedirs(output_folder, exist_ok=True)
+
+        input_key = f"s{scale}"
+        for channel in args.channels:
+            out_path = os.path.join(output_folder, f"{channel}.tif")
+            if os.path.exists(out_path):
+                continue
+
+            print("Exporting channel", channel)
+            internal_path = os.path.join(args.cochlea, "images",  "ome-zarr", f"{channel}.ome.zarr")
+            s3_store, fs = get_s3_path(internal_path, bucket_name=BUCKET_NAME, service_endpoint=SERVICE_ENDPOINT)
+            with zarr.open(s3_store, mode="r") as f:
+                data = f[input_key][:]
+            print("Data shape", data.shape)
+            if args.filter_by_components is not None:
+                data = filter_component(fs, data, args.cochlea, channel, args.filter_by_components)
+            if args.filter_cochlea_channels is not None:
+                us_factor = ds_factor // (2 ** scale)
+                upscaled_filter = upscale_volume(data, filter_volume, upscale_factor=us_factor)
+                data[upscaled_filter == 0] = 0
+                if "PV" in channel:
+                    max_intensity = 1400
+                    data[data > max_intensity] = 0
+
+            if args.binarize:
+                data = (data > 0).astype("uint16")
+            tifffile.imwrite(out_path, data, bigtiff=True, compression="zlib")
 
 
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--cochlea", "-c", required=True)
-    parser.add_argument("--scale", "-s", type=int, required=True)
+    parser.add_argument("--scale", "-s", nargs="+", type=int, required=True)
     parser.add_argument("--output_folder", "-o", required=True)
-    parser.add_argument("--channels", nargs="+", default=["PV", "VGlut3", "CTBP2"])
+    parser.add_argument("--channels", nargs="+", type=str, default=["PV", "VGlut3", "CTBP2"])
     parser.add_argument("--filter_by_components", nargs="+", type=int, default=None)
+    parser.add_argument("--filter_sgn_components", nargs="+", type=int, default=[1])
+    parser.add_argument("--filter_ihc_components", nargs="+", type=int, default=[1])
     parser.add_argument("--binarize", action="store_true")
+    parser.add_argument("--filter_cochlea_channels", nargs="+", type=str, default=None)
+    parser.add_argument("--filter_dilation_iterations", type=int, default=8)
     args = parser.parse_args()
 
     export_lower_resolution(args)
