dataprocesser/Preprocess_MR_Mask_generation.py

import os
import pandas as pd
import numpy as np
import nrrd
import SimpleITK as sitk
import cv2

from dataprocesser.preprocess_MR import step3_vibe_resetsignal
"""
该代码用于处理一组 MR 图像和对应的分割图，应用掩膜、进行归一化，并根据 CSV 文件中的仿真 MR 灰度值对分割图进行替换。最后将处理后的 MR 图像和分割图保存。

主要步骤：
1. **读取数据**：从指定的文件夹中读取 MR 图像和对应的分割图。
2. **归一化处理**：对 MR 图像进行归一化，将其值范围映射到 0 到 255 之间。
3. **轮廓提取**：从归一化后的 MR 图像中提取出主体区域的轮廓（根据给定的阈值分割），创建掩膜。
4. **掩膜应用**：将提取出的掩膜应用到归一化后的 MR 图像上，保留主体区域，抑制背景。
5. **分割图处理**：读取对应的分割图，并与提取出的轮廓进行叠加，之后根据 CSV 文件中的仿真 CT 值替换分割图中的灰度值。
6. **图像保存**：将处理后的 MR 图像和修改后的分割图保存到指定的输出文件夹中，保证其空间属性和几何信息与输入图像一致。
7. **输出**：在 ITK-SNAP 等医学图像工具中打开时, MR 图像和分割图能够保持同步和正确的比例显示。

函数简介：
- `normalize`: 对 MR 图像进行归一化处理，将像素值范围映射到 [0, 255]。
- `create_body_mask`: 从图像中提取出身体的轮廓，生成二值掩膜。
- `apply_mask`: 将提取的掩膜应用到 MR 图像上，保留轮廓内部的区域。
- `process_segmentation`: 读取分割图，并根据 CSV 文件中的仿真 CT 值对其灰度值进行替换。
- `process_image`: 处理单个 MR 图像及其对应的分割图，包括归一化、轮廓提取、掩膜应用、分割图处理等。
- `process_folder`: 处理整个文件夹中的 MR 图像和分割图，逐一处理所有图像并保存结果。
"""

# 归一化函数
def normalize(img, vmin_out=0, vmax_out=1, norm_min_v=None, norm_max_v=None, epsilon=1e-5):
    if norm_min_v is None and norm_max_v is None:
        norm_min_v = np.min(img)
        norm_max_v = np.max(img)
    img = np.clip(img, norm_min_v, norm_max_v)
    img = (img - norm_min_v) / (norm_max_v - norm_min_v + epsilon)
    img = img * (vmax_out - vmin_out) + vmin_out
    return img

