import os
import numpy as np
import nibabel as nib
from util import fileutil, imgutil

if __name__ == '__main__':
    in_dir = '/home/allen/Temp/PlanQA/test'
    out_dir = '/home/allen/Temp/PlanQA/test-normalized'
    normalizer = imgutil.dispersed(min=0, max=60)

    dirs = fileutil.list_dir(in_dir)
    for d in dirs:
        img_path = f'{d}/ct.nii.gz'
        print(img_path)
        nii = nib.load(img_path)
        img_nii_data = np.array(nii.dataobj).astype(np.float32)
        img_nii_data = normalizer(img_nii_data)

        mask_nii = nib.Nifti1Image(img_nii_data, nii.affine)
        out_mask_path = img_path.replace(in_dir, out_dir)
        fileutil.check_and_create_dir_for_new_file(out_mask_path)
        nib.save(mask_nii, out_mask_path)
