synapse_net.imod.to_imod

  1import multiprocessing
  2import os
  3import shutil
  4import tempfile
  5
  6from concurrent import futures
  7from glob import glob
  8from subprocess import run
  9from typing import Optional, Tuple, Union
 10
 11import h5py
 12import imageio.v3 as imageio
 13import mrcfile
 14import numpy as np
 15from scipy.ndimage import distance_transform_edt
 16from skimage.measure import regionprops
 17from skimage.segmentation import find_boundaries
 18from tqdm import tqdm
 19
 20
 21def _load_segmentation(segmentation_path, segmentation_key):
 22    assert os.path.exists(segmentation_path), segmentation_path
 23    if segmentation_key is None:
 24        seg = imageio.imread(segmentation_path)
 25    else:
 26        with h5py.File(segmentation_path, "r") as f:
 27            seg = f[segmentation_key][:]
 28    return seg
 29
 30
 31def _separate_instances(segmentation):
 32    boundaries = find_boundaries(segmentation, mode="outer")
 33    boundaries &= (segmentation != 0)
 34    segmentation = segmentation.copy()
 35    segmentation[boundaries] = 0
 36    return segmentation
 37
 38
 39# TODO: this has still some issues with some tomograms that has an offset info.
 40# For now, this occurs for the inner ear data tomograms; it works for Fidi's STEM tomograms.
 41# Ben's theory is that this might be due to data form JEOL vs. ThermoFischer microscopes.
 42# To test this I can check how it works for data from Maus et al. / Imig et al., which were taken on a JEOL.
 43# Can also check out the mrc documentation here: https://www.ccpem.ac.uk/mrc_format/mrc2014.php
 44def write_segmentation_to_imod(
 45    mrc_path: str,
 46    segmentation: Union[str, np.ndarray],
 47    output_path: str,
 48    segmentation_key: Optional[str] = None,
 49    separate_instances: bool = True,
 50) -> None:
 51    """Write a segmentation to a mod file as closed contour object(s).
 52
 53    Args:
 54        mrc_path: The filepath to the mrc file from which the segmentation was derived.
 55        segmentation: The segmentation (either as numpy array or filepath to a .tif file).
 56        output_path: The output path where the mod file will be saved.
 57        segmentation_key: The key to the segmentation data in case the segmentation is stored in hdf5 files.
 58        separate_instances: Whether to seperate touching instances before the export.
 59    """
 60    cmd = "imodauto"
 61    cmd_path = shutil.which(cmd)
 62    assert cmd_path is not None, f"Could not find the {cmd} imod command."
 63
 64    # Load the segmentation case a filepath was passed.
 65    if isinstance(segmentation, str):
 66        segmentation = _load_segmentation(segmentation, segmentation_key)
 67
 68    # Binarize the segmentation and flip its axes to match the IMOD axis convention.
 69    if separate_instances:
 70        segmentation = _separate_instances(segmentation)
 71    segmentation = (segmentation > 0).astype("uint8")
 72    segmentation = np.flip(segmentation, axis=1)
 73
 74    # Read the voxel size and origin information from the mrc file.
 75    assert os.path.exists(mrc_path)
 76    with mrcfile.open(mrc_path, mode="r") as f:
 77        voxel_size = f.voxel_size
 78        nx, ny, nz = f.header.nxstart, f.header.nystart, f.header.nzstart
 79        origin = f.header.origin
 80
 81    # Write the input for imodauto to a temporary mrc file.
 82    with tempfile.NamedTemporaryFile(suffix=".mrc") as f:
 83        tmp_path = f.name
 84
 85        mrcfile.new(tmp_path, data=segmentation, overwrite=True)
 86        # Write the voxel_size and origin infomration.
 87        with mrcfile.open(tmp_path, mode="r+") as f:
 88            f.voxel_size = voxel_size
 89
 90            f.header.nxstart = nx
 91            f.header.nystart = ny
 92            f.header.nzstart = nz
 93            f.header.origin = (0.0, 0.0, 0.0) * 3 if origin is None else origin
 94
 95            f.update_header_from_data()
 96
 97        # Run the command.
 98        cmd_list = [cmd, "-E", "1", "-u", tmp_path, output_path]
 99        run(cmd_list)
