Get binary output from softmax output

Hi, I’m working on an image segmentation problem. Once I have the output in the form of softmax output, how can I get a binary image from the Variable? More generally, if I have multi-classes, what is the way to convert the softmax output to a single channel image? Thanks!

Maybe taking argmax?

Thanks. I searched over the forum and found out that should be the way. Thank you!