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    # Read the resolution information from the mrcfile.
243    with mrcfile.open(mrc_path, "r") as f:
244        resolution = f.voxel_size.tolist()
245
246    # The resolution is stored in angstrom, we convert it to nanometer.
247    resolution = [res / 10 for res in resolution]
248
249    # Extract the center coordinates and radii from the segmentation.
250    if isinstance(segmentation, str):
251        segmentation = _load_segmentation(segmentation, segmentation_key)
252    coordinates, radii = convert_segmentation_to_spheres(
253        segmentation, resolution=resolution, radius_factor=radius_factor, estimate_radius_2d=estimate_radius_2d
254    )
255
256    # Write the point annotations to imod.
257    write_points_to_imod(coordinates, radii, segmentation.shape, min_radius, output_path)
258
259
260def _get_file_paths(input_path, ext=(".mrc", ".rec")):
261    if not os.path.exists(input_path):
262        raise Exception(f"Input path not found {input_path}.")
263
264    if isinstance(ext, str):
265        ext = (ext,)
266
267    if os.path.isfile(input_path):
268        input_files = [input_path]
269        input_root = None
270    else:
271        input_files = []
272        for ex in ext:
273            input_files.extend(
274                sorted(glob(os.path.join(input_path, "**", f"*{ex}"), recursive=True))
275            )
276        input_root = input_path
277
278    return input_files, input_root
279
280
281def export_helper(
282    input_path: str,
283    segmentation_path: str,
284    output_root: str,
285    export_function: callable,
286    force: bool = False,
287    segmentation_key: Optional[str] = None,
288) -> None:
289    """
290    Helper function to run imod export for files in a directory.
291
292    Args:
293        input_path: The path to the input data.
294            Can either be a folder. In this case all mrc files below the folder will be exported.
295            Or can be a single mrc file. In this case only this mrc file will be segmented.
296        segmentation_path: Filepath to the segmentation results.
297            The filestructure must exactly match `input_path`.
298        output_root: The path to the output directory where the export results will be saved.
299        export_function: The function performing the export to the desired imod format.
300            This function must take the path to the input in a .mrc file,
301            the path to the segmentation in a .tif file and the output path as only arguments.
302            If you want to pass additional arguments to this function the use 'funtools.partial'
303        force: Whether to rerun segmentation for output files that are already present.
304        segmentation_key: The key to the segmentation data in case the segmentation is stored in hdf5 files.
305    """
306    input_files, input_root = _get_file_paths(input_path)
307    segmentation_files, _ = _get_file_paths(segmentation_path, ext=".tif" if segmentation_key is None else ".h5")
308    assert len(input_files) == len(segmentation_files)
309
310    for input_path, seg_path in tqdm(zip(input_files, segmentation_files), total=len(input_files)):
311        # Determine the output file name.
312        input_folder, input_name = os.path.split(input_path)
313        fname = os.path.splitext(input_name)[0] + ".mod"
314        if input_root is None:
315            output_path = os.path.join(output_root, fname)
316        else:  # If we have nested input folders then we preserve the folder structure in the output.
317            rel_folder = os.path.relpath(input_folder, input_root)
318            output_path = os.path.join(output_root, rel_folder, fname)
319
320        # Check if the output path is already present.
321        # If it is we skip the prediction, unless force was set to true.
322        if os.path.exists(output_path) and not force:
323            continue
324
325        os.makedirs(os.path.split(output_path)[0], exist_ok=True)
326        export_function(input_path, seg_path, output_path, segmentation_key=segmentation_key)