语义分割从heatmap恢复出分割mask

最新推荐文章于 2024-07-07 17:15:22 发布

DongHappyyy

最新推荐文章于 2024-07-07 17:15:22 发布

阅读量1.2k

点赞数

分类专栏： torch 文章标签：语义分割深度学习

本文链接：https://blog.csdn.net/weixin_39715012/article/details/120935406

版权

torch 专栏收录该内容

6 篇文章

订阅专栏

本文介绍了如何使用深度学习技术处理和转化热力图，并展示了convert_label函数的应用，通过inverse=True将预测结果映射到预定义的标签集。save_pred函数则用于保存预测结果，以不同颜色对应不同的类别。核心内容涉及图像识别和计算机视觉中的标签映射操作。

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

ignore_label=255,
label_mapping = {-1: ignore_label, 0: ignore_label,
                  1: ignore_label, 2: ignore_label,
                  3: ignore_label, 4: ignore_label,
                  5: ignore_label, 6: ignore_label,
                  7: 0, 8: 1, 9: ignore_label,
                  10: ignore_label, 11: 2, 12: 3,
                  13: 4, 14: ignore_label, 15: ignore_label,
                  16: ignore_label, 17: 5, 18: ignore_label,
                  19: 6, 20: 7, 21: 8, 22: 9, 23: 10, 24: 11,
                  25: 12, 26: 13, 27: 14, 28: 15,
                  29: ignore_label, 30: ignore_label,
                  31: 16, 32: 17, 33: 18}

def convert_label(label, inverse=False):
    temp = label.copy()
    if inverse:
        for v, k in label_mapping.items():
            label[temp == k] = v
    else:
        for k, v in label_mapping.items():
            label[temp == k] = v
    return label

def get_palette(n):
    palette = [0] * (n * 3)
    for j in range(0, n):
        lab = j
        palette[j * 3 + 0] = 0
        palette[j * 3 + 1] = 0
        palette[j * 3 + 2] = 0
        i = 0
        while lab:
            palette[j * 3 + 0] |= (((lab >> 0) & 1) << (7 - i))
            palette[j * 3 + 1] |= (((lab >> 1) & 1) << (7 - i))
            palette[j * 3 + 2] |= (((lab >> 2) & 1) << (7 - i))
            i += 1
            lab >>= 3
    return palette

def save_pred(preds, sv_path, name):
    palette = get_palette(256)
    preds = np.asarray(np.argmax(preds.cpu(), axis=1), dtype=np.uint8)
    for i in range(preds.shape[0]):
        pred = convert_label(preds[i], inverse=True)
        save_img = Image.fromarray(pred)
        save_img.putpalette(palette)
        save_img.save(os.path.join(sv_path, f'{name[i]}.png'))

函数 save_pred 的输入参数，sv_path为保存路径，name为保存图片的名称，preds为预测结果heatmap图，shape为（batchsize, classnumber, width, height）