如何在python中构建OneHot解码器

我有encoded我的图片（口罩），尺寸（img_width X IMG_HEIGHT×1）用OneHotEncoder这种方式：

import numpy as np

def OneHotEncoding(im,n_classes):
  one_hot = np.zeros((im.shape[0], im.shape[1], n_classes),dtype=np.uint8)
  for i, unique_value in enumerate(np.unique(im)):
    one_hot[:, :, i][im == unique_value] = 1
  return one_hot

在使用深度学习进行了一些数据操作之后，softmax激活函数将导致概率而不是0and1值，因此在我的解码器中，我想实现以下方法：

1. 阈值获得0或1仅获得。
2. 将每个通道的权重等于通道索引。
3. 沿通道轴取标签之间的最大值。

import numpy as np

arr = np.array([
    [[0.1,0.2,0,5],[0.2,0.4,0.7],[0.3,0.5,0.8]],
    [[0.3,0.6,0  ],[0.4,0.9,0.1],[0  ,0  ,0.2]],
    [[0.7,0.1,0.1],[0,6,0.1,0.1],[0.6,0.6,0.3]],
    [[0.6,0.2,0.3],[0.4,0.5,0.3],[0.1,0.2,0.7]]
])

# print(arr.dtype,arr.shape)

def oneHotDecoder(img):
    # Thresholding
    img[img<0.5]=0
    img[img>=0.5]=1
    # weigts of the labels
    img = [i*img[:,:,i] for i in range(img.shape[2])]
    # take the max label
    img = np.amax(img,axis=2)
    print(img.shape)
    return img

arr2 = oneHotDecoder(arr)
print(arr2)

我的问题是：

1. 如何摆脱错误： line 15, in oneHotDecoder img[img<0.5]=0 TypeError: '<' not supported between instances of 'list' and 'float'
2. 您建议我改善实施中的其他任何问题吗？

提前致谢。