30686229961

Committed 30 Jul 2026 04:01PM UTC coverage: 78.73% (-0.7%) from 79.392%

Build # 30686229961

Build Type

push

github

Committed by

web-flow

Commit Message

Update torch (and nnunetv2) + add Apple MPS support (#996)

* Bump nnunetv2, remove restriction on torch

* Update new argument of nnUNetPredictor for gpu use

* Remove upper restriction on python version

* Fix tqdm_wrapper in napari plugin for progress bar

* Use MPS backend when available

* Use canonical path to pytorch MPS backend

* Fix napari progress bar update

* Fix napari cancel button

* Update Python version in conda environment creation

* Refactor installation script for ADS

* Avoid warnings with MPS backend

* Update Python version in conda environment creation

* Refactor ADS installation script for clarity

* Update Python version in CI workflow to 3.13

* Revert python pin changes

---------

Co-authored-by: Mathieu Boudreau <emb6150@gmail.com>

Coverage Stats

13 of 48 new or added lines in 2 files covered. (27.08%)

1947 of 2473 relevant lines covered (78.73%)

2.21 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

96.05

/AxonDeepSeg/apply_model.py

from pathlib import Path
import os
import numpy as np
import torch
from PIL import Image
from loguru import logger
from typing import List, Literal, NoReturn

# AxonDeepSeg imports
from AxonDeepSeg.visualization.merge_masks import merge_masks
from AxonDeepSeg import ads_utils
from AxonDeepSeg.ads_utils import _LARGE_IMAGE_PIXEL_LIMIT
from AxonDeepSeg.params import nnunet_suffix, intensity

os.environ['nnUNet_raw'] = 'UNDEFINED'
os.environ['nnUNet_results'] = 'UNDEFINED'
os.environ['nnUNet_preprocessed'] = 'UNDEFINED'
from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor

def get_checkpoint_name(checkpoint_folder_path: Path) -> str:
    '''
    Get the name of the checkpoint file in the given folder, with priority for 
    best validation checkpoint.

    Parameters
    ----------
    checkpoint_folder_path : pathlib.Path
        Path to the folder containing the .pth checkpoint file.

    Returns
    -------
    str
        Name of the checkpoint file, e.g. 'checkpoint_best.pth'.
    '''
    if (checkpoint_folder_path / 'checkpoint_best.pth').exists():
        return 'checkpoint_best.pth'
    elif (checkpoint_folder_path / 'checkpoint_final.pth').exists():
        return 'checkpoint_final.pth'
    else:
        # Return checkpoint with most recent modification time
        checkpoints_namesorted=sorted(checkpoint_folder_path.glob('*.pth'))
        return checkpoints_namesorted[-1].name


def extract_from_nnunet_prediction(pred, pred_path, class_name, class_value) -> str:
    '''
    Extracts the given class from the nnunet raw prediction, saves it in a 
    separate mask and return the path of the extracted mask.

    Parameters
    ----------
    pred : np.ndarray
        The raw prediction from nnunet with values 0, 1, 2, ...
    pred_path : pathlib.Path
        Path to the raw prediction file; We expect its filename to end 
        with '_seg-nnunet.png'
    class_name : str
        Name of the class to extract. e.g. 'axon', 'myelin', etc.
    class_value : int
        Value of the class in the raw prediction.

    Errors
    ------
    ValueError
        If the class value is not found in the raw prediction.
    ValueError
        If the raw nnunet prediction file does not end with '_seg-nnunet.png'.

    Returns
    -------
    new_fname : str
        Path to the extracted class mask saved.
    '''

    pred_path = ads_utils.convert_path(pred_path)

    if not np.any(pred == class_value):
        logger.warning(f'Class value {class_value} not found in the raw prediction.')
    
    if not pred_path.name.endswith(str(nnunet_suffix)):
        raise NameError(f'Raw nnunet pred file does not end with "{nnunet_suffix}".')
    
    extraction = np.zeros_like(pred)
    extraction[pred == class_value] = intensity['binary']
    new_fname = str(pred_path).replace(str(nnunet_suffix), f'_seg-{class_name}.png')
    ads_utils.imwrite(new_fname, extraction)

    return new_fname

def find_folds(
            path_model: Path,
            model_type: Literal['light', 'ensemble']='light',
            ) -> List:
    '''
    For a given model, find the folders containing the folds

    Parameters
    ----------
    path_model : pathlib.Path
        Path to the folder model
    model_type :  Literal['light', 'ensemble'], optional
        Type of model, by default 'light'.       

    Returns
    -------
    List
        List of paths to the folds folders.
    '''
    
    if model_type == 'light':
        folds_avail = ['all']
    else:
        folds_avail = [str(f).split('_')[-1] for f in path_model.glob('fold_*')]

    return folds_avail

def axon_segmentation(
                    path_inputs: List[Path],
                    path_model: Path,
                    model_type: Literal['light', 'ensemble']='light',
                    gpu_id: int=-1,
                    verbosity_level: int=0,
                    allow_large_images: bool=False,
                    ) -> NoReturn:
    '''
    Segment images by applying a nnU-Net pretrained model.

    Parameters
    ----------
    path_inputs : List[pathlib.Path]
        List of images to segment. We assume they all exist and are already in 
        the correct format expected by the model (nb of channels, image format).
    path_model : pathlib.Path
        Path to the folder of the nnU-Net pretrained model. We assume it exists.
    model_type : Literal['light', 'ensemble'], optional
        Type of model, by default 'light'.
    gpu_id : int, optional
        GPU ID to use for cuda acceleration. -1 to use CPU, by default -1.
    verbosity_level : int, optional
        Level of verbosity, by default 0.
    '''
    # Raise PIL's pixel limit before nnUNet spawns its workers so that large images
    # can be read during preprocessing. On Linux (fork), workers inherit the parent's
    # module state. On macOS (spawn), workers start fresh but inherit env vars — the
    # ads_pil_patch.pth hook in site-packages checks this var at startup.
    if allow_large_images:
        Image.MAX_IMAGE_PIXELS = _LARGE_IMAGE_PIXEL_LIMIT
        os.environ["ADS_ALLOW_LARGE_IMAGES"] = "1"

    # find all available folds
    folds_avail = find_folds(path_model, model_type)

    if torch.cuda.is_available() and gpu_id >= 0:
        device = torch.device('cuda', gpu_id)
    elif torch.backends.mps.is_available():
        logger.info('MPS device detected. Using MPS for inference.')
        device = torch.device('mps')
    else:
        device = torch.device('cpu')
    # instantiate predictor
    predictor = nnUNetPredictor(
        perform_everything_on_device=True if device.type == 'cuda' else False,
        device=device,
    )
    logger.info('Running inference on device: {}'.format(predictor.device))

    # find checkpoint name (identical for all folds)
    chkpt_name = get_checkpoint_name(path_model / f'fold_{folds_avail[0]}')
    # init network architecture and load checkpoint
    predictor.initialize_from_trained_model_folder(
        str(path_model),
        use_folds=folds_avail,
        checkpoint_name=chkpt_name,
    )
    logger.info('Model successfully loaded.')

    # create input list
    input_list = [ [str(p)] for p in path_inputs]
    target_suffix = str(nnunet_suffix.with_suffix(''))
    data_format = predictor.dataset_json['file_ending'] # e.g. '.png'
    output_list = [ str(p).replace(data_format, target_suffix) for p in path_inputs ]

    predictor.predict_from_files(
        list_of_lists_or_source_folder=input_list,
        output_folder_or_list_of_truncated_output_files=output_list,
        save_probabilities=False,
        overwrite=True,
    )

    # Clean up env var so it doesn't leak to unrelated subprocesses later
    os.environ.pop("ADS_ALLOW_LARGE_IMAGES", None)

    output_structure = predictor.dataset_json['labels']
    output_classes = sorted(list(output_structure.keys()))
    output_classes.remove('background')
    is_axonmyelin_seg = ['axon', 'myelin'] == output_classes

    # nnUNet outputs a single file will all classes mapped to consecutive ints
    for pred_path in output_list:
        fname = pred_path + data_format
        raw_pred = ads_utils.imread(fname)
        new_masks = []

        for c in output_classes:
            class_value = output_structure[c]
            new_fname = extract_from_nnunet_prediction(raw_pred, fname, c, class_value)
            new_masks.append(new_fname)
        logger.info(f'Successfully saved masks for classes: {output_classes}.')

        if is_axonmyelin_seg:
            merge_masks(new_masks[0], new_masks[1])

        Path(fname).unlink()

1	from pathlib import Path	3✔
2	import os	3✔
3	import numpy as np	3✔
4	import torch	3✔
5	from PIL import Image	3✔
6	from loguru import logger	3✔
7	from typing import List, Literal, NoReturn	3✔
8
9	# AxonDeepSeg imports
10	from AxonDeepSeg.visualization.merge_masks import merge_masks	3✔
11	from AxonDeepSeg import ads_utils	3✔
12	from AxonDeepSeg.ads_utils import _LARGE_IMAGE_PIXEL_LIMIT	3✔
13	from AxonDeepSeg.params import nnunet_suffix, intensity	3✔
14
15	os.environ['nnUNet_raw'] = 'UNDEFINED'	3✔
16	os.environ['nnUNet_results'] = 'UNDEFINED'	3✔
17	os.environ['nnUNet_preprocessed'] = 'UNDEFINED'	3✔
18	from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor	3✔
19
20	def get_checkpoint_name(checkpoint_folder_path: Path) -> str:	3✔
21	'''
22	Get the name of the checkpoint file in the given folder, with priority for
23	best validation checkpoint.
24
25	Parameters
26	----------
27	checkpoint_folder_path : pathlib.Path
28	Path to the folder containing the .pth checkpoint file.
29
30	Returns
31	-------
32	str
33	Name of the checkpoint file, e.g. 'checkpoint_best.pth'.
34	'''
35	if (checkpoint_folder_path / 'checkpoint_best.pth').exists():	3✔
36	return 'checkpoint_best.pth'	3✔
37	elif (checkpoint_folder_path / 'checkpoint_final.pth').exists():	3✔
38	return 'checkpoint_final.pth'	3✔
39	else:
40	# Return checkpoint with most recent modification time
41	checkpoints_namesorted=sorted(checkpoint_folder_path.glob('*.pth'))	3✔
42	return checkpoints_namesorted[-1].name	3✔
43
44
45	def extract_from_nnunet_prediction(pred, pred_path, class_name, class_value) -> str:	3✔
46	'''
47	Extracts the given class from the nnunet raw prediction, saves it in a
48	separate mask and return the path of the extracted mask.
49
50	Parameters
51	----------
52	pred : np.ndarray
53	The raw prediction from nnunet with values 0, 1, 2, ...
54	pred_path : pathlib.Path
55	Path to the raw prediction file; We expect its filename to end
56	with '_seg-nnunet.png'
57	class_name : str
58	Name of the class to extract. e.g. 'axon', 'myelin', etc.
59	class_value : int
60	Value of the class in the raw prediction.
61
62	Errors
63	------
64	ValueError
65	If the class value is not found in the raw prediction.
66	ValueError
67	If the raw nnunet prediction file does not end with '_seg-nnunet.png'.
68
69	Returns
70	-------
71	new_fname : str
72	Path to the extracted class mask saved.
73	'''
74
75	pred_path = ads_utils.convert_path(pred_path)	3✔
76
77	if not np.any(pred == class_value):	3✔
78	logger.warning(f'Class value {class_value} not found in the raw prediction.')	3✔
79
80	if not pred_path.name.endswith(str(nnunet_suffix)):	3✔
81	raise NameError(f'Raw nnunet pred file does not end with "{nnunet_suffix}".')	3✔
82
83	extraction = np.zeros_like(pred)	3✔
84	extraction[pred == class_value] = intensity['binary']	3✔
85	new_fname = str(pred_path).replace(str(nnunet_suffix), f'_seg-{class_name}.png')	3✔
86	ads_utils.imwrite(new_fname, extraction)	3✔
87
88	return new_fname	3✔
89
90	def find_folds(	3✔
91	path_model: Path,
92	model_type: Literal['light', 'ensemble']='light',
93	) -> List:
94	'''
95	For a given model, find the folders containing the folds
96
97	Parameters
98	----------
99	path_model : pathlib.Path
100	Path to the folder model
101	model_type : Literal['light', 'ensemble'], optional
102	Type of model, by default 'light'.
103
104	Returns
105	-------
106	List
107	List of paths to the folds folders.
108	'''
109
110	if model_type == 'light':	3✔
111	folds_avail = ['all']	3✔
112	else:
113	folds_avail = [str(f).split('_')[-1] for f in path_model.glob('fold_*')]	3✔
114
115	return folds_avail	3✔
116
117	def axon_segmentation(	3✔
118	path_inputs: List[Path],
119	path_model: Path,
120	model_type: Literal['light', 'ensemble']='light',
121	gpu_id: int=-1,
122	verbosity_level: int=0,
123	allow_large_images: bool=False,
124	) -> NoReturn:
125	'''
126	Segment images by applying a nnU-Net pretrained model.
127
128	Parameters
129	----------
130	path_inputs : List[pathlib.Path]
131	List of images to segment. We assume they all exist and are already in
132	the correct format expected by the model (nb of channels, image format).
133	path_model : pathlib.Path
134	Path to the folder of the nnU-Net pretrained model. We assume it exists.
135	model_type : Literal['light', 'ensemble'], optional
136	Type of model, by default 'light'.
137	gpu_id : int, optional
138	GPU ID to use for cuda acceleration. -1 to use CPU, by default -1.
139	verbosity_level : int, optional
140	Level of verbosity, by default 0.
141	'''
142	# Raise PIL's pixel limit before nnUNet spawns its workers so that large images
143	# can be read during preprocessing. On Linux (fork), workers inherit the parent's
144	# module state. On macOS (spawn), workers start fresh but inherit env vars — the
145	# ads_pil_patch.pth hook in site-packages checks this var at startup.
146	if allow_large_images:	3✔
147	Image.MAX_IMAGE_PIXELS = _LARGE_IMAGE_PIXEL_LIMIT	×
148	os.environ["ADS_ALLOW_LARGE_IMAGES"] = "1"	×
149
150	# find all available folds
151	folds_avail = find_folds(path_model, model_type)	3✔
152
153	if torch.cuda.is_available() and gpu_id >= 0:	3✔
NEW 154	device = torch.device('cuda', gpu_id)	×
155	elif torch.backends.mps.is_available():	3✔
156	logger.info('MPS device detected. Using MPS for inference.')	1✔
157	device = torch.device('mps')	1✔
158	else:
159	device = torch.device('cpu')	2✔
160	# instantiate predictor
161	predictor = nnUNetPredictor(	3✔
162	perform_everything_on_device=True if device.type == 'cuda' else False,
163	device=device,
164	)
165	logger.info('Running inference on device: {}'.format(predictor.device))	3✔
166
167	# find checkpoint name (identical for all folds)
168	chkpt_name = get_checkpoint_name(path_model / f'fold_{folds_avail[0]}')	3✔
169	# init network architecture and load checkpoint
170	predictor.initialize_from_trained_model_folder(	3✔
171	str(path_model),
172	use_folds=folds_avail,
173	checkpoint_name=chkpt_name,
174	)
175	logger.info('Model successfully loaded.')	3✔
176
177	# create input list
178	input_list = [ [str(p)] for p in path_inputs]	3✔
179	target_suffix = str(nnunet_suffix.with_suffix(''))	3✔
180	data_format = predictor.dataset_json['file_ending'] # e.g. '.png'	3✔
181	output_list = [ str(p).replace(data_format, target_suffix) for p in path_inputs ]	3✔
182
183	predictor.predict_from_files(	3✔
184	list_of_lists_or_source_folder=input_list,
185	output_folder_or_list_of_truncated_output_files=output_list,
186	save_probabilities=False,
187	overwrite=True,
188	)
189
190	# Clean up env var so it doesn't leak to unrelated subprocesses later
191	os.environ.pop("ADS_ALLOW_LARGE_IMAGES", None)	3✔
192
193	output_structure = predictor.dataset_json['labels']	3✔
194	output_classes = sorted(list(output_structure.keys()))	3✔
195	output_classes.remove('background')	3✔
196	is_axonmyelin_seg = ['axon', 'myelin'] == output_classes	3✔
197
198	# nnUNet outputs a single file will all classes mapped to consecutive ints
199	for pred_path in output_list:	3✔
200	fname = pred_path + data_format	3✔
201	raw_pred = ads_utils.imread(fname)	3✔
202	new_masks = []	3✔
203
204	for c in output_classes:	3✔
205	class_value = output_structure[c]	3✔
206	new_fname = extract_from_nnunet_prediction(raw_pred, fname, c, class_value)	3✔
207	new_masks.append(new_fname)	3✔
208	logger.info(f'Successfully saved masks for classes: {output_classes}.')	3✔
209
210	if is_axonmyelin_seg:	3✔
211	merge_masks(new_masks[0], new_masks[1])	3✔
212
213	Path(fname).unlink()	3✔

axondeepseg / axondeepseg / 30686229961

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous