These lines can help some one to transform grayscale mask to multichannel (equal to number of classes) one hot mask.

gray2hot.py

def gray2hot(graymask, num_classes=2, include_background=True):                           
    '''Convert grayscale image of mask to multichannel mask                               
    if include_background=True number of channels in output will be             
    num_classes + 1                                                             
    '''                                                                         
    graymask = graymask.astype(np.int32)                                                  
    onehot_mask = np.zeros((graymask.shape[0], graymask.shape[1], num_classes+include_background), dtype=np.int32) # +1 relate to background                                  
    if include_background:                                                                
        onehot_mask[:, :, 0] = 1 # initialize all as background at first                  
    if np.all(graymask == 0):                                                             
        return onehot_mask                                                                
    indx = list(np.nonzero(graymask))                                                     
    indx.append((graymask[np.nonzero(graymask)]-(not include_background)))                
    onehot_mask[tuple(indx)] = 1                                                          
    if include_background:                                                                
        onehot_mask[:, :, 0] = onehot_mask[:, :, 0] - np.sum(onehot_mask[...,1:], axis=-1) # remove backgound from where foreground exist
    return onehot_mask   

test_gray2hot.ipynb