revert_normalisation.py

import sys
import SimpleITK as sitk
import json
import glob
import os
from tqdm import tqdm
import numpy as np
import torch


# revert normalisation
def get_ct_normalisation_values(ct_plan_path):
    """
    Get the mean and standard deviation for CT normalisation.
    """
    # Load the nnUNet plans file for CT
    with open(ct_plan_path, "r") as f:
        ct_plan = json.load(f)

    ct_mean = ct_plan['foreground_intensity_properties_per_channel']["0"]['mean']
    ct_std = ct_plan['foreground_intensity_properties_per_channel']["0"]['std']
    print(f"CT mean: {ct_mean}, CT std: {ct_std}")
    return ct_mean, ct_std

def revert_normalisation(pred_path, ct_mean, ct_std, save_path=None, mask_path=None, mask_outside_value=-1000):
    """
    Revert the normalisation of a CT image.
    """
    if save_path is None:
        save_path = pred_path + '_revert_norm'
    os.makedirs(save_path, exist_ok=True)
    imgs = glob.glob(os.path.join(pred_path, "*.nii.gz")) + \
        glob.glob(os.path.join(pred_path, "*.mha"))

    if mask_path:
        print(f"Applying mask from {mask_path} with outside value {mask_outside_value}")
    else:
        print("No mask provided, normalisation will be applied to all images.")
    for img in tqdm(imgs):
        img_sitk = sitk.ReadImage(img)
        img_array = sitk.GetArrayFromImage(img_sitk)
        img_array = img_array * ct_std + ct_mean
        img_sitk_reverted = sitk.GetImageFromArray(img_array)
        img_sitk_reverted.CopyInformation(img_sitk)

        # if mask_path is provided, apply the mask
        if mask_path:
            filename = os.path.basename(img)
            filename = filename.replace('_0000', '') if '_0000' in filename else filename
            mask_itk = sitk.ReadImage(os.path.join(mask_path, filename))
            img_sitk_reverted = sitk.Mask(img_sitk_reverted, mask_itk, outsideValue=mask_outside_value)
        sitk.WriteImage(img_sitk_reverted, os.path.join(save_path, os.path.basename(img)))
        # print(f"Reverted saved to {os.path.join(save_path, os.path.basename(img))}")
import SimpleITK as sitk
import numpy as np

def print_sitk_space(img: sitk.Image, name: str = "img"):
    if not isinstance(img, sitk.Image):
        print(f"[{name}] 不是 SimpleITK.Image（得到 {type(img)}），没有空间信息可打印。")
        return
    size = img.GetSize()            # (x, y, z)
    spacing = img.GetSpacing()      # (x, y, z)
    origin = img.GetOrigin()        # (x, y, z)
    direction = np.array(img.GetDirection())
    dim = img.GetDimension()
    if direction.size == dim*dim:
        direction = direction.reshape(dim, dim)

    print(f"[{name}] size (x,y,z)     = {size}")
    print(f"[{name}] spacing (x,y,z)  = {spacing}")
    print(f"[{name}] origin  (x,y,z)  = {origin}")
    print(f"[{name}] direction matrix =\n{direction}")
    print(f"[{name}] pixel type       = {img.GetPixelIDTypeAsString()}")

def revert_normalisation_modified(pred_path, ct_mean, ct_std, save_path=None,
                         mask_path=None, mask_sitk=None, mask_outside_value=-1000):
    if save_path is None:
        save_path = pred_path + '_revert_norm'
    os.makedirs(save_path, exist_ok=True)

    imgs = glob.glob(os.path.join(pred_path, "*.nii.gz")) + \
           glob.glob(os.path.join(pred_path, "*.mha"))

    if mask_path:
        print(f"Applying mask from {mask_path} with outside value {mask_outside_value}")
    elif mask_sitk is not None:
        print(f"Applying provided mask_sitk with outside value {mask_outside_value}")
    else:
        print("No mask provided, normalisation will be applied to all images.")

    for img in tqdm(imgs):
        img_sitk = sitk.ReadImage(img)
        img_array = sitk.GetArrayFromImage(img_sitk)
        img_array = img_array * ct_std + ct_mean
        img_sitk_reverted = sitk.GetImageFromArray(img_array)
        img_sitk_reverted.CopyInformation(img_sitk)

        if mask_path:
            filename = os.path.basename(img)
            filename = filename.replace('_0000', '') if '_0000' in filename else filename
            mask_itk = sitk.ReadImage(os.path.join(mask_path, filename))
            img_sitk_reverted = sitk.Mask(img_sitk_reverted, mask_itk, outsideValue=mask_outside_value)
        elif mask_sitk is not None:
            img_sitk_reverted = sitk.Mask(img_sitk_reverted, mask_sitk, outsideValue=mask_outside_value)

        sitk.WriteImage(img_sitk_reverted, os.path.join(save_path, os.path.basename(img)))
