synapse_net.inference

184def compute_scale_from_voxel_size(
185    voxel_size: Dict[str, float],
186    model_type: str
187) -> List[float]:
188    """Compute the appropriate scale factor for inference with a given pretrained model.
189
190    Args:
191        voxel_size: The voxel size of the data for inference.
192        model_type: The name of the pretrained model.
193
194    Returns:
195        The scale factor, as a list in zyx order.
196    """
197    training_voxel_size = get_model_training_resolution(model_type)
198    scale = [
199        voxel_size["x"] / training_voxel_size["x"],
200        voxel_size["y"] / training_voxel_size["y"],
201    ]
202    if len(voxel_size) == 3 and len(training_voxel_size) == 3:
203        scale.append(
204            voxel_size["z"] / training_voxel_size["z"]
205        )
206    return scale

 97def get_model(model_type: str, device: Optional[Union[str, torch.device]] = None) -> torch.nn.Module:
 98    """Get the model for a specific segmentation type.
 99
100    Args:
101        model_type: The model for one of the following segmentation tasks:
102            'vesicles_3d', 'active_zone', 'compartments', 'mitochondria', 'ribbon', 'vesicles_2d', 'vesicles_cryo'.
103        device: The device to use.
104
105    Returns:
106        The model.
107    """
108    if device is None:
109        device = get_device(device)
110    model_path = get_model_path(model_type)
111    model = _restore_legacy_attributes(torch.load(model_path, weights_only=False))
112    model.to(device)
113    return model

276def get_segmentation_function(model_type: str) -> Callable:
277    """Get the segmentation function associated with a model type.
278
279    Args:
280        model_type: The name of the pretrained model.
281
282    Returns:
283        The segmentation function used for the model type.
284
285    Raises:
286        ValueError: If the model type is unknown.
287    """
288    if model_type.startswith("vesicles"):
289        return segment_vesicles
290    if model_type in ("mitochondria", "mitochondria2"):
291        return segment_mitochondria
292    if model_type == "active_zone":
293        return segment_active_zone
294    if model_type == "compartments":
295        return segment_compartments
296    if model_type == "ribbon":
297        return _segment_ribbon_AZ
298    if "cristae" in model_type:
299        return segment_cristae
300    raise ValueError(f"Unknown model type: {model_type}")

303def run_segmentation(
304    image: np.ndarray,
305    model: torch.nn.Module,
306    model_type: str,
307    tiling: Optional[Dict[str, Dict[str, int]]] = None,
308    scale: Optional[List[float]] = None,
309    verbose: bool = False,
310    **kwargs,
311) -> np.ndarray | Dict[str, np.ndarray]:
312    """Run synaptic structure segmentation.
313
314    Args:
315        image: The input image or image volume.
316        model: The segmentation model.
317        model_type: The model type. This will determine which segmentation post-processing is used.
318        tiling: The tiling settings for inference. By default, the tiling is derived from the patch shape
319            that was used for training the model.
320        scale: A scale factor for resizing the input before applying the model.
321            The output will be scaled back to the initial size.
322        verbose: Whether to print detailed information about the prediction and segmentation.
323        kwargs: Optional parameters for the segmentation function.
324
325    Returns:
326        The segmentation. For models that return multiple segmentations, this function returns a dictionary.
327    """
328    if tiling is None:
329        tiling = get_default_tiling(patch_shape=get_model_training_patch_shape(model_type))
330    segmentation_function = get_segmentation_function(model_type)
331    if segmentation_function is segment_cristae:
332        training_resolution = get_model_training_resolution(model_type)
333        voxel_size = np.mean(list(training_resolution.values()))
334        segmentation = segmentation_function(
335            image, model=model, tiling=tiling, scale=scale, verbose=verbose, voxel_size=voxel_size, **kwargs
336        )
337    else:
338        segmentation = segmentation_function(
339            image, model=model, tiling=tiling, scale=scale, verbose=verbose, **kwargs
340        )
341    return segmentation