import SimpleITK as sitk
import os, glob
import json

data = {}
out_path = '/shared/temp_code/infer_res/img_size/g1_ori.json'
nii_dir = '/data/job_520/job_data_preprocess/output/preprocess/preprocess_file_details/nii'
for i, path in enumerate(glob.glob(nii_dir+'/*.nii.gz')):
    if i > 2:
        break
    uid = path.split('/')[-1].split('.nii.gz')[0]
    print(i, uid)
    itk_image = sitk.ReadImage(path)
    img = sitk.GetArrayFromImage(itk_image)
    z, y, x = img.shape
    data[uid] = [z, y, x]
with open(out_path, 'w+') as f:
    json.dump(data, f)