def revert_normalisation_single_modified(pred_sitk, ct_mean, ct_std, mask_sitk=None, mr_sitk = None,outside_value=-1000):
    print(type(pred_sitk))
    # print()
    # arr = sitk.GetArrayFromImage(pred_sitk).astype(np.float32)
    # print(arr)
    arr = pred_sitk * float(ct_std) + float(ct_mean)

    # out = sitk.GetImageFromArray(arr)
    # out.CopyInformation(mr_sitk)

    if mask_sitk is not None:
        out = sitk.Mask(arr, mask_sitk, outsideValue=outside_value)

    return out
# def revert_normalisation_single_modified(pred_sitk, ct_mean, ct_std, mask_sitk=None, mr_sitk=None, outside_value=-1000):
#     import SimpleITK as sitk
#     import numpy as np

#     print_sitk_space(pred_sitk, "pred_sitk (in)")  # 打印传入影像的空间信息

#     arr = sitk.GetArrayFromImage(pred_sitk).astype(np.float32)  # (z, y, x)
#     arr = arr * float(ct_std) + float(ct_mean)

#     out = sitk.GetImageFromArray(arr)   # 这里生成的新图默认 spacing=(1,1,1), origin=(0,0,0), direction=I

#     # 用参考影像复制空间信息：优先用 mr_sitk（如果你希望与原始 MR 对齐）
#     ref = mr_sitk if mr_sitk is not None else pred_sitk
#     out.CopyInformation(ref)

#     print_sitk_space(out, "out (after CopyInformation)")  # 打印复制后的空间信息

#     if mask_sitk is not None:
#         # 如果 out 和 mask 的网格不完全一致，可以先重采样到 mask 的网格
#         if (out.GetSize()!=mask_sitk.GetSize() or
#             out.GetSpacing()!=mask_sitk.GetSpacing() or
#             out.GetOrigin()!=mask_sitk.GetOrigin() or
#             out.GetDirection()!=mask_sitk.GetDirection()):
#             out = sitk.Resample(out, mask_sitk, sitk.Transform(), sitk.sitkLinear, outside_value, out.GetPixelID())
#         out = sitk.Mask(out, sitk.Cast(mask_sitk, sitk.sitkUInt8), outsideValue=outside_value)

#     return out

def revert_normalisation_single(pred_sitk, ct_mean, ct_std):
    arr = sitk.GetArrayFromImage(pred_sitk)
    arr = arr * ct_std + ct_mean
    reverted = sitk.GetImageFromArray(arr)
    reverted.CopyInformation(pred_sitk)
    return reverted
if __name__ == "__main__":
    ct_plan_path = "/datasets/work/hb-synthrad2023/work/synthrad2025/bw_workplace/data/nnunet_struct/preprocessed/Dataset203_synthrad2025_task1_CT/nnUNetPlans.json"
    ct_mean, ct_std = get_ct_normalisation_values(ct_plan_path)
    pred_path = "/datasets/work/hb-synthrad2023/work/synthrad2025/bw_workplace/data/nnunet_struct/results/Dataset202_synthrad2025_task1_MR_mask/nnUNetTrainerMRCT__nnUNetPlans__3d_fullres/fold_0/validation"
    revert_normalisation(pred_path, ct_mean, ct_std, save_path=pred_path + "_revert_norm")