nnunetv2系列：预测掩码可视化

实例分割的掩码输入的值，一般都得转成0， 1， 2这样的数值，预测出来后人无法直接可视化查看结果，这里给出代码示例，根据自定义的值进行转换，适合单通道灰度图和三通道彩色图。

代码示例

import cv2
import os


def mask_predict_recover(
    mask_predict_dir_path,
    mask_recover_dir_path,
    predict_recover_value_dict,
    gray_flag=True,
):
    os.makedirs(mask_recover_dir_path, exist_ok=True)
    mask_name_list = os.listdir(mask_predict_dir_path)
    for mask_name in mask_name_list:
        mask_path = os.path.join(mask_predict_dir_path, mask_name)
        if gray_flag:
            mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
        else:
            mask = cv2.imread(mask_path)
        for predict_value, recover_value in predict_recover_value_dict.items():
            mask[mask == predict_value] = recover_value
        mask_recover_path = os.path.join(mask_recover_dir_path, mask_name)
        cv2.imwrite(mask_recover_path, mask)


if __name__ == "__main__":
    mask_predict_dir_path = "./imagesTr_predict_test"
    mask_recover_dir_path = "./imagesTr_predict_recover_test"
    predict_recover_value_dict = {
        1: 128,
        2: 255,
    }
    mask_predict_recover(
        mask_predict_dir_path,
        mask_recover_dir_path,
        predict_recover_value_dict,
    )