常用代码块

合并两个nii.gz

def merge_nii_gz(file1_path, file2_path, output_path):
    """
    将两个 .nii.gz 文件合并为一个,保持与原文件相同的尺寸,进行标签的并集操作,并保存为新的 .nii.gz 文件。

    :param file1_path: 第一个 .nii.gz 文件路径
    :param file2_path: 第二个 .nii.gz 文件路径
    :param output_path: 输出合并后的 .nii.gz 文件路径
    """
    # 加载第一个 .nii.gz 文件
    nii1 = nib.load(file1_path)
    data1 = nii1.get_fdata()

    # 加载第二个 .nii.gz 文件
    nii2 = nib.load(file2_path)
    data2 = nii2.get_fdata()

    # 检查两个文件的维度是否匹配
    if data1.shape != data2.shape:
        raise ValueError("两个 .nii.gz 文件的维度不一致,无法合并。")

    # 进行标签的并集操作 (按位逻辑或)
    merged_data = np.logical_or(data1, data2).astype(np.float32)

    # 使用第一个 NIfTI 文件的仿射矩阵和头信息
    merged_nii = nib.Nifti1Image(merged_data, nii1.affine, nii1.header)

    # 保存合并后的 NIfTI 图像为 .nii.gz 文件
    nib.save(merged_nii, output_path)
    print(f"标签并集后的 .nii.gz 文件已保存为 {output_path}")

Dicom转nii.gz

def dicom_to_nii_gz(dicom_dir, output_path):
    """
    使用 SimpleITK 将 DICOM 文件夹中的所有 DICOM 文件转换为 nii.gz 格式并保存,
    并保留原始 DICOM 的 spacing, direction 等信息。
    已处理中文路径可能的异常

    :param dicom_dir: 包含 DICOM 文件的文件夹路径
    :param output_path: 输出的 nii.gz 文件路径
    """
    # 读取 DICOM 文件夹中的 DICOM 图像
    reader = SimpleITK.ImageSeriesReader()
    dicom_series = reader.GetGDCMSeriesFileNames(dicom_dir)
    reader.SetFileNames(dicom_series)

    # 读取 DICOM 切片并组合成 3D 图像
    image = reader.Execute()

    # 获取原始的 spacing 和 direction
    spacing = image.GetSpacing()
    direction = image.GetDirection()
    origin = image.GetOrigin()

    print(f"Spacing: {spacing}")
    print(f"Direction: {direction}")
    print(f"Origin: {origin}")

    # 改变当前工作目录为输出文件夹所在目录
    if not os.path.exists("\\".join(output_path.split("\\")[:-1])):
        os.makedirs("\\".join(output_path.split("\\")[:-1]))

    os.chdir("\\".join(output_path.split("\\")[:-1]))

    # 将图像转换为 NIfTI 格式并保存为 .nii.gz
    output_file = output_path.split("\\")[-1]
    SimpleITK.WriteImage(image, output_file)

    print(f"DICOM 文件已成功转换并保存为 {output_path}")

按比例缩小bbox

def shrink_bbox_in_mask(mask_path, output_path, shrink_percentage):
    """
    函数功能:缩小给定nii.gz掩码文件中的边界框(bbox)
    
    参数:
    mask_path (str): 输入nii.gz掩码文件的路径
    output_path (str): 保存修改后掩码文件的路径
    shrink_percentage (float): 缩小边界框的百分比,例如 0.1 表示缩小10%
    
    返回:
    无(将处理结果保存为新的nii.gz文件)
    """
    
    # 第1步: 加载nii.gz掩码文件
    img = nib.load(mask_path)
    mask_data = img.get_fdata()
    
    # 第2步: 获取掩码中非零区域的bounding box
    mask_nonzero = np.argwhere(mask_data > 0)
    min_coord = mask_nonzero.min(axis=0)  # 获取非零体素的最小坐标
    max_coord = mask_nonzero.max(axis=0)  # 获取非零体素的最大坐标
    
    # 第3步: 计算bbox的中心和尺寸
    center = (min_coord + max_coord) / 2  # bbox中心坐标
    bbox_size = max_coord - min_coord      # bbox的尺寸
    
    # 第4步: 根据给定的百分比缩小bbox
    shrink_factor = (1 - shrink_percentage)  # 缩小因子
    new_bbox_size = bbox_size * shrink_factor  # 缩小后的bbox尺寸
    new_min_coord = np.round(center - new_bbox_size / 2).astype(int)  # 新的最小坐标
    new_max_coord = np.round(center + new_bbox_size / 2).astype(int)  # 新的最大坐标
    
    # 确保新的bbox在原始mask的维度范围内
    new_min_coord = np.maximum(new_min_coord, 0)
    new_max_coord = np.minimum(new_max_coord, mask_data.shape)
    
    # 第5步: 创建包含缩小后的bbox的新掩码
    new_mask_data = np.zeros_like(mask_data)  # 初始化新的掩码
    new_mask_data[new_min_coord[0]:new_max_coord[0],
                  new_min_coord[1]:new_max_coord[1],
                  new_min_coord[2]:new_max_coord[2]] = mask_data[
                      new_min_coord[0]:new_max_coord[0],
                      new_min_coord[1]:new_max_coord[1],
                      new_min_coord[2]:new_max_coord[2]]  # 填充缩小后的bbox区域
    
    # 第6步: 将新的掩码保存为nii.gz文件
    new_img = nib.Nifti1Image(new_mask_data, img.affine)
    nib.save(new_img, output_path)
    
    print(f'缩小后的掩码已保存到 {output_path}')

