text: #    Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany
  #
  #    Licensed under the Apache License, Version 2.0 (the "License");
  #    you may not use this file except in compliance with the License.
  #    You may obtain a copy of the License at
  #
  #        http://www.apache.org/licenses/LICENSE-2.0
  #
  #    Unless required by applicable law or agreed to in writing, software
  #    distributed under the License is distributed on an "AS IS" BASIS,
  #    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  #    See the License for the specific language governing permissions and
  #    limitations under the License.
  
  from multiprocessing import Pool
  import SimpleITK as sitk
  import numpy as np
  from batchgenerators.utilities.file_and_folder_operations import *
  from nnunet.paths import nnUNet_raw_data
  from nnunet.paths import preprocessing_output_dir
  from skimage.io import imread
  
  
  def load_tiff_convert_to_nifti(img_file, lab_file, img_out_base, anno_out, spacing):
      img = imread(img_file)
      img_itk = sitk.GetImageFromArray(img.astype(np.float32))
      img_itk.SetSpacing(np.array(spacing)[::-1])
      sitk.WriteImage(img_itk, join(img_out_base + "_0000.nii.gz"))
  
      if lab_file is not None:
          l = imread(lab_file)
          l[l > 0] = 1
          l_itk = sitk.GetImageFromArray(l.astype(np.uint8))
          l_itk.SetSpacing(np.array(spacing)[::-1])
          sitk.WriteImage(l_itk, anno_out)
  
  
  def prepare_task(base, task_id, task_name, spacing):
      p = Pool(16)
  
      foldername = "Task%03.0d_%s" % (task_id, task_name)
  
      out_base = join(nnUNet_raw_data, foldername)
      imagestr = join(out_base, "imagesTr")
      imagests = join(out_base, "imagesTs")
      labelstr = join(out_base, "labelsTr")
      maybe_mkdir_p(imagestr)
      maybe_mkdir_p(imagests)
      maybe_mkdir_p(labelstr)
  
      train_patient_names = []
      test_patient_names = []
      res = []
  
      for train_sequence in [i for i in subfolders(base + "_train", join=False) if not i.endswith("_GT")]:
          train_cases = subfiles(join(base + '_train', train_sequence), suffix=".tif", join=False)
          for t in train_cases:
              casename = train_sequence + "_" + t[:-4]
              img_file = join(base + '_train', train_sequence, t)
              lab_file = join(base + '_train', train_sequence + "_GT", "SEG", "man_seg" + t[1:])
              if not isfile(lab_file):
                  continue
              img_out_base = join(imagestr, casename)
              anno_out = join(labelstr, casename + ".nii.gz")
              res.append(
                  p.starmap_async(load_tiff_convert_to_nifti, ((img_file, lab_file, img_out_base, anno_out, spacing),)))
              train_patient_names.append(casename)
  
      for test_sequence in [i for i in subfolders(base + "_test", join=False) if not i.endswith("_GT")]:
          test_cases = subfiles(join(base + '_test', test_sequence), suffix=".tif", join=False)
          for t in test_cases:
              casename = test_sequence + "_" + t[:-4]
              img_file = join(base + '_test', test_sequence, t)
              lab_file = None
              img_out_base = join(imagests, casename)
              anno_out = None
              res.append(
                  p.starmap_async(load_tiff_convert_to_nifti, ((img_file, lab_file, img_out_base, anno_out, spacing),)))
              test_patient_names.append(casename)
  
      _ = [i.get() for i in res]
  
      json_dict = {}
      json_dict['name'] = task_name
      json_dict['description'] = ""
      json_dict['tensorImageSize'] = "4D"
      json_dict['reference'] = ""
      json_dict['licence'] = ""
      json_dict['release'] = "0.0"
      json_dict['modality'] = {
          "0": "BF",
      }
      json_dict['labels'] = {
          "0": "background",
          "1": "cell",
      }
  
      json_dict['numTraining'] = len(train_patient_names)
      json_dict['numTest'] = len(test_patient_names)
      json_dict['training'] = [{'image': "./imagesTr/%s.nii.gz" % i, "label": "./labelsTr/%s.nii.gz" % i} for i in
                               train_patient_names]
      json_dict['test'] = ["./imagesTs/%s.nii.gz" % i for i in test_patient_names]
  
      save_json(json_dict, os.path.join(out_base, "dataset.json"))
      p.close()
      p.join()
  
  
  if __name__ == "__main__":
      base = "/media/fabian/My Book/datasets/CellTrackingChallenge/Fluo-C3DH-A549_ManAndSim"
      task_id = 75
      task_name = 'Fluo_C3DH_A549_ManAndSim'
      spacing = (1, 0.126, 0.126)
      prepare_task(base, task_id, task_name, spacing)
  
      task_name = "Task075_Fluo_C3DH_A549_ManAndSim"
      labelsTr = join(nnUNet_raw_data, task_name, "labelsTr")
      cases = subfiles(labelsTr, suffix='.nii.gz', join=False)
      splits = []
      splits.append(
          {'train': [i[:-7] for i in cases if i.startswith('01_') or i.startswith('02_SIM')],
           'val': [i[:-7] for i in cases if i.startswith('02_') and not i.startswith('02_SIM')]}
      )
      splits.append(
          {'train': [i[:-7] for i in cases if i.startswith('02_') or i.startswith('01_SIM')],
           'val': [i[:-7] for i in cases if i.startswith('01_') and not i.startswith('01_SIM')]}
      )
      splits.append(
          {'train': [i[:-7] for i in cases if i.startswith('01_') or i.startswith('02_') and not i.startswith('02_SIM')],
           'val': [i[:-7] for i in cases if i.startswith('02_SIM')]}
      )
      splits.append(
          {'train': [i[:-7] for i in cases if i.startswith('02_') or i.startswith('01_') and not i.startswith('01_SIM')],
           'val': [i[:-7] for i in cases if i.startswith('01_SIM')]}
      )
      save_pickle(splits, join(preprocessing_output_dir, task_name, "splits_final.pkl"))
  
  