save image

[mouxinyue1]

Hello, I want to change the saved result image into grayscale rather than binary graph, may I ask what part of the code to modify

[milesial]

The image saving is right there in predict.py https://github.com/milesial/Pytorch-UNet/blob/40d5ba797c5689cd9560233baa0e52f28f92727c/predict.py#L62

If you have num_classes > 1 you should remove the argmax in the mask prediction and replace with whatever fits you to convert multiple classes to greyscale. If you have num_classes ==1 you should remove the thresholding in the mask prediction.