How to aggreate multi-channel labels back to one channel

[LIKP0]

Hello all! Thanks for your excellent work.
I'm doing multi-class segmentation work and I train a basic UNet with out_channels=30, which denotes I have 30 labels including background. After use the saver provided by the turtorial saver = SaveImage(output_dir="./output", output_ext=".nii.gz", output_postfix="seg"), I find the label shape is (30, 224, 192, 160). I want to aggreate the labels to be (1, 224, 192, 160).
I have searched the transforms in dictionary.py but I can't find any proper method. I'm new to the lib and could anyone tell me if there is a flexible method to do this?
Really appreciate your help!

[LIKP0]

I think a simple script can help. Notice that use np.where() is a lot faster than a rude "for"

import numpy as np
from glob import glob
import SimpleITK as sitk

if __name__ == "__main__":

    in_dir = '/mnt/lijl/MONAI/Result/Upperbound/inferTs/'
    out_dir = '/mnt/lijl/MONAI/Result/Upperbound/trans_seg/'

    img_list = sorted(glob(in_dir + '*.nii.gz'))
    for img_name in img_list:
        index = img_name.split('/')[-1]
        print(index)
        img = sitk.GetArrayFromImage(sitk.ReadImage(img_name))

        sh = img.shape
        print(sh)
        res = np.zeros([sh[1], sh[2], sh[3]])
        for i in range(1, sh[0]):
            channel = img[i]
            res[np.where(channel == 1)] = i

        sitk.WriteImage(sitk.GetImageFromArray(res), out_dir + index)