圖像分割中mask的保存

​在圖像分割中，最後的輸出結果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')

​