| from ddmr.utils.nifti_utils import save_nifti | |
| from tqdm import tqdm | |
| import os | |
| import h5py | |
| import ddmr.utils.constants as C | |
| DATASET_LOCATION = '/mnt/EncryptedData1/Users/javier/vessel_registration/3Dirca/dataset/EVAL' | |
| DATASET_NAMES = ['Affine', 'None', 'Translation'] | |
| DATASET_FILENAME = 'volume' | |
| if __name__ == '__main__': | |
| for dataset_name in DATASET_NAMES: | |
| dataset_loc = os.path.join(DATASET_LOCATION, dataset_name) | |
| dataset_files = os.listdir(dataset_loc) | |
| dataset_files.sort() | |
| dataset_files = [os.path.join(dataset_loc, f) for f in dataset_files if DATASET_FILENAME in f] | |
| iterator = tqdm(dataset_files) | |
| for fn in iterator: | |
| f = os.path.split(fn)[-1].split('.hd5')[0] | |
| vol_file = h5py.File(fn, 'r') | |
| fix_vessels = vol_file[C.H5_FIX_VESSELS_MASK][..., 0] | |
| mov_vessels = vol_file[C.H5_MOV_VESSELS_MASK][..., 0] | |
| dst_folder = os.path.join(os.getcwd(), 'VESSELS', dataset_name) | |
| os.makedirs(dst_folder, exist_ok=True) | |
| save_nifti(fix_vessels, os.path.join(dst_folder, f+'_fix.nii.gz')) | |
| save_nifti(mov_vessels, os.path.join(dst_folder, f+'_mov.nii.gz')) | |
| vol_file.close() | |