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,
)