Can augmentedImageDatastores (or augmentedImageSources) be used in a semantic segmentation neural network?

Hi,

Having read through the documentation for these data types, it's not super obvious whether they can be used for a semantic segmentation CNN. My guess is no, because I've fed the augmentedImageSource my images in a 4D array (512x512x1x80) and the corresponding segmentation labels (also 512x512x1x80). Now when I feed the auImSo the labels as the raw images (which are just 512x512 images with pixel values either 1 or 2 for the two different classes) the error I get when attempting to train a network is "The output size (2) of the last layer doesn't match the number of classes (1)". So I then tried converting the 4D label image array to a categorical, but the error then is that "The number of observations in X and Y disagree", where X is the 4D array of images and Y the labels. Is there something I'm missing here?

Thanks all.