Spaces:
Running
Running
| 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) | |