Spaces:
Running
Running
File size: 6,135 Bytes
f5288df |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 |
from __future__ import generators
import logging
import glob, os, functools
import sys
sys.path.append('../')
import SimpleITK as sitk
import numpy as np
import scipy
import nibabel as nib
import skimage
import matplotlib.pyplot as plt
import scipy.misc
from scipy import ndimage
from skimage.transform import resize,rescale
import cv2
import itk
import subprocess
from tqdm import tqdm
import pandas as pd
import warnings
import statistics
import torch
import csv
import os
import yaml
from HD_BET.run import run_hd_bet # git clone HDBET repo
from dataset.preprocess_utils import enhance, enhance_noN4
from dataset.preprocess_datasets_T1_to_2d import create_quantile_from_brain
warnings.filterwarnings('ignore')
cuda_device = '1'
os.environ['CUDA_VISIBLE_DEVICES'] = cuda_device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#torch.cuda.set_device(1) # Set CUDA device to 0 (first GPU)
#net = net.cuda()
def select_template_based_on_age(age):
for golden_file_path, age_values in age_ranges.items():
if age_values['min_age'] <= int(age) and int(age) <= age_values['max_age']:
print(golden_file_path)
return golden_file_path
def register_to_template(input_image_path, output_path, fixed_image_path,rename_id,create_subfolder=True):
fixed_image = itk.imread(fixed_image_path, itk.F)
# Import Parameter Map
parameter_object = itk.ParameterObject.New()
parameter_object.AddParameterFile('/media/sdb/divyanshu/divyanshu/aidan_segmentation/pediatric-brain-age-main/dataset/golden_image/mni_templates/Parameters_Rigid.txt')
if "nii" in input_image_path and "._" not in input_image_path:
print(input_image_path)
# Call registration function
try:
moving_image = itk.imread(input_image_path, itk.F)
result_image, result_transform_parameters = itk.elastix_registration_method(
fixed_image, moving_image,
parameter_object=parameter_object,
log_to_console=False)
image_id = input_image_path.split("/")[-1]
itk.imwrite(result_image, output_path+"/"+rename_id+".nii.gz")
print("Registered ", rename_id)
except:
print("Cannot transform", rename_id)
def outlier_voting(numbers):
mean = statistics.mean(numbers)
stdev = statistics.stdev(numbers)
threshold = stdev # *2 #*3
good_nums_avg =[]
for n in numbers:
if n > mean + threshold or n < mean - threshold:
continue
else:
good_nums_avg.append(n)
#if len(good_nums_avg)<=3:
# print(len(good_nums_avg))
return np.average(good_nums_avg)
#Data https://openneuro.org/datasets/ds000228/versions/1.1.0
img_path = '/media/sdb/divyanshu/divyanshu/aidan_segmentation/dummy_t1_preprocess/OAS2_0001_MR1.nii.gz'
data_path = "/media/sdb/divyanshu/divyanshu/longitudinal_fm/datasets/abide/data"
gt_age = 86 # age of subject
gender = "M" # gender
path_to = "/media/sdb/divyanshu/divyanshu/longitudinal_fm/datasets/abide/preprocessed_data" # save to
# MNI templates http://nist.mni.mcgill.ca/pediatric-atlases-4-5-18-5y/
age_ranges = {"/media/data/BrainIAC/src/BrainIAC/golden_image/mni_templates/nihpd_asym_04.5-08.5_t1w.nii" : {"min_age":3, "max_age":7},
"/media/data/BrainIAC/src/BrainIAC/golden_image/mni_templates/nihpd_asym_07.5-13.5_t1w.nii": {"min_age":8, "max_age":13},
"/media/data/BrainIAC/src/BrainIAC/golden_image/mni_templates/nihpd_asym_13.0-18.5_t1w.nii": {"min_age":14, "max_age":100}}
for eachimage in tqdm(os.listdir(data_path), desc="Processing images", unit="image"):
if 1:#"sub" in eachimage:
## load image
img_path = os.path.join(data_path, eachimage)
nii= nib.load(img_path)
image, affine = nii.get_fdata(), nii.affine
#plt.imshow(image[:,:,100])
#print(nib.aff2axcodes(affine))
# path to store registered image in
new_path_to = path_to#path_to+"/"+img_path.split("/")[-1].split(".")[0]
if eachimage in os.listdir(new_path_to):
print("yay")
else:
if not os.path.exists(path_to):
os.mkdir(path_to)
if not os.path.exists(new_path_to):
os.mkdir(new_path_to)
# register image to MNI template
golden_file_path = select_template_based_on_age(gt_age)
print("Registering to template:", golden_file_path)
#fun fact: the registering to the template pipeline is not deterministic
register_to_template(img_path, new_path_to, golden_file_path,eachimage.split(".")[0]+"_"+"registered.nii.gz", create_subfolder=False)
# enchance and normalize image
#if not os.path.exists(new_path_to+"/no_z"):
# os.mkdir(new_path_to+"/no_z")
image_sitk = sitk.ReadImage(os.path.join(new_path_to, eachimage.split(".")[0]+"_"+"registered.nii.gz"))
image_array = sitk.GetArrayFromImage(image_sitk)
image_array = enhance(image_array) # or enhance_noN4(image_array) if no bias field correction is needed
image3 = sitk.GetImageFromArray(image_array)
sitk.WriteImage(image3,os.path.join(new_path_to, eachimage.split(".")[0]+"_"+"registered_no_z.nii.gz"))
#skull strip ## when running this with rest of the preprocessing, change the src path to include the registered image path!!!!
new_path_to = path_to
run_hd_bet(os.path.join(new_path_to, eachimage.split(".")[0]+"_"+"registered_no_z.nii.gz"),os.path.join(new_path_to, eachimage),
mode="accurate",
config_file='/media/sdb/divyanshu/divyanshu/aidan_segmentation/pediatric-brain-age-main/HD_BET/config.py',
device=device,
postprocess=False,
do_tta=True,
keep_mask=True,
overwrite=True)
|