normalize_degree_powers.py

import argparse
import json
import numpy as np
from common import get_degree_powers
from pathlib import Path
from dipy.io.image import load_nifti, save_nifti

# === Parse args ===

parser = argparse.ArgumentParser(description='normalize degree powers that were generated by compute_degree_powers')
parser.add_argument('degree_powers_path', type=str, help='path to folder to which compute_degree_powers wrote its output')
parser.add_argument('output_dir', type=str, help='path to folder in which to save the computed degree power images')
args = parser.parse_args()
degree_powers_path = Path(args.degree_powers_path)
output_dir = Path(args.output_dir)
l_values_in_path = degree_powers_path/'l_values.txt'
l_values_out_path = output_dir/'l_values.txt'
degree_power_image_inpath = degree_powers_path/'degree_power_images'
degree_power_image_outpath = output_dir/'degree_power_images'
degree_power_image_outpath.mkdir(exist_ok=True)

if not degree_power_image_inpath.exists():
    raise FileNotFoundError(f"Could not find {degree_power_image_inpath}")

with open(l_values_in_path, 'r') as l_values_file:
    l_values_in = np.array(json.load(l_values_file), dtype=int)

if 0 not in l_values_in:
    raise Exception(f"There is no l=0 in {l_values_in_path}.")

l_values_zero_index = np.where(l_values_in==0)[0][0]
l_values_nonzero_indices = np.where(l_values_in!=0)[0]

l_values_out = l_values_in[l_values_nonzero_indices]

with open(l_values_out_path, 'w') as l_values_file:
    json.dump(
        l_values_out.astype(int).tolist(), # list of ints suitable for json serialization
        l_values_file
    )

for dp_filepath in degree_power_image_inpath.glob('*.nii.gz'):

    basename = '_'.join(str(dp_filepath.name).split("_")[:-1]) # remove the trailing _degreepowers.nii.gz
    output_file_path = degree_power_image_outpath/f'{basename}_degreepowersnormalized.nii.gz'
    if output_file_path.exists():
        print(f"Skipping degree power computation for {basename} since the following output file exists:\n{output_file_path}")
        continue

    dp_data, affine, img = load_nifti(dp_filepath, return_img=True)

    dp_data_normalized = dp_data[...,l_values_nonzero_indices] / dp_data[...,[l_values_zero_index]]
    # If the l=0 power was 0 then the whole thing should be 0, but the division wouldn't make sense, so we just force zero:
    dp_data_normalized[ dp_data[...,l_values_zero_index] <= 0 ] = 0 

    save_nifti(output_file_path, dp_data_normalized, affine, img.header)
    print(f"normalized and saved degree powers for for {basename}...")
print("done")