import os
import argparse
import SimpleITK as sitk
import imageio

def vis_img(job_data_root):
    path = os.path.join(job_data_root, 'output/preprocess/qc')
    save_path = os.path.join(job_data_root, 'output/preprocess/qc_vis')
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    nii_imgs = os.listdir(path)

    for img in nii_imgs:
        nii_path = os.path.join(path, img)
        sitk_img = sitk.ReadImage(nii_path)
        np_img = sitk.GetArrayFromImage(sitk_img)

        num_layer = np_img.shape[0]
        start_index = num_layer // 2 - 5
        end_index = start_index + 10

        ten_slice = np_img[start_index:end_index,:,:]

        fps = 24.0
        imageio.mimsave(os.path.join(save_path, img[:-7]+'.gif'), ten_slice, 'GIF', duration=1/fps)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='PyTorch qc vis')
    parser.add_argument('--job_data_root', default='/data/job_715/job_data_preprocess',type=str)
    args = parser.parse_args()

    vis_img(args.job_data_root)