Преобразование изображений в горячий массив и обратно возвращает черные изображения

У меня есть массив numpy массива numpy из 3 изображений в градациях серого только со значениями 0 и 255 (shape: (3, 512, 512)). Я превратил их в горячие кодировки в 2 класса, используя tf.one_hot (shape: (3, 512, 512, 2)). Сейчас пытаюсь вернуть их в первоначальный вид через np.argmax. Но когда я читаю изображения, они все черные.

    labels = []
    labels.append(label_1)
    labels.append(label_2)
    labels.append(label_3)
    labels = np.asarray(labels)

    print(labels.shape)                # (3, 512, 512)

    sess = tf.InteractiveSession()

    labels = tf.one_hot(labels, 2)     # (3, 512, 52, 2)
    print(labels.shape)
    #print(labels)

    labels = labels.eval()             # to convert to numpy array from tensor

    imgs = np.argmax(labels, axis=3)   # also tried tf.argmax
    print(imgs.shape)                  # (3, 512, 512)

    for i in range(imgs.shape[0]):
        image = imgs[i]
        print(img.shape)               # (512, 512)
        indices = np.where(image > 0) 
        print(indices)                 # array([], dtype=int64), array([], dtype=int64)
        print(indices)
        image = Image.fromarray(image, 'L')
        image.show()                   # black images, values all zero

Я уверен, что мне не хватает чего-то довольно простого, но я не могу понять это. Любая помощь будет оценена. Спасибо.

Изменить:

Я проверил индексы, где массив имеет ненулевые значения. Но все значения кажутся нулями. Таким образом, не проблема отображения изображения. Думаю проблема с argmax но не знаю в чем.