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