Hi, I am working on Segmentation of Camvid dataset. Masks are of shape [h,w,3]. Indeed for each class instead of the label of a unique class we have a color_code. I found a code for segmentation of this dataset which in order to use CrossEntropyLoss, a snippet was used that convert each mask from shape [h,w,3] to [32,h,w], where 32 is the number of classes. There is also another snippet that converts a mask of shape [32,h,w] to a 3-channel mask of shape [h,w,3]. Now, when I finished training and visualize predicted mask, the result is not good at all. I am exploring where I did a mistake. First I checked the two snippet for converting a RGB mask to a Binary mask. The codes are below.
def Color_map(df):
    
    '''
    Returns the reversed String.
    Parameters:
        dataframe: A Dataframe with rgb values with class maps.
    Returns:
        code2id: A dictionary with color as keys and class id as values.   
        id2code: A dictionary with class id as keys and color as values.
        name2id: A dictionary with class name as keys and class id as values.
        id2name: A dictionary with class id as keys and class name as values.
    '''
    cls = pd.read_csv(df)
    
    # thic line of code tuples the code of colors
    # len(cls.name) is the number of classes that we have: 32
    # output: [(64, 128, 64), (192, 0, 128)]
    color_code = [tuple(cls.drop("name",axis=1).loc[idx]) for idx in range(len(cls.name))]
    
    # it gives a number to each code
    # assigns color codes to id numbers
    code2id = {v: k for k, v in enumerate(list(color_code))}
    
    # it assigns numbers(classes) to codes
    id2code = {k: v for k, v in enumerate(list(color_code))}
    
    # it collects name of each class 
    color_name = [cls['name'][idx] for idx in range(len(cls.name))]
    
    # it gives to each class a number
    name2id = {v: k for k, v in enumerate(list(color_name))}
    
    # it gives
    id2name = {k: v for k, v in enumerate(list(color_name))}  
    
    return(code2id, id2code, name2id, id2name)

def mask_to_rgb(mask, id2code):
    ''' 
        Converts a Binary Mask of shape: [batch_size,num_classes,h,w] 
        to RGB image mask of shape [batch_size, h, w, color_code]
        
        Parameters:
            img: A Binary mask
            color_map: Dictionary representing color mappings
        returns:
            out: A RGB mask of shape [batch_size, h, w, color_code]
    '''
    ## Since our mask is one-hot encoding
    ## the argmax returns the output class for each pixel
    ## It returns the label of each pixel that is a number in range : 0-31
    ## dim 0 :batch_size
    
    single_layer = np.argmax(mask, axis=1)
    
    ## it converts each mask to [batch_size, h,w, color_code]
    output = np.zeros((mask.shape[0],mask.shape[2],mask.shape[3],3))
    
    for k in id2code.keys():
        
        output[single_layer==k] = id2code[k]
        
    return(output.astype(np.float32))

def rgb_to_mask(img, id2code):
    ''' 
        Converts a RGB image mask of shape [batch_size,h, w, color_code], to a mask of shape
        [batch_size,n_classes,h,w]
        
        Parameters:
            img: A RGB img mask
            color_map: Dictionary representing color mappings: ecah class assigns to a unique color code
        returns:
            out: A Binary Mask of shape [batch_size, classes, h, w]
    '''
    
    # num_classes is equal to len(mask) 
    num_classes = len(id2code)
    
    # it makes a tensor of shape h,w,num_classes:(720,960,num_classes)
    shape = img.shape[:2]+(num_classes,)
    
    # it makes a tensor with given shape and with type float64
    out = np.zeros(shape, dtype=np.float64)
    
    # 
    for i, cls in enumerate(id2code):
        
        #print(f'i: {i}, cls: {cls}')
        
        # img.reshape((-1,3)) flats mask except in channels
        
        # it reads thecolor code for a multiplication of higght and width and if it is one of the color code of 
        # the classes that we have then the third dimension takes the label of that class and the first
        # two dimsnions return to the hight and width  
        out[:,:,i] = np.all(np.array(img).reshape((-1,3)) == id2code[i], axis=1).reshape(shape[:2])
        
        # out: hight, width, class
        # returns class, hight, width
    return(out.transpose(2,0,1))

I expect for a mask in training set, when I convert it to a binary mask using function rgb_to_mask and then convert it again to an rgb-mask using function mask_to_rgb the result would be the same as the original mask. But they are not the same as it can be seen in the following code. I do not know where the problem is.
print(f'mask_sample_shape: {mask_sample.shape} mask_sample: {mask_sample.dtype}, mask_type: {type(mask_sample)}')
_, id2code,_,_ = Color_map(os.path.join(path,'class_dict.csv'))
mask_cls = rgb_to_mask(mask_sample, id2code)
print(f'mask_cls: {mask_cls.shape} mask_cls:{mask_cls.dtype} mask_cls_type: {type(mask_cls)}')
## Now converting mask_cls to mask_rgb
mask_rgb = mask_to_rgb(mask_cls[np.newaxis,...], id2code)
mask_rgb = mask_rgb.squeeze(0)
print(f'mask_rgb_shape: {mask_rgb.shape} mask_rgb: {mask_rgb.dtype}, mask_rgb_type: {type(mask_rgb)}')
comparison = mask_rgb==mask_sample
print(comparison.all())

out:
mask_sample_shape: (720, 960, 3) mask_sample: float32, mask_type: <class 'numpy.ndarray'>
mask_cls: (32, 720, 960) mask_cls:float32 mask_cls_type: <class 'numpy.ndarray'>
mask_rgb_shape: (720, 960, 3) mask_rgb: float32, mask_rgb_type: <class 'numpy.ndarray'>
False

Dose anyone have any idea where is the problem? I also visualized both masks at the end but the result was not the same.
Thanks