Pytorch语义分割理解

贴出部分pytorch的语义分割代码理解

 

#读取图片
image =PIL.Image.open("D:\\DataSets\\PascalVOC2012\\VOCdevkit\\VOC2012\\JPEGImages\\2010_003055.jpg")
#图片预处理，转化为0-1之间
image_transf =transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485,0.456,0.406],
                        std =[0.229,0.224,0.225])
])
image_tensor =image_transf(image).unsqueeze(0)
output =model(image_tensor)["out"]
#将输出转化为二维图像
outputarg =torch.argmax(output.squeeze(),dim=0).numpy()
outputarg

上述代码将图片放入网络模型中，输出图片中每个像素点的预测值，因为有20类，所以输出范围在【0-20】之间，此图片输出结果为：

#对得到的输出结果进行编码
def decode_segmaps(image,label_colors,nc=21):
    #函数将输出为2D图像，会将不同的类编码为不同的颜色
    r =np.zeros_like(image).astype(np.uint8)
    g =np.zeros_like(image).astype(np.uint8)
    b =np.zeros_like(image).astype(np.uint8)
    print("r:",r)
    print("r[0]:",r[0])
    for cla in range(0,nc):
        idx =image ==cla
        r[idx]=label_colors[cla,0]
        g[idx] =label_colors[cla,1]
        b[idx] =label_colors[cla,2]
"""
        print("image:",image)
        print("idx:",idx)
        print("cla:",cla)
        print("r[idx]:",r[idx])
        print("g[idx]:",g[idx])
        print("b[idx]:",b[idx])
"""
    rgbimage =np.stack([r,g,b],axis =2)
   # print(rgbimage)
    return rgbimage

上述代码令我头疼了老半天，可算是理解了

 r =np.zeros_like(image).astype(np.uint8)
 g =np.zeros_like(image).astype(np.uint8)
 b =np.zeros_like(image).astype(np.uint8)

 用于生成R、G、B三通道各自与原图片相对应的矩阵

idx =image ==cla

idx是一个bool变量，true为1，false为0，解释：

image作为一个经过模型预测后输出的图片分类矩阵，其每个数都落在【0-20】之间，

cla作为一个for循环的数值，范围同样是落在【0-20】之间

这里判断image矩阵上的每个数是否与cla相等，如果相等，则代表image图像上的这个像素点的分类就是cla，假设当cla=0时，image的图像分类矩阵为：

 此时上面那一部分为0，与cla相等，说明上面一部分的分类类别为0，也就是label对应的“背景”，这一部分的idx为True（对应的为Image图像中为True的像素点的位置），下面的r,g,b会有相对应的值

而下面对应的15与cla不相等，说明这一部分的idx为False，下面的r,g,b为空

 r[idx]=label_colors[cla,0]
 g[idx] =label_colors[cla,1]
 b[idx] =label_colors[cla,2]

这里的r[idx]代表的是第几行，label_colors[cla,0]代表第cla类别上的第0通道（r通道）上的灰度值

rgbimage =np.stack([r,g,b],axis =2)

np.stack（）将上述三元素组合，转化成一个三通道的图像（RGB），输出便为语义分割后的图像。