图像分割中mask的保存_pytorch保存图像分割mask-优快云博客

博客介绍了图像分割中mask结果的处理步骤。mask结果是N x W x H的tensor，处理时先转成numpy的array，再将N x W x H转为1 x W x H，原理是判断像素得分最高的类别并置为其序号，最后转成Image调色板着色，第一步和第二步可互换，pytorch和numpy有相关计算函数。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

在图像分割中，最后的输出结果mask结果是一个N x W x H的tensor，其中N表示有多少个类别，W是图像的宽度，H是图像的高度，每一个W x H都是一个二值矩阵

第一步，转成numpy中的array

第二步，N x W x H --> 1 x W x H（W x H），原理就是，对图像中的每一个像素，判断它在哪一类中的得分最高，然后把像素值置为得分最高类的序号。

第三步，转成Image，调色板着色

其中，第一步和第二步可以互换，pytorch和numpy中都有计算第二部的函数

pytorch中，axis是根据哪一维来计算最大值，[1]表示保存索引，[0]是保存数据，后面是转成numpy

torch.max(outpout, axis)[1].cpu().numpy()

numpy.argmax

numpy.argmax(output, axis=0)

palette = [random.randint(0, 255) for x in range(256 * 3)] # 随机颜色的调色板

def colorize_mask(mask):
    # mask: numpy array of the mask
    new_mask = Image.fromarray(mask.astype(np.uint8)).convert('P')
    new_mask.putpalette(palette)

    return new_mask
mask = torch.max(output, 0)[1].cpu().numpy()
mask_img = colorize_mask(mask)
mask_img.save('filepath + feilename')