【超线程，重采样，插值法，内存监控】

超线程，重采样，插值法，内存监控

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


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"D:\BaiduNetdiskDownload\Breast_data\A"

# Define the output directory for the resampled files
output_dir = r"D:\BaiduNetdiskDownload\Breast_data\B"
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)
    image = nib.load(nifti_path)
    new_shape = (340, 340, 204)
    new_spacing = (1.0, 1.0, 1.0)
    resampled_image = resample_image(image, new_shape, new_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)



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

                # Define the new file path and name


                # Save the resampled image to the new file path




print("All images resampled and saved to:", output_dir)