synapse_net.morphology

  1import multiprocessing as mp
  2import warnings
  3from concurrent import futures
  4from typing import Dict, List, Optional, Tuple
  5
  6import trimesh
  7
  8import numpy as np
  9import pandas as pd
 10
 11from scipy.ndimage import distance_transform_edt, convolve
 12from skimage.graph import MCP
 13from skimage.measure import regionprops, marching_cubes, find_contours
 14from skimage.morphology import skeletonize, medial_axis, label
 15from skimage.segmentation import find_boundaries
 16
 17
 18def _size_filter_ids(ids, props, min_size):
 19    if ids is None:
 20        ids = [prop.label for prop in props if prop.area > min_size]
 21    else:
 22        sizes = {prop.label: prop.area for prop in props}
 23        ids = [vid for vid in ids if sizes[vid] > min_size]
 24    return ids
 25
 26
 27def _compute_radii_naive(vesicles, resolution, ids, min_size):
 28    props = regionprops(vesicles)
 29
 30    ids = _size_filter_ids(ids, props, min_size)
 31    radii = {prop.label: resolution[0] * (prop.axis_minor_length + prop.axis_major_length) / 4
 32             for prop in props if prop.label in ids}
 33    assert len(radii) == len(ids)
 34    return ids, radii
 35
 36
 37def _compute_radii_distances(vesicles, resolution, ids, min_size, derive_distances_2d):
 38    vesicle_boundaries = find_boundaries(vesicles, mode="thick")
 39
 40    if derive_distances_2d:
 41
 42        def dist(input_):
 43            return distance_transform_edt(input_ == 0, sampling=resolution[1:])
 44
 45        with futures.ThreadPoolExecutor(mp.cpu_count()) as tp:
 46            distances = list(tp.map(dist, vesicle_boundaries))
 47        distances = np.stack(distances)
 48
 49    else:
 50        distances = distance_transform_edt(vesicle_boundaries == 0, sampling=resolution)
 51
 52    props = regionprops(vesicles, intensity_image=distances)
 53    ids = _size_filter_ids(ids, props, min_size)
 54    radii = {prop.label: prop.intensity_max for prop in props if prop.label in ids}
 55    assert len(radii) == len(ids)
 56
 57    return ids, radii
 58
 59
 60def compute_radii(
 61    vesicles: np.ndarray,
 62    resolution: Tuple[float, float, float],
 63    ids: Optional[List[int]] = None,
 64    derive_radius_from_distances: bool = True,
 65    derive_distances_2d: bool = True,
 66    min_size: int = 500,
 67) -> Tuple[List[int], Dict[int, float]]:
 68    """Compute the radii for a vesicle segmentation.
 69
 70    Args:
 71        vesicles: The vesicle segmentation.
 72        resolution: The pixel / voxel size of the data.
 73        ids: Vesicle ids to restrict the radius computation to.
 74        derive_radius_from_distances: Whether to derive the radii
 75            from the distance to the vesicle boundaries, or from the
 76            axis fitted to the vesicle by scikit-image regionprops.
 77        derive_distances_2d: Whether to derive the radii individually in 2d
 78            or in 3d. Deriving the radii in 3d is beneficial for anisotropic
 79            data or data that suffers from the missing wedge effect.
 80        min_size: The minimal size for extracting the radii.
 81
 82    Returns:
 83        The ids of the extracted radii.
 84        The radii that were computed.
 85    """
 86    if derive_radius_from_distances:
 87        ids, radii = _compute_radii_distances(
 88            vesicles, resolution,
 89            ids=ids, min_size=min_size, derive_distances_2d=derive_distances_2d
 90        )
 91    else:
 92        assert not derive_distances_2d
 93        ids, radii = _compute_radii_naive(vesicles, resolution, ids=ids, min_size=min_size)
 94    return ids, radii
 95
 96
 97def compute_object_morphology(
 98    object_: np.ndarray,
 99    structure_name: str,
100    resolution: Tuple[float, float, float] = None
101) -> pd.DataFrame:
102    """Compute the volume and surface area of a 2D or 3D object.
103
104    Args:
105        object_: 2D or 3D binary object array.
106        structure_name: Name of the structure being analyzed.
107        resolution: The pixel / voxel size of the data.
108
109    Returns:
110        Morphology information containing volume and surface area.
111    """
112    if object_.ndim == 2:
113        # Use find_contours for 2D data
114        contours = find_contours(object_, level=0.5)
115
116        # Compute perimeter (total length of all contours)
117        perimeter = sum(
118            np.sqrt(np.sum(np.diff(contour, axis=0)**2, axis=1)).sum()
119            for contour in contours
120        )
121
122        # Compute area (number of positive pixels)
123        area = np.sum(object_ > 0)
124
125        # Adjust for resolution if provided
126        if resolution is not None:
127            area *= resolution[0] * resolution[1]
128            perimeter *= resolution[0]
129
130        morphology = pd.DataFrame({
131            "structure": [structure_name],
132            "area [pixel^2]" if resolution is None else "area [nm^2]": [area],
133            "perimeter [pixel]" if resolution is None else "perimeter [nm]": [perimeter],
134        })
135
136    elif object_.ndim == 3:
137        # Use marching_cubes for 3D data
138        verts, faces, normals, _ = marching_cubes(
139            object_,
140            spacing=(1.0, 1.0, 1.0) if resolution is None else resolution,
141        )
142
143        mesh = trimesh.Trimesh(vertices=verts, faces=faces, vertex_normals=normals)
144        surface = mesh.area
145        if mesh.is_watertight:
146            volume = np.abs(mesh.volume)
147        else:
148            warnings.warn("Could not compute mesh surface for the volume; setting it to NaN.")
149            volume = np.nan
150
151        morphology = pd.DataFrame({
152            "structure": [structure_name],
153            "volume [pixel^3]" if resolution is None else "volume [nm^3]": [volume],
154            "surface [pixel^2]" if resolution is None else "surface [nm^2]": [surface],
155        })
156
157    else:
158        raise ValueError("Input object must be a 2D or 3D numpy array.")
159
160    return morphology
161
162
163def _find_endpoints(component):
164    # Define a 3x3 kernel to count neighbors
165    kernel = np.ones((3, 3), dtype=int)
166    neighbor_count = convolve(component.astype(int), kernel, mode="constant", cval=0)
167    endpoints = np.argwhere((component == 1) & (neighbor_count == 2))  # Degree = 1
168    return endpoints
169
170
171def _compute_longest_path(component, endpoints):
172    # Use the first endpoint as the source
173    src = tuple(endpoints[0])
174    cost = np.where(component, 1, np.inf)  # Cost map: 1 for skeleton, inf for background
175    mcp = MCP(cost)
176    _, traceback = mcp.find_costs([src])
177
178    # Use the second endpoint as the destination
179    dst = tuple(endpoints[-1])
180
181    # Trace back the path
182    path = np.zeros_like(component, dtype=bool)
183    current = dst
184
185    # Extract offsets from the MCP object
186    offsets = np.array(mcp.offsets)
187    nrows, ncols = component.shape
188
189    while current != src:
190        path[current] = True
191        current_offset_index = traceback[current]
192        if current_offset_index < 0:
193            # No valid path found
194            break
195        offset = offsets[current_offset_index]
196        # Move to the predecessor
197        current = (current[0] - offset[0], current[1] - offset[1])
198        # Ensure indices are within bounds
199        if not (0 <= current[0] < nrows and 0 <= current[1] < ncols):
200            break
201
202    path[src] = True  # Include the source
203    return path
204
205
206def _prune_skeleton_longest_path(skeleton):
207    pruned_skeleton = np.zeros_like(skeleton, dtype=bool)
208
209    # Label connected components in the skeleton
210    labeled_skeleton, num_labels = label(skeleton, return_num=True)
211
212    for label_id in range(1, num_labels + 1):
213        # Isolate the current connected component
214        component = (labeled_skeleton == label_id)
215
216        # Find the endpoints of the component
217        endpoints = _find_endpoints(component)
218        if len(endpoints) < 2:
219            continue  # Skip if there are no valid endpoints
220        elif len(endpoints) == 2:  # Nothing to prune
221            pruned_skeleton |= component
222            continue
223
224        # Compute the longest path using MCP
225        longest_path = _compute_longest_path(component, endpoints)
226        pruned_skeleton |= longest_path
227
228    return pruned_skeleton.astype(skeleton.dtype)
229
230
231def skeletonize_object(
232    segmentation: np.ndarray,
233    method: str = "skeletonize",
234    prune: bool = True,
235    min_prune_size: int = 10,
236) -> np.ndarray:
237    """Skeletonize a 3D object by inidividually skeletonizing each slice.
238
239    Args:
240        segmentation: The segmented object.
241        method: The method to use for skeletonization. Either 'skeletonize' or 'medial_axis'.
242        prune: Whether to prune the skeleton.
243        min_prune_size: The minimal size of components after pruning.
244
245    Returns:
246        The skeletonized object.
247    """
248    assert method in ("skeletonize", "medial_axis")
249    seg_thin = np.zeros_like(segmentation)
250    skeletor = skeletonize if method == "skeletonize" else medial_axis
251    # Parallelize?
252    for z in range(segmentation.shape[0]):
253        skeleton = skeletor(segmentation[z])
254
255        if prune:
256            skeleton = _prune_skeleton_longest_path(skeleton)
257            if min_prune_size > 0:
258                skeleton = label(skeleton)
259                ids, sizes = np.unique(skeleton, return_counts=True)
260                ids, sizes = ids[1:], sizes[1:]
261                skeleton = np.isin(skeleton, ids[sizes >= min_prune_size])
262
263        seg_thin[z] = skeleton
264    return seg_thin