synapse_net.imod.export

  1import shutil
  2import tempfile
  3from subprocess import run
  4from typing import Dict, List, Optional, Tuple
  5
  6import imageio.v3 as imageio
  7import numpy as np
  8from elf.io import open_file
  9from scipy.ndimage import maximum_filter1d
 10from skimage.morphology import ball
 11
 12from tqdm import tqdm
 13
 14
 15def get_label_names(
 16    imod_path: str,
 17    return_types: bool = False,
 18    return_counts: bool = False,
 19    filter_empty_labels: bool = True,
 20) -> Dict[str, str]:
 21    """Get the ids and names of annotations in an imod file.
 22
 23    Uses the imodinfo command to parse the contents of the imod file.
 24
 25    Args:
 26        imod_path: The filepath to the imod annotation file.
 27        return_types: Whether to also return the annotation types.
 28        return_counts: Whether to also return the number of annotations per object.
 29        filter_empty_labels: Whether to filter out empty objects from the object list.
 30
 31    Returns:
 32        Dictionary that maps the object ids to their names.
 33    """
 34    cmd = "imodinfo"
 35    cmd_path = shutil.which(cmd)
 36    assert cmd_path is not None, f"Could not find the {cmd} imod command."
 37
 38    label_names, label_types, label_counts = {}, {}, {}
 39
 40    with tempfile.NamedTemporaryFile() as f:
 41        tmp_path = f.name
 42        run([cmd, "-f", tmp_path, imod_path])
 43
 44        object_id = None
 45        with open(tmp_path) as f:
 46            lines = f.readlines()
 47            for line in lines:
 48
 49                if line.startswith("OBJECT"):
 50                    object_id = int(line.rstrip("\n").strip().split()[-1])
 51
 52                if line.startswith("NAME"):
 53                    name = line.rstrip("\n").strip().split()[-1]
 54                    assert object_id is not None
 55                    label_names[object_id] = name
 56
 57                if "object uses" in line:
 58                    type_ = " ".join(line.rstrip("\n").strip().split()[2:]).rstrip(".")
 59                    label_types[object_id] = type_
 60
 61                # Parse the number of annotions for this object.
 62                if "contours" in line:
 63                    try:
 64                        count = int(line.rstrip("\n").strip().split()[0])
 65                        label_counts[object_id] = count
 66                    # The condition also hits for 'closed countours' / 'open contours'.
 67                    # We just catch the exception thrown for these cases and continue.
 68                    except ValueError:
 69                        pass
 70
 71                if "CONTOUR" in line and label_types[object_id] == "scattered points":
 72                    count = int(line.rstrip("\n").strip().split()[2])
 73                    label_counts[object_id] = count
 74
 75    if filter_empty_labels:
 76        empty_labels = [k for k, v in label_counts.items() if v == 0]
 77        label_names = {k: v for k, v in label_names.items() if k not in empty_labels}
 78        label_types = {k: v for k, v in label_types.items() if k not in empty_labels}
 79        label_counts = {k: v for k, v in label_counts.items() if k not in empty_labels}
 80
 81    if return_types and return_counts:
 82        return label_names, label_types, label_counts
 83    elif return_types:
 84        return label_names, label_types
 85    elif return_counts:
 86        return label_names, label_counts
 87
 88    return label_names
 89
 90
 91def export_segmentation(
 92    imod_path: str,
 93    mrc_path: str,
 94    object_id: Optional[int] = None,
 95    output_path: Optional[int] = None,
 96    require_object: bool = True,
 97    depth_interpolation: Optional[int] = None,
 98) -> Optional[np.ndarray]:
 99    """Export a segmentation mask from IMOD annotations.
100
101    This function uses the imodmop command to extract masks.
102
103    Args:
104        imod_path: The filepath to the imod annotation file.
105        mrc_path: The filepath to the mrc file with the corresponding tomogram data.
106        object_id: The id of the object to be extracted.
107            If none is given all objects in the mod file will be extracted.
108        output_path: Optional path to a file where to save the extracted mask as tif.
109            If not given the segmentation mask will be returned as a numpy array.
110        require_object: Whether to throw an error if the object could not be extractd.
111        depth_interpolation: Window for interpolating the extracted mask across the depth
112            axis. This can be used to close masks which are only annotated in a subset of slices.
113
114    Returns:
115        The extracted mask. Will only be returned if `output_path` is not given.
116    """
117    cmd = "imodmop"
118    cmd_path = shutil.which(cmd)
119    assert cmd_path is not None, f"Could not find the {cmd} imod command."
120
121    with tempfile.NamedTemporaryFile() as f:
122        tmp_path = f.name
123
124        if object_id is None:
125            cmd = [cmd, "-ma", "1", imod_path, mrc_path, tmp_path]
126        else:
127            cmd = [cmd, "-ma", "1", "-o", str(object_id), imod_path, mrc_path, tmp_path]
128
129        run(cmd)
130        with open_file(tmp_path, ext=".mrc", mode="r") as f:
131            data = f["data"][:]
132
133    segmentation = data == 1
134    if require_object and segmentation.sum() == 0:
135        id_str = "" if object_id is None else f"for object {object_id}"
136        raise RuntimeError(f"Segmentation extracted from {imod_path} {id_str} is empty.")
137
138    if depth_interpolation is not None:
139        segmentation = maximum_filter1d(segmentation, size=depth_interpolation, axis=0, mode="constant")
140
141    if output_path is None:
142        return segmentation
143
144    imageio.imwrite(output_path, segmentation.astype("uint8"), compression="zlib")
145
146
147def draw_spheres(
148    coordinates: np.ndarray,
149    radii: np.ndarray,
150    shape: Tuple[int, int, int],
151    verbose: bool = True,
152) -> np.ndarray:
153    """Create a volumetric segmentation by painting spheres around the given coordinates.
154
155    Args:
156        coordinates: The center coordinates of the spheres.
157        radii: The radii of the spheres.
158        shape: The shape of the volume.
159        verbose: Whether to print the progress bar.
160
161    Returns:
162        The segmentation volume with painted spheres.
163    """
164    labels = np.zeros(shape, dtype="uint32")
165    for label_id, (coord, radius) in tqdm(
166        enumerate(zip(coordinates, radii), start=1), total=len(coordinates), disable=not verbose
167    ):
168        radius = int(radius)
169        mask = ball(radius)
170        full_mask = np.zeros(shape, dtype="bool")
171        full_slice = tuple(
172            slice(max(co - radius, 0), min(co + radius, sh)) for co, sh in zip(coord, shape)
173        )
174        radius_clipped_left = [co - max(co - radius, 0) for co in coord]
175        radius_clipped_right = [min(co + radius, sh) - co for co, sh in zip(coord, shape)]
176        mask_slice = tuple(
177            slice(radius + 1 - rl, radius + 1 + rr) for rl, rr in zip(radius_clipped_left, radius_clipped_right)
178        )
179        full_mask[full_slice] = mask[mask_slice]
180        labels[full_mask] = label_id
181    return labels
182
183
184def load_points_from_imodinfo(
185    imod_path: str,
186    full_shape: Tuple[int, int, int],
187    bb: Optional[Tuple[slice, slice, slice]] = None,
188    exclude_labels: Optional[List[int]] = None,
189    exclude_label_patterns: Optional[List[str]] = None,
190    resolution: Optional[float] = None,
191) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Dict[int, str]]:
192    """Load point coordinates, radii and label information from a .mod file.
193
194    The coordinates and sizes returned will be scaled so that they are in
195    the voxel coordinate space if the 'resolution' parameter is passed.
196    If it is not passed then the radius will be returned in the physical resolution.
197
198    Args:
199        imod_path: The filepath to the .mod file.
200        full_shape: The voxel shape of the volume.
201        bb: Optional bounding box to limit the extracted points to.
202        exclude_labels: Label ids to exclude from the export.
203            This can be used to exclude specific labels / classes, specifying them by their id.
204        exclude_label_patterns: Label names to exclude from the export.
205            This can be used to exclude specific labels / classes, specifying them by their name.
206        resolution: The resolution / voxel size of the data. Will be used to scale the radii.
207
208    Returns:
209        The center coordinates of the sphere annotations.
210        The radii of the spheres.
211        The ids of the semantic labels.
212        The names of the semantic labels.
213    """
214    coordinates, sizes, labels = [], [], []
215    label_names = {}
216
217    if bb is not None:
218        start = [b.start for b in bb]
219        shape = [b.stop - sta for b, sta in zip(bb, start)]
220
221    # first round: load the point sizes and labels
222    with tempfile.NamedTemporaryFile() as f:
223        tmp_path = f.name
224        run(["imodinfo", "-p", "-f", tmp_path, imod_path])
225
226        label_id = None
227        label_name = None
228        with open(tmp_path) as f:
229            lines = f.readlines()
230            for line in lines:
231
232                if line.startswith("#Object"):
233                    line = line.rstrip("\n")
234                    label_id = int(line.split(" ")[1])
235                    label_name = line.split(",")[1]
236                    label_names[label_id] = label_name
237                    continue
238
239                if label_id is None:
240                    continue
241                if exclude_labels is not None and label_id in exclude_labels:
242                    continue
243                if exclude_label_patterns is not None and any(
244                    pattern.lower() in label_name.lower() for pattern in exclude_label_patterns
245                ):
246                    continue
247
248                try:
249                    size = float(line.rstrip("\n"))
250                    sizes.append(size)
251                    labels.append(label_id)
252                except ValueError:
253                    continue
254
255    label_names = {lid: label_name for lid, label_name in label_names.items() if lid in labels}
256
257    in_bounds = []
258    with tempfile.NamedTemporaryFile() as f:
259        tmp_path = f.name
260        run(["imodinfo", "-vv", "-f", tmp_path, imod_path])
261
262        label_id = None
263        with open(tmp_path) as f:
264            lines = f.readlines()
265            for line in lines:
266                if line.startswith("OBJECT"):
267                    label_id = int(line.rstrip("\n").split(" ")[-1])
268                    continue
269                if label_id not in label_names:
270                    continue
271
272                values = line.strip().rstrip("\n").split("\t")
273                if len(values) != 3:
274                    continue
275
276                try:
277                    x, y, z = float(values[0]), float(values[1]), float(values[2])
278                    coords = [x, y, z]
279                    coords = [int(np.round(float(coord))) for coord in coords]
280                    # IMOD uses a very weird coordinate system:
281                    coords = [coords[2], full_shape[1] - coords[1], coords[0]]
282
283                    in_bound = True
284                    if bb is not None:
285                        coords = [co - sta for co, sta in zip(coords, start)]
286                        if any(co < 0 or co > sh for co, sh in zip(coords, shape)):
287                            in_bound = False
288
289                    in_bounds.append(in_bound)
290                    coordinates.append(coords)
291                except ValueError:
292                    continue
293
294    # labelx = {seg_id: int(label_id) for seg_id, label_id in enumerate(labels, 1)}
295    # print("Extracted the following labels:", label_names)
296    # print("With counts:", {k: v for k, v in zip(*np.unique(list(labelx.values()), return_counts=True))})
297    # breakpoint()
298    assert len(coordinates) == len(sizes) == len(labels) == len(in_bounds), \
299        f"{len(coordinates)}, {len(sizes)}, {len(labels)}, {len(in_bounds)}"
300
301    coordinates, sizes, labels = np.array(coordinates), np.array(sizes), np.array(labels)
302    in_bounds = np.array(in_bounds)
303
304    # get rid of empty annotations
305    in_bounds = np.logical_and(in_bounds, sizes > 0)
306
307    coordinates, sizes, labels = coordinates[in_bounds], sizes[in_bounds], labels[in_bounds]
308
309    if resolution is not None:
310        sizes /= resolution
311
312    if len(coordinates) == 0:
313        raise RuntimeError(f"Empty annotations: {imod_path}")
314    return coordinates, sizes, labels, label_names
315
316
317def export_point_annotations(
318    imod_path: str,
319    shape: Tuple[int, int, int],
320    bb: Optional[Tuple[slice, slice, slice]] = None,
321    exclude_labels: Optional[List[int]] = None,
322    exclude_label_patterns: Optional[List[str]] = None,
323    return_coords_and_radii: bool = False,
324    resolution: Optional[float] = None,
325) -> Tuple[np.ndarray, np.ndarray, Dict[int, str]]:
326    """Create a segmentation by drawing spheres corresponding to objects from a .mod file.
327
328    Args:
329        imod_path: The filepath to the .mod file.
330        shape: The voxel shape of the volume.
331        bb: Optional bounding box to limit the extracted points to.
332        exclude_labels: Label ids to exclude from the segmentation.
333            This can be used to exclude specific labels / classes, specifying them by their id.
334        exclude_label_patterns: Label names to exclude from the segmentation.
335            This can be used to exclude specific labels / classes, specifying them by their name.
336        return_coords_and_radii: Whether to also return the underlying coordinates
337            and radii of the exported spheres.
338        resolution: The resolution / voxel size of the data. Will be used to scale the radii.
339
340    Returns:
341        The exported segmentation.
342        The label ids for the instance ids in the segmentation.
343        The map of label ids to corresponding obejct names.
344    """
345    coordinates, radii, labels, label_names = load_points_from_imodinfo(
346        imod_path, shape, bb=bb,
347        exclude_labels=exclude_labels,
348        exclude_label_patterns=exclude_label_patterns,
349        resolution=resolution,
350    )
351    labels = {seg_id: int(label_id) for seg_id, label_id in enumerate(labels, 1)}
352    segmentation = draw_spheres(coordinates, radii, shape)
353    if return_coords_and_radii:
354        return segmentation, labels, label_names, coordinates, radii
355    return segmentation, labels, label_names