how to perform augmentation when segmentation mask has 4 channels

[ai1361720220000]

when it comes to segmentation task, the mask usually has 1 channels. but if the mask has 4 channels in my task, and i find it gets wrong as following code: the shape of image is (w, h, 3), the shape of mask is (w, h ,4)

segmap = np.zeros((128, 128, 4), dtype=np.int32)
segmap = SegmentationMapsOnImage(segmap , shape=image.shape)

please give me some advice, thanks!!

[half-potato]

split into 4 maps