Код: Выделить всё
import os
import ants
import nibabel as nib
import numpy as np
from nipy.algorithms.registration import HistogramRegistration
from scipy.ndimage import gaussian_filter
from skimage import exposure
import logging
# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Utility functions
def load_nifti(file_path):
if not os.path.exists(file_path):
logger.error(f"File not found: {file_path}")
raise FileNotFoundError(f"File not found: {file_path}")
return nib.load(file_path)
`your text`
def save_nifti(data, affine, header, output_path):
nib.save(nib.Nifti1Image(data, affine, header), output_path)
def n4_bias_correction(input_file, output_file):
try:
img = ants.image_read(input_file)
img_corrected = ants.n4_bias_field_correction(img)
ants.image_write(img_corrected, output_file)
logger.info(f'N4 bias correction completed for {input_file}')
except Exception as e:
logger.error(f'Error during N4 bias correction: {e}')
def brain_extraction(input_file, output_file):
try:
img = ants.image_read(input_file)
brain_mask = ants.get_mask(img)
img_brain = img * brain_mask
ants.image_write(img_brain, output_file)
logger.info(f'Brain extraction completed for {input_file}')
except Exception as e:
logger.error(f'Error during brain extraction: {e}')
def motion_correction(input_file, reference_file, output_file):
try:
input_img = load_nifti(input_file)
reference_img = load_nifti(reference_file)
reg = HistogramRegistration(input_img.get_fdata(), reference_img.get_fdata(), similarity="cc")
transformation = reg.optimize('rigid')
corrected_data = transformation.apply(input_img.get_fdata())
# Use the reference image's header and affine for consistency
save_nifti(corrected_data, reference_img.affine, reference_img.header, output_file)
logger.info(f'Motion correction completed for {input_file}')
except Exception as e:
logger.error(f'Error during motion correction: {e}')
def reduce_artifacts(data):
return gaussian_filter(data, sigma=1)
def enhance_contrast(data):
return exposure.equalize_adapthist(data, clip_limit=0.03)
def preprocess_subject(subject_dir, reference_file, output_dir):
if not os.path.exists(output_dir):
os.makedirs(output_dir)
input_file = os.path.join(subject_dir, 'anat', 'sub-01_acq-1_T1w.nii')
n4_output = os.path.join(output_dir, 'n4_bias_corrected.nii')
brain_extracted_output = os.path.join(output_dir, 'brain_extracted.nii')
motion_corrected_output = os.path.join(output_dir, 'motion_corrected.nii')
final_output = os.path.join(output_dir, 'final_preprocessed.nii')
n4_bias_correction(input_file, n4_output)
brain_extraction(n4_output, brain_extracted_output)
motion_correction(brain_extracted_output, reference_file, motion_corrected_output)
if not os.path.exists(motion_corrected_output):
raise FileNotFoundError(f"Motion correction output file not created: {motion_corrected_output}")
data = load_nifti(motion_corrected_output).get_fdata()
artifact_reduced_data = reduce_artifacts(data)
contrast_enhanced_data = enhance_contrast(artifact_reduced_data)
header = load_nifti(motion_corrected_output).header
affine = load_nifti(motion_corrected_output).affine
save_nifti(contrast_enhanced_data, affine, header, final_output)
def preprocess_dataset(dataset_dir, output_base_dir):
subjects = [d for d in os.listdir(dataset_dir) if os.path.isdir(os.path.join(dataset_dir, d)) and d.startswith('sub-')]
for subject in subjects:
subject_dir = os.path.join(dataset_dir, subject)
output_dir = os.path.join(output_base_dir, subject)
reference_file = os.path.join(subject_dir, 'anat', 'sub-01_acq-1_T1w.nii')
preprocess_subject(subject_dir, reference_file, output_dir)
# Set the dataset directory and output base directory
dataset_dir = r"C:\Users\arnav\Downloads\SKIP Dataset"
output_base_dir = r"C:\Users\arnav\Downloads"
preprocess_dataset(dataset_dir, output_base_dir)
Код: Выделить всё
INFO:__main__:N4 bias correction completed for C:\Users\arnav\Downloads\SKIP Dataset\sub-01\anat\sub-01_acq-1_T1w.nii
INFO:__main__:Brain extraction completed for C:\Users\arnav\Downloads\sub-01\n4_bias_corrected.nii
ERROR:__main__:Error during motion correction: 'memmap' object has no attribute 'get_affine'
Подробнее здесь: https://stackoverflow.com/questions/786 ... otion-corr