# 创建轮廓掩膜
def create_body_mask_simple(numpy_img, body_threshold=50):
    numpy_img = numpy_img.astype(np.int16)
    body_mask = np.where(numpy_img > body_threshold, 1, 0).astype(np.uint8)
    contours, _ = cv2.findContours(body_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    mask = np.zeros_like(body_mask, dtype=np.uint8)

    if contours:
        largest_contour = max(contours, key=cv2.contourArea)
        mask = np.ascontiguousarray(mask)
        largest_contour = np.ascontiguousarray(largest_contour)
        cv2.drawContours(mask, [largest_contour], -1, 1, thickness=cv2.FILLED)

    return mask

def create_body_mask(numpy_img, body_threshold=-500, min_contour_area=10000):
    """
    Create a binary body mask from a CT image tensor, using a specific threshold for the body parts.

    Args:
    tensor_img (torch.Tensor): A tensor representation of a grayscale CT image, with intensity values from -1024 to 1500.

    Returns:
    torch.Tensor: A binary mask tensor where the entire body region is 1 and the background is 0.
    """
    # Convert tensor to numpy array
    numpy_img = np.ascontiguousarray(numpy_img.astype(np.int16))  # Ensure we can handle negative values correctly
    #numpy_img = numpy_img.astype(np.int16)

    # Threshold the image at -500 to separate potential body from the background
    binary_img = np.where(numpy_img > body_threshold, 1, 0).astype(np.uint8)

    # Find contours from the binary image
    contours, _ = cv2.findContours(binary_img, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)

    # Create an empty mask
    mask = np.zeros_like(binary_img)

    VERBOSE = False
    # Fill all detected body contours
    if contours:
        for contour in contours:
            if cv2.contourArea(contour) >= min_contour_area:
                if VERBOSE:
                    print('current contour area: ', cv2.contourArea(contour), 'threshold: ', min_contour_area)
                cv2.drawContours(mask, [contour], -1, 1, thickness=cv2.FILLED)

    return mask

def apply_mask(normalized_image_array, mask_array):
    return normalized_image_array * mask_array

def print_all_info(data, title):
    print(f'min, max of {title}:', np.min(data), np.max(data))

# process the segmentation, replace the classes with simulated MR values
def process_segmentation(combined_array, csv_simulation_values, mr_signal_formula=step3_vibe_resetsignal.calculate_signal_vibe):
    combined_array = combined_array.astype(np.int16)
    print_all_info(combined_array, 'combine')
    # two columns of  unique value 和 simulation value
    # the first element will not be included
    organ_indexs = csv_simulation_values[1:, 0]  # first column: organ index
    T1_values = csv_simulation_values[1:, 1]  # second column: simulate MRI value
    T2_values = csv_simulation_values[1:, 2]
    Rho_values = csv_simulation_values[1:, 3]
    order_begin_from_0 = True if organ_indexs.astype(int).min()==0 else False
    #print('organ order number begin from 0:', order_begin_from_0)
    #print(organ_indexs)
    assign_value_mask = np.zeros_like(combined_array)

    step=0
    for step in range(len(organ_indexs)):
        organ_index = organ_indexs[step] # in csv file, organs begin with 1       
        t1 = float(T1_values[step])
        t2 = float(T2_values[step])
        rho = float(Rho_values[step])

        simulation_value = mr_signal_formula(t1, t2, rho)
        organ_index = int(organ_index)
        if order_begin_from_0:
            #print("order in csv begin from 0")
            assign_value_mask[combined_array == organ_index+1] = simulation_value #  organ_index+ 1
        else:
            #print("order in csv begin from 1")
            assign_value_mask[combined_array == organ_index] = simulation_value
        step+=1
    print_all_info(assign_value_mask, 'assignment')
    return assign_value_mask

# 处理单个图像和分割图
def process_image(input_path, contour_path, seg_path, csv_simulation_values, output_path1, output_path2, body_threshold):
    # 读取原始 MR 图像和分割图
    if input_path.endswith('.nrrd'):
        img, header = nrrd.read(input_path)
        segmentation_img, header_seg = nrrd.read(seg_path)
    elif input_path.endswith('.nii.gz') or input_path.endswith('.nii'):
        import nibabel as nib
        img_metadata = nib.load(input_path)
        img = img_metadata.get_fdata()
        affine = img_metadata.affine
        
        seg_metadata = nib.load(seg_path)
        segmentation_img = seg_metadata.get_fdata()
    
    # 归一化处理
    norm_max=255 #255
    low_percentile = 5
    high_percentile = 90
    img_normalized = normalize(img, 0, norm_max, np.percentile(img, low_percentile), np.percentile(img, high_percentile), epsilon=0)
    
    # 提取轮廓图
    body_contour = np.zeros_like(img, dtype=np.int16)
    for i in range(img.shape[2]):
        slice_data = img[:, :, i]
        body_contour[:, :, i] = create_body_mask(slice_data, body_threshold=body_threshold)
    
    # 应用掩膜到归一化 MR 图像
    masked_image = apply_mask(img_normalized, body_contour)
    
    # 处理分割图
    # add contour background to the segmentation (all region inside body + 1)
    combined_array = segmentation_img + body_contour
    combined_array = np.clip(combined_array, 0, np.max(segmentation_img) + 1)
    print_all_info(segmentation_img, 'seg')
    processed_segmentation = process_segmentation(combined_array, csv_simulation_values)
    
    # normalize to 0-1
    # masked_image = masked_image/norm_max
    # processed_segmentation = processed_segmentation/norm_max

    if input_path.endswith('.nrrd'):
        # 保存处理后的 MR 图像
        nrrd.write(output_path1, masked_image, header)
        
        # 保存处理后的分割图
        nrrd.write(output_path2, processed_segmentation, header_seg)

        # save the body contour mask

    elif input_path.endswith('.nii.gz') or input_path.endswith('.nii'):
        img_processed = nib.Nifti1Image(masked_image, affine)
        nib.save(img_processed, output_path1)
        seg_processed = nib.Nifti1Image(processed_segmentation, affine)
        nib.save(seg_processed, output_path2)
        contour_processed = nib.Nifti1Image(body_contour, affine)
        
        # Split the path into directory and filename
        directory, filename = os.path.split(output_path2)
        new_filename = filename.replace('seg', 'contour')
        contour_path = os.path.join(directory, new_filename)

        nib.save(contour_processed, contour_path)
    return processed_segmentation

# 处理文件夹
def process_folder(input_folder1, input_folder2, output_folder1, output_folder2, csv_simulation_file, body_threshold=50):
    # 读取CSV文件获取仿真CT灰度值 (两列)
    csv_simulation_values = pd.read_csv(csv_simulation_file, header=None).to_numpy()

    # 检查 csv_simulation_values 是否是二维数组
    if csv_simulation_values.ndim == 1:
        raise ValueError("CSV 文件格式不正确，应该包含两列：organ_index 和 simulation_value")

    # 确保输出文件夹存在
    os.makedirs(output_folder1, exist_ok=True)
    os.makedirs(output_folder2, exist_ok=True)

    for filename in os.listdir(input_folder1):
        if filename.endswith('.nrrd'):
            input_file_path = os.path.join(input_folder1, filename)
            seg_file_path = os.path.join(input_folder2, filename)
            output_file_path1 = os.path.join(output_folder1, filename)
            output_file_path2 = os.path.join(output_folder2, filename)

            print(f"Processing {input_file_path} with segmentation {seg_file_path}")
            processed_seg = process_image(input_file_path, None, seg_file_path, csv_simulation_values, output_file_path1, output_file_path2, body_threshold)

def analyse_hist(input_path):
    if input_path.endswith('.nrrd'):
        img, header = nrrd.read(input_path)
    elif input_path.endswith('.nii.gz'):
        import nibabel as nib
        img_metadata = nib.load(input_path)
        img = img_metadata.get_fdata()
        affine = img_metadata.affine
    import numpy as np
    import matplotlib.pyplot as plt

    # Plot the histogram
    print('shape of img: ', img.shape)
    plt.hist(img[:, :, 50], bins=30, edgecolor='black', alpha=0.7)
    plt.xlabel('Value')
    plt.ylabel('Frequency')
    plt.title('Value Distribution')
    plt.show()


def process_csv(csv_file, output_folder1, output_folder2, csv_simulation_file, body_threshold=50, output_mr_csv_file='processed_mr_csv_file.csv'):
    # 读取CSV文件获取仿真CT灰度值 (两列)
    csv_simulation_values = pd.read_csv(csv_simulation_file, header=None).to_numpy()
    #csv_simulation_values = pd.read_csv(csv_simulation_file)

    # 检查 csv_simulation_values 是否是二维数组
    if csv_simulation_values.ndim == 1:
        raise ValueError("CSV 文件格式不正确，应该包含两列：organ_index 和 simulation_value")

    # 确保输出文件夹存在
    os.makedirs(output_folder1, exist_ok=True)
    os.makedirs(output_folder2, exist_ok=True)
    
    from step1_init_data_list import list_img_seg_ad_pIDs_from_new_simplified_csv
    patient_IDs, Aorta_diss, segs, images = list_img_seg_ad_pIDs_from_new_simplified_csv(csv_file)
    from tqdm import tqdm
    dataset_list = []
    for idx in tqdm(range(len(images))):
        if (images[idx].endswith('.nii.gz') and segs[idx].endswith('.nii.gz')) or \
            (images[idx].endswith('.nii') and segs[idx].endswith('.nii')):
            input_file_path = images[idx]
            seg_file_path = segs[idx]
            patient_id = patient_IDs[idx]
            ad = Aorta_diss[idx]
            root_dir = os.path.dirname(input_file_path)
            
            output_file_path1 = os.path.join(output_folder1, os.path.relpath(input_file_path, start=root_dir))
            
            synthrad_basic_mr_name = 'mr'
            synthrad_basic_seg_name = 'mr_merged_seg'
            if os.path.basename(output_file_path1) == f'{synthrad_basic_mr_name}.nii.gz' or \
                os.path.basename(output_file_path1) == f'{synthrad_basic_mr_name}.nii':
                # Insert the patient ID in the filename
                output_file_path1 = output_file_path1.replace(f'{synthrad_basic_mr_name}', f'mr_{patient_id}')

            output_file_path2 = os.path.join(output_folder2, os.path.relpath(seg_file_path, start=root_dir))
            
            if os.path.basename(output_file_path2) == f'{synthrad_basic_seg_name}.nii.gz' or \
                os.path.basename(output_file_path2) == f'{synthrad_basic_seg_name}.nii':
                # Insert the patient ID in the filename
                output_file_path2 = output_file_path2.replace(f'{synthrad_basic_seg_name}', f'mr_seg_{patient_id}')

            print(f"Processing {input_file_path} with segmentation {seg_file_path}")
            print(f"Save results to {output_file_path1} and {output_file_path2}")
            
            processed_seg = process_image(input_file_path, None, seg_file_path, csv_simulation_values, output_file_path1, output_file_path2, body_threshold)

            # processed_mr_csv_file = ...
            csv_mr_line = [patient_id,ad,output_file_path2,output_file_path1]
            dataset_list.append(csv_mr_line)

    import csv
    with open(output_mr_csv_file, 'w', newline='') as f:
        csvwriter = csv.writer(f)
        csvwriter.writerow(['id', 'Aorta_diss', 'seg', 'img'])
        csvwriter.writerows(dataset_list) 

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Process MR images and segmentation maps, apply masks and replace grayscale values.")
    parser.add_argument('--input_folder1', required=True, help="Path to the folder containing input MR .nrrd files.")
    parser.add_argument('--input_folder2', required=True, help="Path to the folder containing segmentation .nrrd files.")
    parser.add_argument('--output_folder1', required=True, help="Path to the folder to save the output MR files.")
    parser.add_argument('--output_folder2', required=True, help="Path to the folder to save the output segmentation files.")
    parser.add_argument('--csv_simulation_file', required=True, help="CSV file containing simulated CT grayscale values.")
    parser.add_argument('--body_threshold', type=int, default=50, help="Threshold to separate body from background.")
    args = parser.parse_args()

    process_folder(args.input_folder1, args.input_folder2, args.output_folder1, args.output_folder2, args.csv_simulation_file, args.body_threshold)