100
101
102def convert_segmentation_to_spheres(
103    segmentation: np.ndarray,
104    verbose: bool = False,
105    num_workers: Optional[int] = None,
106    resolution: Optional[Tuple[float, float, float]] = None,
107    radius_factor: float = 1.0,
108    estimate_radius_2d: bool = True,
109    props: Optional[list] = None,
110) -> Tuple[np.ndarray, np.ndarray]:
111    """Extract spheres parameterized by center and radius from a segmentation.
112
113    Args:
114        segmentation: The segmentation.
115        verbose: Whether to print a progress bar.
116        num_workers: Number of workers to use for parallelization.
117        resolution: The physical resolution of the data.
118        radius_factor: Factor for increasing the radius to account for too small exported spheres.
119        estimate_radius_2d: If true the distance to boundary for determining the centroid and computing
120            the radius will be computed only in 2d rather than in 3d. This can lead to better results
121            in case of deformation across the depth axis.
122        props: Optional list of regionprops
123
124    Returns:
125        The center coordinates.
126        The radii.
127    """
128    num_workers = multiprocessing.cpu_count() if num_workers is None else num_workers
129    if props is None:
130        props = regionprops(segmentation)
131
132    def coords_and_rads(prop):
133        seg_id = prop.label
134        bbox = prop.bbox
135        ndim = len(bbox) // 2
136
137        # Handle 2D bounding box
138        if len(bbox) == 4:
139            bb = np.s_[bbox[0]:bbox[2], bbox[1]:bbox[3]]
140            mask = segmentation[bb] == seg_id
141            if resolution:
142                dists = distance_transform_edt(mask, sampling=resolution)
143            else:
144                dists = distance_transform_edt(mask)
145
146        # Handle 3D bounding box
147        elif len(bbox) == 6:
148            bb = np.s_[bbox[0]:bbox[3], bbox[1]:bbox[4], bbox[2]:bbox[5]]
149            mask = segmentation[bb] == seg_id
150
151            if estimate_radius_2d:
152                if resolution:
153                    dists = np.array([distance_transform_edt(ma, sampling=resolution[1:]) for ma in mask])
154                else:
155                    dists = np.array([distance_transform_edt(ma) for ma in mask])
156            else:
157                dists = distance_transform_edt(mask, sampling=resolution)
158        else:
159            raise ValueError(f"Unsupported bounding box dimensions: {len(bbox)}")
160
161        max_coord = np.unravel_index(np.argmax(dists), mask.shape)
162        radius = dists[max_coord] * radius_factor
163
164        offset = np.array(bbox[:ndim])
165
166        coord = np.array(max_coord) + offset
167
168        return coord, radius
169
170    with futures.ThreadPoolExecutor(num_workers) as tp:
171        res = list(tqdm(
172            tp.map(coords_and_rads, props), disable=not verbose, total=len(props),
173            desc="Compute coordinates and radii"
174        ))
175
176    coords = [re[0] for re in res]
177    rads = [re[1] for re in res]
178    return np.array(coords), np.array(rads)
179
180
181def write_points_to_imod(
182    coordinates: np.ndarray,
183    radii: np.ndarray,
184    shape: Tuple[int, int, int],
185    min_radius: Union[float, int],
186    output_path: str,
187    color: Optional[Tuple[int, int, int]] = None,
188) -> None:
189    """Write point annotations to a .mod file for IMOD.
190
191    Args:
192        coordinates: Array with the point coordinates. #2D or 3D
193        radii: Array with the point radii.
194        shape: Shape of the volume to be exported.
195        min_radius: Minimum radius for export.
196        output_path: Where to save the .mod file.
197        color: Optional color for writing out the points.
198    """
199    cmd = "point2model"
200    cmd_path = shutil.which(cmd)
201    assert cmd_path is not None, f"Could not find the {cmd} imod command."
202
203    def _pad(inp, n=3):
204        inp = int(round(inp))
205        plen = 6 + (n - len(str(inp)))
206        pw = plen * " "
207        return f"{pw}{inp}.00"
208
209    is_3d = coordinates.shape[1] == 3
210    if not is_3d and coordinates.shape[1] != 2:
211        raise ValueError("Coordinates must have shape (N, 2) for 2D or (N, 3) for 3D.")
212
213    with tempfile.NamedTemporaryFile() as tmp_file:
214        fname = tmp_file.name
215        with open(fname, "w") as f:
216            for coord, radius in zip(coordinates, radii):
217                if radius < min_radius:
218                    continue
219
220                if is_3d:
221                    # (z, y, x) indexing
222                    x = _pad(coord[2])
223                    y = _pad(shape[1] - coord[1])
224                    z = _pad(coord[0])
225
226                else:
227                    # (y, x) indexing, single z-plane
228                    x = _pad(coord[1])
229                    y = _pad(shape[0] - coord[0])
230                    z = _pad(0)
231
232                f.write(f"{x}{y}{z}{_pad(radius, 2)}\n")
233        cmd = [cmd, "-si", "-scat", fname, output_path]
234
235        if color is not None:
236            assert len(color) == 3
237            r, g, b = [str(co) for co in color]
238            cmd += ["-co", f"{r} {g} {b}"]
239
240        run(cmd)
241
242
243def write_segmentation_to_imod_as_points(
244    mrc_path: str,
245    segmentation: Union[str, np.ndarray],
246    output_path: str,
247    min_radius: Union[int, float],
248    radius_factor: float = 1.0,
249    estimate_radius_2d: bool = True,
250    segmentation_key: Optional[str] = None,
251) -> None:
252    """Write segmentation results to .mod file with imod point annotations.
253
254    This approximates each segmented object as a sphere.
255
256    Args:
257        mrc_path: Filepath to the mrc volume that was segmented.
258        segmentation: The segmentation (either as numpy array or filepath to a .tif file).
259        output_path: Where to save the .mod file.
260        min_radius: Minimum radius for export.
261        radius_factor: Factor for increasing the radius to account for too small exported spheres.
262        estimate_radius_2d: If true the distance to boundary for determining the centroid and computing
263            the radius will be computed only in 2d rather than in 3d. This can lead to better results
264            in case of deformation across the depth axis.
265        segmentation_key: The key to the segmentation data in case the segmentation is stored in hdf5 files.
266    """
267
268    # Extract the center coordinates and radii from the segmentation.
269    if isinstance(segmentation, str):
270        segmentation = _load_segmentation(segmentation, segmentation_key)
271    coordinates, radii = convert_segmentation_to_spheres(
272        segmentation, radius_factor=radius_factor, estimate_radius_2d=estimate_radius_2d
273    )
274
275    # Write the point annotations to imod.
276    write_points_to_imod(coordinates, radii, segmentation.shape, min_radius, output_path)
277
278
279def _get_file_paths(input_path, ext=(".mrc", ".rec")):
280    if not os.path.exists(input_path):
281        raise Exception(f"Input path not found {input_path}.")
282
283    if isinstance(ext, str):
284        ext = (ext,)
285
286    if os.path.isfile(input_path):
287        input_files = [input_path]
288        input_root = None
289    else:
290        input_files = []
291        for ex in ext:
292            input_files.extend(
293                sorted(glob(os.path.join(input_path, "**", f"*{ex}"), recursive=True))
294            )
295        input_root = input_path
296
297    return input_files, input_root
298
299
300def export_helper(
301    input_path: str,
302    segmentation_path: str,
303    output_root: str,
304    export_function: callable,
305    force: bool = False,
306    segmentation_key: Optional[str] = None,
307) -> None:
308    """
309    Helper function to run imod export for files in a directory.
310
311    Args:
312        input_path: The path to the input data.
313            Can either be a folder. In this case all mrc files below the folder will be exported.
314            Or can be a single mrc file. In this case only this mrc file will be segmented.
315        segmentation_path: Filepath to the segmentation results.
316            The filestructure must exactly match `input_path`.
317        output_root: The path to the output directory where the export results will be saved.
318        export_function: The function performing the export to the desired imod format.
319            This function must take the path to the input in a .mrc file,
320            the path to the segmentation in a .tif file and the output path as only arguments.
321            If you want to pass additional arguments to this function the use 'funtools.partial'
322        force: Whether to rerun segmentation for output files that are already present.
323        segmentation_key: The key to the segmentation data in case the segmentation is stored in hdf5 files.
324    """
325    input_files, input_root = _get_file_paths(input_path)
326    segmentation_files, _ = _get_file_paths(segmentation_path, ext=".tif" if segmentation_key is None else ".h5")
327    assert len(input_files) == len(segmentation_files)
328
329    for input_path, seg_path in tqdm(zip(input_files, segmentation_files), total=len(input_files)):
330        # Determine the output file name.
331        input_folder, input_name = os.path.split(input_path)
332        fname = os.path.splitext(input_name)[0] + ".mod"
333        if input_root is None:
334            output_path = os.path.join(output_root, fname)
335        else:  # If we have nested input folders then we preserve the folder structure in the output.
336            rel_folder = os.path.relpath(input_folder, input_root)
337            output_path = os.path.join(output_root, rel_folder, fname)
338
339        # Check if the output path is already present.
340        # If it is we skip the prediction, unless force was set to true.
341        if os.path.exists(output_path) and not force:
342            continue
343
344        os.makedirs(os.path.split(output_path)[0], exist_ok=True)
345        export_function(input_path, seg_path, output_path, segmentation_key=segmentation_key)