# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: env-atlases
#     language: python
#     name: env-atlases
# ---

# +
import os
import json

import nibabel as nib
import numpy as np

np.warnings.filterwarnings('ignore')
# -

# ## Converting the unique values in atlases and mapping to a new space. 
#
# There are two folders in `ndmg_atlases.zip` that needs to be converted. First is in `./labels` folder and the other is in `./func_atlases/label`. Within each folder, there will be a json folder that maps the old space to new space in the form of `"new_value": old_value`. Values in new space is the key since JSON does not allow non-string keys, while values can be numberics.

# #### Below maps files in `./labels/`

# +
path = './label/'
files = os.listdir(path)

# Make folder for JSON files
os.makedirs(path + 'json/', exist_ok=True)
json_path = path + 'json/'

for file in files:
    if 'nii' in file:
        parcellation_name = file.split('.nii')[0]
        tmp = nib.load(path + file)
        
        uniques = np.unique(tmp.get_data())
        dtype = uniques.dtype

        print('Creating JSON map for {}. Has dtype {}.'.format(parcellation_name, dtype))
        
        # Convert to new values. 
        json_dump = []
        for new_val, old_val in enumerate(uniques):
            if np.issubdtype(dtype, np.float):
                json_dump.append((str(new_val), float(old_val)))
            else:
                json_dump.append((str(new_val), int(old_val)))
        
        # Write to json file
        json_dump = dict(json_dump)
        with open(json_path + '{}.json'.format(parcellation_name), 'w') as f:
            json.dump(json_dump, f, indent = 4, )


# -

# ## Rewrite all niftis using the new space
# All nifti images are in uint32.

# #### Rewriting files in `./labels`

def update_nifti(path, json_path, updated_path):
    files = os.listdir(path)
    
    for file in files:
        if 'nii.gz' in file:
            parcellation_name = file.split('.nii')[0]

            print("Updating {} atlas.".format(parcellation_name))

            img = nib.load(path + file)
            dat = img.get_data() # Image in current nifti

            # Open JSON map
            with open(json_path + '{}.json'.format(parcellation_name)) as f:
                newdict = json.load(f)

            # Check if all old_val == new_val. In this case, we just cast to uint32.
            compare = np.all(np.asarray(list(newdict.keys())).astype(np.float) == np.asarray(list(newdict.values())))

            if compare:
                output = dat.astype(np.uint32)
            else:
                output = np.zeros(dat.shape, dtype = np.uint32) # Image for output

                # Convert to new values
                for new_val, old_val in newdict.items():
                    #idx = np.where(dat == old_val)
                    #np.put(output, idx, int(new_val))
                    np.putmask(output, dat == old_val, int(new_val))
            
            img.header['datatype'] = 768

            # Save nifti
            nifti = nib.Nifti1Image(dataobj = output, affine = img.affine, header = img.header)
            nib.save(nifti, updated_path + file)


# +
path = './label/'
json_path = path + 'json/'
updated_path = path + 'updated/'

# Make folder for updated niftis
os.makedirs(updated_path, exist_ok=True)

update_nifti(path=path, json_path=json_path, updated_path=updated_path)
# -

# ## Checking the updated niftis
#
# Compare the counts of each value from old nifti to new nifti

# +
from collections import Counter
from numpy.testing import assert_array_equal

def get_counts(arr):
    c = Counter(arr.ravel())
    return np.array([v for k,v in sorted(c.items())])

def make_comparisons(old_path, updated_path):
    old_files = [x for x in os.listdir(old_path) if 'nii' in x]
    updated_files = [x for x in os.listdir(updated_path) if 'nii' in x]
    
    assert len(old_files) == len(updated_files)
    
    for file in old_files:
        if 'nii' in file:
            parcellation_name = file.split('.nii')[0]

            old_file = nib.load(old_path + file).get_data()
            updated_file = nib.load(updated_path + file).get_data()
            
            old_counts = get_counts(old_file)
            updated_counts = get_counts(updated_file)
            
            assert_array_equal(old_counts, updated_counts)
            print("{} identical.".format(parcellation_name))


# -

path = './label/'
json_path = path + 'json/'
updated_path = path + 'updated/'

make_comparisons(path, updated_path)

# #### Just look at the first 10 unique values to see if there are any floats

files = os.listdir(updated_path)
for file in files:
    if 'nii' in file:
        parcellation_name = file.split('.nii')[0]
        tmp = nib.load(updated_path + file).get_data()
        uniques = np.unique(tmp)
        print(parcellation_name, uniques[:10])



StarcoderdataJupyterScriptsDedupFiltered