torch,GPU，nii重采样，多线程，内存保护，检索对应的标签信息，根据分割标签信息重采样MRI

由于，MRI的nii文件shape和spacing与对应的segmentation的shape和spacing对应不上，同时每个segmentation的shape也都不一样，所以需要根据对应编号各自的shape和spacing进行重采样。

我这里是DCE-MRI文件dir下面有多个文件夹，每个文件夹下都是一个时间点的nii MRI文件。同时segmentation的只有一张，标识是数字一样，根据数字有不同的对应关系。同时有两种分割，器官和肿瘤，因此需要做检查。

input_dir = r"C:\breast_data\A" 输入文件夹
output_dir = r"D:\BaiduNetdiskDownload\Breast_data\Y_gt" 输出文件夹
seg_t_dir=r"D:\BaiduNetdiskDownload\Breast_data\Y_t_mask" 分割肿瘤
seg_o_dir=r"D:\BaiduNetdiskDownload\Breast_data\Y_o_mask" 分割器官
ysy<=50:控制内存大小，插值很依赖于CPU

import torch
import torch.nn.functional as F
import nibabel as nib
import numpy as np
import threading
import time
import psutil
import os
import subprocess
import numpy
import torch
from torch import nn
from d2l import torch as d2l
import re

def resample_image(image, new_shape, new_spacing):
    # Move tensors to the GPU device
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    tensor = torch.from_numpy(image.get_fdata()).unsqueeze(0).unsqueeze(0).float().to(device)
    # Calculate the resize factor
    current_shape = tensor.shape[-3:]
    resize_factor = [n / o for n, o in zip(new_shape, current_shape)]
    # Perform the resampling using interpolate function
    resampled_tensor = F.interpolate(
        tensor,
        size=new_shape,
        mode='trilinear',
        align_corners=False
    )
    # Create a new NIfTI image with updated data and spacing
    resampled_nifti = nib.Nifti1Image(resampled_tensor.squeeze().cpu().numpy(), affine=image.affine)
    # Update the affine matrix for the desired spacing
    current_spacing = image.header.get_zooms()[:3]
    scale_factor = [c / (r * rf) for c, r, rf in zip(current_spacing, new_spacing, resize_factor)]
    resampled_nifti.header.set_zooms(new_spacing + (image.header.get_zooms()[3:]))
    # Scale the affine matrix
    resampled_nifti.affine[:3, :3] = np.diag(scale_factor)
    return resampled_nifti

# Define the directory containing the folders of NIfTI files

input_dir = r"C:\breast_data\A"
output_dir = r"D:\BaiduNetdiskDownload\Breast_data\Y_gt"
seg_t_dir=r"D:\BaiduNetdiskDownload\Breast_data\Y_t_mask"
seg_o_dir=r"D:\BaiduNetdiskDownload\Breast_data\Y_o_mask"
os.makedirs(output_dir, exist_ok=True)

# Iterate over the folders in the input directory

def write_file(filename,folder_path,output_folder_path):
    nifti_path = os.path.join(folder_path, filename)
    number=extract_number_from_path(folder_path)
    image_breast_seg_path=os.path.join(seg_breast_dir, "Yunzhong_"+str(number)+".nii.gz")
    image_gt_seg_path=os.path.join(seg_gt_dir, "Yunzhong_"+str(number)+".nii.gz")

    image_MRI = nib.load(nifti_path)
    image_breast_seg_shape=nib.load(image_breast_seg_path).shape
    image_gt_seg_shape=nib.load(image_gt_seg_path).shape

    image_breast_seg_spacing=nib.load(image_breast_seg_path).header.get_zooms()
    image_gt_seg_spacing=nib.load(image_gt_seg_path).header.get_zooms()
    if image_breast_seg_shape==image_gt_seg_shape and image_breast_seg_spacing==image_gt_seg_spacing:
        resampled_image = resample_image(image_MRI, image_gt_seg_shape, image_gt_seg_spacing)
        output_filename = filename.replace(".nii.gz", "_resampled.nii.gz")
        output_path = os.path.join(output_folder_path, output_filename)
        nib.save(resampled_image, output_path)
        print("Resampled image saved:", output_path,image_gt_seg_shape,image_gt_seg_spacing)
    else:
        print(str(nifti_path)+"seg data error")


def extract_number_from_path(path):
    # Use regular expression to find the number in the path string
    match = re.search(r'\d+', path)

    if match:
        number = match.group()
        return int(number)
    else:
        return None

for foldername in os.listdir(input_dir):
    folder_path = os.path.join(input_dir, foldername)
    if os.path.isdir(folder_path):
        # Create a corresponding output folder
        output_folder_path = os.path.join(output_dir, foldername)
        os.makedirs(output_folder_path, exist_ok=True)

        # Iterate over the NIfTI files in the current folder
        for filename in os.listdir(folder_path):
            if filename.endswith(".nii") or filename.endswith(".nii.gz"):
                # Load the NIfTI image

                t=threading.Thread(name="writing_file:"+str(foldername)+str(filename),target=write_file,args=(filename,folder_path,output_folder_path))
                t.start()
                time.sleep(0.3)
                while True:
                    mem = psutil.virtual_memory()
                                # 系统总计内存
                    zj = float(mem.total) / 1024 / 1024 / 1024
                                # 系统已经使用内存
                    ysy = float(mem.used) / 1024 / 1024 / 1024
                                # 系统空闲内存
                    kx = float(mem.free) / 1024 / 1024 / 1024
                    #print('系统总计内存:%d.3GB' % zj)
                    #print('系统已经使用内存:%d.3 GB' % ysy)
                    #print('系统空闲内存:%d.3GB' % kx)
                    if ysy<=50:
                        #print("start threading:"+str(foldername)+str(filename))
                        break
                    else:
                        #print('waiting')
                        time.sleep(0.1)
                # Perform resampling
print("All images resampled and saved to:", output_dir)