synapse_net.inference

149def compute_scale_from_voxel_size(
150    voxel_size: Dict[str, float],
151    model_type: str
152) -> List[float]:
153    """Compute the appropriate scale factor for inference with a given pretrained model.
154
155    Args:
156        voxel_size: The voxel size of the data for inference.
157        model_type: The name of the pretrained model.
158
159    Returns:
160        The scale factor, as a list in zyx order.
161    """
162    training_voxel_size = get_model_training_resolution(model_type)
163    scale = [
164        voxel_size["x"] / training_voxel_size["x"],
165        voxel_size["y"] / training_voxel_size["y"],
166    ]
167    if len(voxel_size) == 3 and len(training_voxel_size) == 3:
168        scale.append(
169            voxel_size["z"] / training_voxel_size["z"]
170        )
171    return scale

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

241def get_segmentation_function(model_type: str) -> Callable:
242    """Get the segmentation function associated with a model type.
243
244    Args:
245        model_type: The name of the pretrained model.
246
247    Returns:
248        The segmentation function used for the model type.
249
250    Raises:
251        ValueError: If the model type is unknown.
252    """
253    if model_type.startswith("vesicles"):
254        return segment_vesicles
255    if model_type in ("mitochondria", "mitochondria2"):
256        return segment_mitochondria
257    if model_type == "active_zone":
258        return segment_active_zone
259    if model_type == "compartments":
260        return segment_compartments
261    if model_type == "ribbon":
262        return _segment_ribbon_AZ
263    if "cristae" in model_type:
264        return segment_cristae
265    raise ValueError(f"Unknown model type: {model_type}")

268def run_segmentation(
269    image: np.ndarray,
270    model: torch.nn.Module,
271    model_type: str,
272    tiling: Optional[Dict[str, Dict[str, int]]] = None,
273    scale: Optional[List[float]] = None,
274    verbose: bool = False,
275    **kwargs,
276) -> np.ndarray | Dict[str, np.ndarray]:
277    """Run synaptic structure segmentation.
278
279    Args:
280        image: The input image or image volume.
281        model: The segmentation model.
282        model_type: The model type. This will determine which segmentation post-processing is used.
283        tiling: The tiling settings for inference.
284        scale: A scale factor for resizing the input before applying the model.
285            The output will be scaled back to the initial size.
286        verbose: Whether to print detailed information about the prediction and segmentation.
287        kwargs: Optional parameters for the segmentation function.
288
289    Returns:
290        The segmentation. For models that return multiple segmentations, this function returns a dictionary.
291    """
292    segmentation_function = get_segmentation_function(model_type)
293    if segmentation_function is segment_cristae:
294        training_resolution = get_model_training_resolution(model_type)
295        voxel_size = np.mean(list(training_resolution.values()))
296        segmentation = segmentation_function(
297            image, model=model, tiling=tiling, scale=scale, verbose=verbose, voxel_size=voxel_size, **kwargs
298        )
299    else:
300        segmentation = segmentation_function(
301            image, model=model, tiling=tiling, scale=scale, verbose=verbose, **kwargs
302        )
303    return segmentation