I changed the training set of the original code to a class 6 matrix input of 1x512x512

And the corresponding modification of the original Unet_conditional and other more code, but the training effect is very poor, there is a great god can guide it?