在图像分割中,最后的输出结果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')