Doubt on nn.NLLLoss () for segmentation probelm


I’m applying nn.NLLLoss () for segmentation problems.
However, I don’t understand how to extract the segmentation from the final layer. I’ll explain.

  • Given an encoder-decoder model, the final decoder level releases a dimension tensor (N, C, W, H)
    to which softmax is applied, where N=minibach, C= number of class, WxH= image size.

  • the tensor target has the following structure (N,W,H), wherein WxH class 0 contains the object to be segmented (a), while class 1 contains the object to be segmented (b), a + b = total segmentation.

  • now my doubt: from the output decoder of the decoder I want to extract my two objects to be segmented (my two classes) therefore strictly logic:

    • output tesnor 1: (N,[0],W,H) – > object a
    • output tensor 2: (N,[1],W,H) – > object b
      to extract the elementation I will have to make a threshold or search for the value 0 or 1 in output tensor 1,2?

For example:

      t1_seg              = tensor_out_decoder[0]
       t2_seg              = tensor_out_decoder[1]
       a                       = np.zeros(shape=(seg.shape[1],seg.shape[2]))
       b                       = np.zeros(shape=(seg.shape[1],seg.shape[2]))
       idx_a                = t1_seg[:,:] > 0.5 
       idx_b                = t2_seg[:,:]  > 0.5

       a[idx_myo]        = 50
       b[idx_rv]            = 100