Working with multi-class mask

[ronlos]

Hi,

I used your code for semantic segmentation on my own data and it worked great!
Now I want to work on data that has more than 1 class.
not sure how am I supposed to build the masks so your code will work with them.
Lets say I have 6 classes - My mask should be 1918 x 1280 x 3 x 6 with each layer being a binary mask for each class?
or does it need to be 1918 x 1280 x 1 with different grayscale value for each class?
or maybe something else?

Thanks,
Rony.

[milesial]

If you're using RGB images as input and 6 classes, your images are 1918 x 1280 x 3 JPEG or PNG, or whatever PIL.Image.open can read.
Your masks can be 1918 x 1280 numpy or pytorch tensors with class indices 0..5, stored as .npy or .pt. They could also be grayscale image with grey values 0..5 if PIL.Image.open reads that correctly.

Check out https://github.com/milesial/Pytorch-UNet/blob/master/utils/data_loading.py.
You can also modify the data loading code so that it matches you specific data format.

[ronlos]

That's what i'm doing...
for some reason I keep getting this error - RuntimeError: 1only batches of spatial targets supported (3D tensors) but got targets of size: : [1, 1, 640, 959].
When I check my images size with img.shape I get - (1280, 1918, 3).

What am I missing here?
tried to run with squeeze(1) and it seems to run there's no actually much learning going on..
How do I get my loader to load the images as 3D and not 4D?