text: from PIL import Image
  from utility import *
  import skimage.morphology
  import ants
  from tqdm import tqdm
  from glob import glob
  
  
  # Program description:
  # This script preprocess the CT and MRI pair for bone synthesis project.
  
  
  def check_resolution(root_dir: str = '/mnt/sdb2/Backup/Downloads/CT_MRI'):
      ct_dict, mr_dict = {}, {}
      for pid in tqdm(sorted(os.listdir(root_dir))):
          folder = os.path.join(root_dir, pid)
          file_fps = glob(folder + '/*.*')
          ct_fp, mr_fp = None, None
          for file_fp in file_fps:
              if os.path.basename(file_fp)[5:7] == 'CT':
                  ct_fp = file_fp
              elif os.path.basename(file_fp)[5:7] == 'T1':
                  mr_fp = file_fp
          ct_zoom = get_resolution(ct_fp)
          mr_zoom = get_resolution(mr_fp)
          if ct_zoom not in ct_dict:
              ct_dict[ct_zoom] = 1
          else:
              ct_dict[ct_zoom] += 1
          if mr_zoom not in mr_dict:
              mr_dict[mr_zoom] = 1
          else:
              mr_dict[mr_zoom] += 1
      print('*****************')
      print(ct_dict)
      print('*****************')
      print(mr_dict)
  
  
  def register_ct_mr(root_dir: str = '/mnt/sdb2/Backup/Downloads/CT_MRI',
                     dir: str = 'mr_to_ct'):
      """Rigid registration for paired CT and MR
  
      Parameters:
        root_dir:
          The directory that contains the paired CT and MR images.
        dir:
          The direction of rigid registration. Options: ['mr_to_ct', 'ct_to_mr'].
          Defaults to 'mr_to_ct'.
      """
      assert dir in ['mr_to_ct', 'ct_to_mr'], "Error: dir must be one of 'mr_to_ct' or 'ct_to_mr'."
      for pid in tqdm(sorted(os.listdir(root_dir))):
          folder = os.path.join(root_dir, pid)
          file_fps = glob(folder + '/*.*')
          ct, mr = None, None
          for file_fp in file_fps:
              if os.path.basename(file_fp)[5:7] == 'CT':
                  ct = ants.image_read(file_fp)
              elif os.path.basename(file_fp)[5:7] == 'T1':
                  mr = ants.image_read(file_fp)
          if dir == 'ct_to_mr':
              fixed = mr
              moving = ct
              name_tag = '_reg_CT'
          else:
              fixed = ct
              moving = mr
              name_tag = '_reg_MR'
  
          mytx = ants.registration(fixed=fixed, moving=moving, type_of_transform='Rigid')
          warped = ants.apply_transforms(fixed=fixed, moving=moving, transformlist=mytx['fwdtransforms'])
          ants.image_write(warped, os.path.join(folder, os.path.basename(pid)[:4] + name_tag + '.nii'), ri=True)
  
  
  def preprocess(ct_path: str, mr_path: str, T_ct: int = 400, radius: int = 4) -> list:
      """Apply preprocessing to paired of CT and MRI images
         and save the processed volumes in current folder.
  
      Parameters:
        ct_path:
          The filepath of CT image.
        mr_path:
          The filepath of MR image.
        T_ct:
          The general thresholding value for CT bone extraction. Defaults to 400.
        radius:
          The radius of the ball structure for morphological dilation.
  
      Returns:
          A list of processed volumes: [ct, mr]
      """
      ct = read_img(ct_path)[0]
      mr = read_img(mr_path)[0]
  
      # 1. determine the ROI of CT and apply it to MR
      mr[ct == 0] = 0
  
      # 2. apply a general thresholding for CT bone extraction
      ct_msk = ct.copy()
      ct_msk[ct_msk > T_ct] = 1
      ct_msk[ct_msk != 1] = 0
  
      # 3. obtain the largest connected component (remove isolated noise)
      ct_msk = get_cc3d(ct_msk, top=1)
  
      # 4. apply morphological dilation with a ball structure
      # (keep the content between bone plates)
      struct = skimage.morphology.ball(radius=radius)
      ct_msk = skimage.morphology.binary_dilation(ct_msk, struct)
      ct *= ct_msk
  
      return [ct, mr]
  
  
  def generate_25d_slices(data_dir: str, save_dir: str, load_size: int = 512, n_slices: int = 5):
      """ Given the dataset of paired CT and MR, generate 2D slices into pix2pix format.
  
      Parameters:
          data_dir:
              The dataset folder which includes sub-folders for each patient CT and MR.
          save_dir:
              Saving directory
          load_size:
              The dimension of generated slices. Defaults to 512.
          n_slices:
              Number of slices to generate 2.5D data
      """
      # create saving directories if they did not exist.
      if not os.path.exists(save_dir):
          os.mkdir(save_dir)
  
      h = n_slices // 2
  
      for id_dir in tqdm(sorted(glob(data_dir + '/*'))[:66]):
          pid = os.path.basename(id_dir)
          ct_path = os.path.join(data_dir, pid, pid + '_CT' + '.nii')
          mr_path = os.path.join(data_dir, pid, pid + '_reg_MR' + '.nii')
          [ct, mr] = preprocess(ct_path, mr_path)
  
          assert ct.shape == mr.shape, \
              'Error found: registered CT has a different dimension as MR'
  
          if ct.shape[0] != load_size or ct.shape[1] != load_size:
              print('Reshape volume...')
              ct = resize(ct, (load_size, load_size, ct.shape[2]), order=3)
              mr = resize(mr, (load_size, load_size, ct.shape[2]), order=3)
  
  
          for slc in range(h, ct.shape[2] - h):
              # rotate and flip to adjust slice orientation
              ct_slcs = ct[:, :, slc - h:slc + h + 1]
              mr_slcs = mr[:, :, slc - h:slc + h + 1]
  
              # skip the empty slices (no foreground)
              if np.max(ct_slcs) == 0:
                  continue
  
              np.savez_compressed(os.path.join(save_dir, f'{pid}_{slc}'), mr=mr_slcs, ct=ct_slcs)
  
  
  if __name__ == '__main__':
  
      # A: MRI B:CT
      data_dir = '/mnt/sdb2/Backup/Downloads/CT_MRI'
      # check_resolution(root_dir=data_dir)
      # register_ct_mr(root_dir=data_dir, dir='mr_to_ct')
      generate_25d_slices(data_dir, '/mnt/sdb2/Research/Bone/data/25D', n_slices=7)
  
  
  
  
  