Multi-class segmentation about medical-transformer

[jeya-maria-jose]

Yes, you can just change the num_classes argument to number of classes in

Medical-Transformer/lib/models/axialnet.py

Line 511 in c9c7064

def __init__(self, block, block_2, layers, num_classes=2, zero_init_residual=True,

to convert the code for multi-class segmentation.

[shanpriya3]

Thanks for the response. I have a ground-truth with 3 classes including background with values 0,127,255. As mentioned in your previous comment I changed num_classes=3 in axialnet.py In utils.py,

Medical-Transformer/utils.py

Lines 156 to 157 in 703a080

	mask[mask<=127] = 0
	mask[mask>127] = 1

mask[mask<=127] = 0
mask[mask>127] = 1

which makes the ground-truth with values 0 and 1 but I should have 0,1 and 2 for my case(with 3 classes).

I tried doing this
mask[mask<127] = 0
mask[mask==127] = 1
mask[mask>127] = 2

But I got this error. Could you please help me with this?