Accessing the Logits

[AlexanderZeilmann]

Is it possible to access the segmentation logits, i.e., the raw output of the neural network before the final argmax operation?

[AlexanderZeilmann]

The reason why I am asking is that this would make it easier to

1. estimate the uncertainty of the segmentation at each pixel and
2. support multi-class segmentation.

From what I see here :

nnInteractive/nnInteractive/inference/inference_session.py

Line 443 in 9192ecf

pred = self.network(input_for_predict[None])[0].argmax(0).detach()

and here:

nnInteractive/nnInteractive/inference/inference_session.py

Line 566 in 9192ecf

pred = self.network(preallocated_input[None])[0].argmax(0).detach().cpu()

it looks like as if this is currently not possible.

Maybe adding something like

logits = self.network(input_for_predict[None])[0]
if write_logits:
   target_buffer_logits = logits
pred = logits.argmax(0).detach() 

would make sense to support the above mentioned use cases.

[sasha-360]

is this model supportted for multi-class segmentation?

[AlexanderZeilmann]

is this model supportted for multi-class segmentation?

No, but with access to the logits one can at least implement a simplistic multi-class segmentation.

[FabianIsensee]

Accessing the logits is not easy as the predictions are always only generated in a part of the image and we immediately convert them into a segmentation. I pushed a new version of the inference code today which untangles a lot of the mess I had previously but it will still be a challenge to modify it so that the logits can be exported