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. #2D or 3D
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    is_3d = coordinates.shape[1] == 3
197    if not is_3d and coordinates.shape[1] != 2:
198        raise ValueError("Coordinates must have shape (N, 2) for 2D or (N, 3) for 3D.")
199
200    with tempfile.NamedTemporaryFile() as tmp_file:
201        fname = tmp_file.name
202        with open(fname, "w") as f:
203            for coord, radius in zip(coordinates, radii):
204                if radius < min_radius:
205                    continue
206
207                if is_3d:
208                    # (z, y, x) indexing
209                    x = _pad(coord[2])
210                    y = _pad(shape[1] - coord[1])
211                    z = _pad(coord[0])
212                    
213                else:
214                    # (y, x) indexing, single z-plane
215                    x = _pad(coord[1])
216                    y = _pad(shape[0] - coord[0])
217                    z=_pad(0)
218                    
219                f.write(f"{x}{y}{z}{_pad(radius, 2)}\n")
220        cmd = [cmd, "-si", "-scat", fname, output_path]
221
222        if color is not None:
223            assert len(color) == 3
224            r, g, b = [str(co) for co in color]
225            cmd += ["-co", f"{r} {g} {b}"]
226
227        run(cmd)
228
229
230
231def write_segmentation_to_imod_as_points(
232    mrc_path: str,
233    segmentation: Union[str, np.ndarray],
234    output_path: str,
235    min_radius: Union[int, float],
236    radius_factor: float = 1.0,
237    estimate_radius_2d: bool = True,
238    segmentation_key: Optional[str] = None,
239) -> None:
240    """Write segmentation results to .mod file with imod point annotations.
241
242    This approximates each segmented object as a sphere.
243
244    Args:
245        mrc_path: Filepath to the mrc volume that was segmented.
246        segmentation: The segmentation (either as numpy array or filepath to a .tif file).
247        output_path: Where to save the .mod file.
248        min_radius: Minimum radius for export.
249        radius_factor: Factor for increasing the radius to account for too small exported spheres.
250        estimate_radius_2d: If true the distance to boundary for determining the centroid and computing
251            the radius will be computed only in 2d rather than in 3d. This can lead to better results
252            in case of deformation across the depth axis.
253        segmentation_key: The key to the segmentation data in case the segmentation is stored in hdf5 files.
254    """
255
256    # Extract the center coordinates and radii from the segmentation.
257    if isinstance(segmentation, str):
258        segmentation = _load_segmentation(segmentation, segmentation_key)
259    coordinates, radii = convert_segmentation_to_spheres(
260        segmentation, radius_factor=radius_factor, estimate_radius_2d=estimate_radius_2d
261    )
262
263    # Write the point annotations to imod.
264    write_points_to_imod(coordinates, radii, segmentation.shape, min_radius, output_path)
265
266
267def _get_file_paths(input_path, ext=(".mrc", ".rec")):
268    if not os.path.exists(input_path):
269        raise Exception(f"Input path not found {input_path}.")
270
271    if isinstance(ext, str):
272        ext = (ext,)
273
274    if os.path.isfile(input_path):
275        input_files = [input_path]
276        input_root = None
277    else:
278        input_files = []
279        for ex in ext:
280            input_files.extend(
281                sorted(glob(os.path.join(input_path, "**", f"*{ex}"), recursive=True))
282            )
283        input_root = input_path
284
285    return input_files, input_root
286
287
288def export_helper(
289    input_path: str,
290    segmentation_path: str,
291    output_root: str,
292    export_function: callable,
293    force: bool = False,
294    segmentation_key: Optional[str] = None,
295) -> None:
296    """
297    Helper function to run imod export for files in a directory.
298
299    Args:
300        input_path: The path to the input data.
301            Can either be a folder. In this case all mrc files below the folder will be exported.
302            Or can be a single mrc file. In this case only this mrc file will be segmented.
303        segmentation_path: Filepath to the segmentation results.
304            The filestructure must exactly match `input_path`.
305        output_root: The path to the output directory where the export results will be saved.
306        export_function: The function performing the export to the desired imod format.
307            This function must take the path to the input in a .mrc file,
308            the path to the segmentation in a .tif file and the output path as only arguments.
309            If you want to pass additional arguments to this function the use 'funtools.partial'
310        force: Whether to rerun segmentation for output files that are already present.
311        segmentation_key: The key to the segmentation data in case the segmentation is stored in hdf5 files.
312    """
313    input_files, input_root = _get_file_paths(input_path)
314    segmentation_files, _ = _get_file_paths(segmentation_path, ext=".tif" if segmentation_key is None else ".h5")
315    assert len(input_files) == len(segmentation_files)
316
317    for input_path, seg_path in tqdm(zip(input_files, segmentation_files), total=len(input_files)):
318        # Determine the output file name.
319        input_folder, input_name = os.path.split(input_path)
320        fname = os.path.splitext(input_name)[0] + ".mod"
321        if input_root is None:
322            output_path = os.path.join(output_root, fname)
323        else:  # If we have nested input folders then we preserve the folder structure in the output.
324            rel_folder = os.path.relpath(input_folder, input_root)
325            output_path = os.path.join(output_root, rel_folder, fname)
326
327        # Check if the output path is already present.
328        # If it is we skip the prediction, unless force was set to true.
329        if os.path.exists(output_path) and not force:
330            continue
331
332        os.makedirs(os.path.split(output_path)[0], exist_ok=True)
333        export_function(input_path, seg_path, output_path, segmentation_key=segmentation_key)