import numpy as np
import SimpleITK as sitk
import os

files = ['ourdata/' + f for f in os.listdir('ourdata') if f.endswith('.nii.gz')]

print(files)
#0 intact 1 not intact
age = np.array([53.19917864,53.36071184,56.51471595,70.10540726,58.67761807,52.47091034,56.93908282,26.52156057,65.99589322,63,55.42231348,53.68925394,56.24914442,51.52635181,37.49760438,50.48323066,50.45311431,67.68514716,53.70020534,63.95893224,69.43463381,61.44010951,71.61396304,54.15195072,63.45242984,56.12046543,68.11225188,45.98494182,51.1430527,68.81314168,63.96167009,60.49281314,58.46132786,60.6954141,61.11978097,73.97399042,57.18001369,45.14989733,34.77891855,56.8733744,66.42299795,53.8973306,55.31827515,66.17385352,72.81587953,67.53730322,72.30663929,53.04312115,66.29979466,51.09103354,46,64,44,73,60,55,45,60,74])
#0 HC #1 MND
diagnosis = np.array([0,1,1,0,1,1,0,0,0,0,1,1,1,0,0,1,1,1,1,1,1,1,0,0,1,1,1,0,0,1,1,1,1,0,0,1,1,0,1,0,1,0,0,1,1,1,1,1,1,1,1,0,0,0,0,0,1,1,1])

for i, file in enumerate(files):
    simg = sitk.ReadImage(file)

    #read and crop
    npy_img = np.zeros([48,80,80])
    npy_img[0:47,0:61,0:66] = sitk.GetArrayFromImage(simg)
    out_file = file.replace('.nii.gz','.npz')

    np.savez_compressed(
        out_file,
        vol=npy_img,
        age=age[i],
        attribute=diagnosis[i],
    )