3维前缀和

def prefix_sum_3d(arr):
    # 创建一个和原数组大小相同的数组来存储前缀和
    prefix_sum = np.zeros_like(arr)

    # 使用三重循环计算每个位置的前缀和
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            for k in range(arr.shape[2]):
                # 计算前缀和
                prefix_sum[i, j, k] = arr[i, j, k]
                
                if i > 0:
                    prefix_sum[i, j, k] += prefix_sum[i - 1, j, k]
                if j > 0:
                    prefix_sum[i, j, k] += prefix_sum[i, j - 1, k]
                if k > 0:
                    prefix_sum[i, j, k] += prefix_sum[i, j, k - 1]
                if i > 0 and j > 0:
                    prefix_sum[i, j, k] -= prefix_sum[i - 1, j - 1, k]
                if i > 0 and k > 0:
                    prefix_sum[i, j, k] -= prefix_sum[i - 1, j, k - 1]
                if j > 0 and k > 0:
                    prefix_sum[i, j, k] -= prefix_sum[i, j - 1, k - 1]
                if i > 0 and j > 0 and k > 0:
                    prefix_sum[i, j, k] += prefix_sum[i - 1, j - 1, k - 1]

    return prefix_sum

截取mask中有值的部分与图像的对应部分

def crop_image_and_mask(image_path, mask_path, output_image_path, output_mask_path):
    # 加载图像和掩膜数据
    image_nii = nib.load(image_path)
    mask_nii = nib.load(mask_path)
    
    # 获取图像和掩膜的 numpy 数组
    image_data = image_nii.get_fdata()
    mask_data = mask_nii.get_fdata()
    
    # 获取掩膜中非零值的坐标
    mask_nonzero_coords = np.argwhere(mask_data > 0)
    
    # 找到掩膜的最小和最大坐标
    min_coords = mask_nonzero_coords.min(axis=0)
    max_coords = mask_nonzero_coords.max(axis=0)
    
    # 截取图像和掩膜中对应的部分
    cropped_image_data = image_data[min_coords[0]:max_coords[0]+1, 
                                    min_coords[1]:max_coords[1]+1, 
                                    min_coords[2]:max_coords[2]+1]
    cropped_mask_data = mask_data[min_coords[0]:max_coords[0]+1, 
                                  min_coords[1]:max_coords[1]+1, 
                                  min_coords[2]:max_coords[2]+1]
    
    # 创建新的 nifti 图像
    cropped_image_nii = nib.Nifti1Image(cropped_image_data, image_nii.affine)
    cropped_mask_nii = nib.Nifti1Image(cropped_mask_data, mask_nii.affine)
    
    # 保存截取后的图像和掩膜
    nib.save(cropped_image_nii, output_image_path)
    nib.save(cropped_mask_nii, output_mask_path)

除mask不为零的位置,其他位置置为该nii图像的最小值

def apply_mask_and_replace(nii_path, mask_path):
    # 读取 nii 图像和对应的 mask
    nii_img = nib.load(nii_path)
    mask_img = nib.load(mask_path)
    
    nii_data = nii_img.get_fdata()  # 获取nii图像的数据
    mask_data = mask_img.get_fdata()  # 获取mask的数据
    
    # 找到nii图像的最小值
    min_value = np.min(nii_data)
    
    # 创建一个新的图像数据,复制原始图像数据
    modified_data = nii_data.copy()
    
    # 将mask外的区域设置为最小值
    modified_data[mask_data == 0] = min_value
    
    # 创建新的nii图像
    modified_nii = nib.Nifti1Image(modified_data, nii_img.affine)
    
    return modified_nii

常用模型

Swin

预训练模型:

from functools import partial
from collections import OrderedDict
import torch
import torch.nn as nn
from monai.networks.nets.swin_unetr import SwinTransformer, SwinUNETR
from monai.utils import ensure_tuple_rep
import torch.nn.functional as F

from modules.VisionTransformer import VisionTransformer

class Classifier(nn.Module):
    def __init__(self, input_dim, num_classes):
        super(Classifier, self).__init__()
        # 使用全局平均池化
        self.global_pool = nn.AdaptiveAvgPool3d(1)  # 输出维度为 (batch_size, in_channels, 1, 1, 1)

        self.fc1 = nn.Linear(input_dim, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, num_classes)  # 最后一层是分类层

    def forward(self, x):
        x = self.global_pool(x)  # 输出形状:[batch_size, 384, 1, 1, 1]
        x = x.view(x.size(0), -1)  # 输出形状:[batch_size, 384]
        x = F.leaky_relu(self.fc1(x))
        x = F.leaky_relu(self.fc2(x))
        x = self.fc3(x)
        return x

class swinTrans_model(nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.model = SwinUNETR(
            img_size = [64, 64, 64],
            in_channels = 1,
            out_channels = 14,
            feature_size = 48,
        )
        self.model.load_state_dict(torch.load("./models/pretrain/monai_swin.pt"))
        self.classifier = Classifier(input_dim=768, num_classes=2)

    def forward(self, x, classify = True):
        hidden_states_out = self.model.swinViT(x, True)
        if classify:
            return self.classifier(hidden_states_out[4])
        return hidden_states_out[4]

def SwinTrans():
    model = swinTrans_model()
    return model

发表评论