How is upsampling perform (at inference) ?

[romainVala]

Hello

I trained for (MRI) brain segmentation a nnUNet model with a dataset at 0.75 mm iso.
At inference , I tested a very hig resolution input volume (0.25 mm)

If I understand correctly the way nnunet process, is to first resample the input volume to the training set resolution (0.75) and then upsample back the result.

Can you describe (or point me the corresponding code) how this is perform ?

I try to reproduce :
If I perform this two step myself, the easiest (usual) way to perform it is to use nearest interpolation to upsample the (binary) label. The result is an effective resolution much bigger compare to result of nn-unet -> nnunet gives me very smooth boundary whereas the nearest interpolation does not improve the resolution (ie I do see pixelesize boudary (~0.75mm) )

Actually I found a way to improve the upsampling of a label map, (not sure if this is well know ?) :
The idea is to first transform the 3D label volume to a 4D onehot version and then to upsample each channel (with any interpolation scheme, spline, or cubic) and finally perform the argmax to end up with a binary 3D upsapled volume.

The biggest issue with this method, is time and memory: if there are too much label and too high resolution, the onehot is just too big ... is there a way to do it patch wise ?

I would be very interested to understand how you handle this in nnUNet.

Many thanks

Romain

[fhrzic]

Hello Romain,

The scripts you are seeking fore are in nnunetv2->inference->predict_from_raw_data.py. Main script called from them is in nnunetv2->inference->export_prediction.py (convert_predicted_logits_to_segmentation_with_correct_shape). That script is calling predicted_logits = configuration_manager.resampling_fn_probabilities() where it uses resampling_fn_probabilities which is partial function of proper resample method.

How is inference done? Essentially it is overlapping windows (script predict_sliding_window_return_logits ) which also implement gaussian weighting. It creates volume logits predictions and then carefully reassemble volume back from the patches.

Resample can be found in preprocessing->resampling. Now based on anisotropic and isotropic spacing different resample interpolations are used. Those are described in paper and are buried in code behind plans.

Generally:
1.) The nii.gz for prediction is resampled and normalize the same way as it is caluclated in the fingerprint of the whole experiment(preprocessing scripts).
2.) The prediction is done is sliding window manner where at some point everything is stored either on CPU or GPU. It can go OOM if there is no space. This is answer to your question, and to the best of my knowledge, it is possible that for some huge scans there is no enough memory to store everything in memory (especially if patch size is big).
3.) Once predictions are gathered, they are reassembled in main volume space by interpolations chosen at the very beginning of the experiment (take note that they are not the same for z and x,y axes as described in paper).

Time is slow, there is no argue about it because it heavily depends on input CT shape - the bigger the CT the more time and memory it consumes. In my opinion, if you want to do workaround and can sacrifice some precision - reshape the nifti file prior giving it to the nnUnet or break it down into several smaller ones. It is bad workaround but at least you should avoid memory problems. Without heavy alterations to the code, I am not seeing how you can easily adjust nnUnet code to your problem.

Best,
Franko

P.S. I am also reading nnUnet code and try to understand everything, there may be that I have not understand something properly, but this is all I know and understand for now. Code is hard, it is a lot of everything everywhere and you need to dwell really deep into it to understand how everything works. For 100% correct answer, Fabian is the boss.

[romainVala]

Hi Franko
thanks for taking time to explain.
Pointing the code was indeed helpfull, I can now confirm that label resampling is done after a onehote

I quote the resize_segmentation comment:
" Resizes a segmentation map. Supports all orders (see skimage documentation). Will transform segmentation map to one hot encoding which is resized and transformed back to a segmentation map.
This prevents interpolation artifacts ([0, 0, 2] -> [0, 1, 2])"

The small difference I do still see is that argmax is not use at the end to revert back to a 3D label map
instead a threshold is done ( >= 0.5 )

In case of voxel with more than 2 label in partial volume this could make a difference, (not sure if it is a big deal or not) . Anyway avoiding the argmax also allow to no have to store the full 4D onhot image and thus it spares a lot of memory ! ...

The last point which is not fully clear to me, is wether the resampling is done on the full volume (ie after resambling the sliding widows inference patches) or directly at the patch level before merging the patches ?

I forgot about possible anisotropy, but in case of isotropic voxel size for training I guess every thing is done isotropic

[fhrzic]

Hi Roman,

just let me clear a bit of code navigation to see where the prediction is done. Now I must make a disclosure, I am not sure which rules applies to your problem - but here is where you can find an answer:

After the code obtain predictions from e.g. predict_logits_from_preprocessed_data (predict_from_raw_data.py) it saves them -- to be precise export them to a seg map. This is done by function calling the convert_predicted_logits_to_segmentation_with_correct_shape (call can be seen in predict_from_raw_data.py, but function is located i nnnunetv2->inference->export_prediction.py)

In export_prediction.py you can see that segmentation is obtained from label_manager.convert_probabilities_to_segmentation(predicted_probabilities) which is located in LabelManager (nnunetv2->utilities->label_handling->label_handling.py).

By opening that script you can indeed see the following code snippet:

if self.has_regions: if isinstance(predicted_probabilities, np.ndarray): segmentation = np.zeros(predicted_probabilities.shape[1:], dtype=np.uint16) else: # no uint16 in torch segmentation = torch.zeros(predicted_probabilities.shape[1:], dtype=torch.int16, device=predicted_probabilities.device) for i, c in enumerate(self.regions_class_order): segmentation[predicted_probabilities[i] > 0.5] = c else: segmentation = predicted_probabilities.argmax(0)

At the end, you see argmax (but I am not sure if ti applies to your code). In case of multilabel, it should.

Based on my investigation and experience from other projects - I can confirm that you are right: avoiding saving logits and saving prediction saves a ton of memory, but you must be aware to successfully compute sliding windows you must keep logits in memory in case that windows overlap. If there is no overlapping between windows, you can avoid the storage of logits and store only the argmax values. Another approach that came to my mind is keep only the logits for windows of interest in memory, once when I know that for that particular voxel no further window will make adjustment, I calculate prediction and save that. However - to implement this it could be quite a journey with questionable impact on time utilization (it will save memory though).

Your last question, when the resampling is done can be answered by looking in export_prediction.py : convert_predicted_logits_to_segmentation_with_correct_shape . Based on what I can see, resampling is done on the logits of the whole volume, and then it is transferred to one class. (configuration_manager.resampling_fn_probabilities is followed by label_manager.convert_probabilities_to_segmentation).

For isotropic/anisotropic it is only the matter of which interpolation is being utilized. But TBH, I am not sure if it plays any major roles...

Franko