# use a fixed random seed to guarantee that when you run the code twice you will get the same outcome
manual_seed: 0
model:
  name: UNet3D
  # number of input channels to the model
  in_channels: 1
  # number of output channels
  out_channels: 1
  # determines the order of operators in a single layer (crg - Conv3d+ReLU+GroupNorm)
  layer_order: gcr
  # initial number of feature maps
  f_maps: 32
  # number of groups in the groupnorm
  num_groups: 8
  # apply element-wise nn.Sigmoid after the final 1x1x1 convolution, otherwise apply nn.Softmax
  final_sigmoid: true
# loss function to be used during training
loss:
  name: BCEDiceLoss
  # a target value that is ignored and does not contribute to the input gradient
  ignore_index: null
  # skip the last channel in the target (i.e. when last channel contains data not relevant for the loss)
  skip_last_target: true
optimizer:
  # initial learning rate
  learning_rate: 0.0002
  # weight decay
  weight_decay: 0.00001
# evaluation metric
eval_metric:
  # use AdaptedRandError metric
  name: BoundaryAdaptedRandError
  # probability maps threshold
  threshold: 0.4
  # use the last target channel to compute the metric
  use_last_target: true
  # use only the first channel for computing the metric
  use_first_input: true
lr_scheduler:
  name: ReduceLROnPlateau
  # make sure to use the 'min' mode cause lower AdaptedRandError is better
  mode: min
  factor: 0.5
  patience: 15
trainer:
  # model with lower eval score is considered better
  eval_score_higher_is_better: False
  # path to the checkpoint directory
  checkpoint_dir: '/g/kreshuk/wolny/workspace/for-pytorch-3dunet/FOR_paper_root/final/unet/bce_dice/ds1x'
  # path to latest checkpoint; if provided the training will be resumed from that checkpoint
  resume: null
  # how many iterations between validations
  validate_after_iters: 1000
  # how many iterations between tensorboard logging
  log_after_iters: 100
  # max number of epochs
  epochs: 1000
  # max number of iterations
  iters: 150000
# Configure training and validation loaders
loaders:
  # how many subprocesses to use for data loading
  num_workers: 8
  # path to the raw data within the H5
  raw_internal_path: /raw
  # path to the the label data withtin the H5
  label_internal_path: /label
  # configuration of the train loader
  train:
    # paths to the training datasets
    file_paths:
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie1_t00000_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie1_t00002_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie1_t00003_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie1_t00008_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie1_t00009_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie1_t00010_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie1_t00012_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie1_t00014_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie1_t00016_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie1_t00018_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie1_t00020_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie1_t00025_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie1_t00035_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie1_t00040_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie1_t00049_crop_gt.h5'
      # Beautiful
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie2_T00000_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie2_T00002_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie2_T00004_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie2_T00006_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie2_T00008_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie2_T00012_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie2_T00014_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie2_T00016_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie2_T00018_crop_gt.h5'
      # Nice
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie3_T00000_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie3_T00002_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Train/Movie3_T00004_crop_gt.h5'

    # SliceBuilder configuration, i.e. how to iterate over the input volume patch-by-patch
    slice_builder:
      name: SliceBuilder
      # train patch size given to the network (adapt to fit in your GPU mem, generally the bigger patch the better)
      patch_shape: [80, 170, 170]
      # train stride between patches
      stride_shape: [20, 40, 40]

    transformer:
      raw:
        - name: Standardize
        - name: RandomFlip
        - name: RandomRotate90
        - name: RandomRotate
          # rotate only in ZY plane due to anisotropy
          axes: [[2, 1]]
          angle_spectrum: 15
          mode: reflect
        - name: ElasticDeformation
          spline_order: 3
        - name: AdditiveGaussianNoise
          execution_probability: 0.25
        - name: AdditivePoissonNoise
          execution_probability: 0.25
        - name: ToTensor
          expand_dims: true
      label:
        - name: RandomFlip
        - name: RandomRotate90
        - name: RandomRotate
          # rotate only in ZY plane due to anisotropy
          axes: [[2, 1]]
          angle_spectrum: 15
          mode: reflect
        - name: ElasticDeformation
          spline_order: 0
        - name: StandardLabelToBoundary
          # append original ground truth labels to the last channel (to be able to compute the eval metric)
          append_label: true
          # guassian blur and threshold the boundary map in order to help with overfitting
          blur: true
          # stddev of Gaussian kernel
          sigma: 1.0
        - name: ToTensor
          expand_dims: false

  # configuration of the val loader
  val:
    # paths to the val datasets
    file_paths:
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Val/Movie1_t00004_crop_gt.h5'
      - '/g/kreshuk/wolny/Datasets/LateralRoot/Val/Movie1_t00030_crop_gt.h5'

    # SliceBuilder configuration, i.e. how to iterate over the input volume patch-by-patch
    slice_builder:
      name: SliceBuilder
      # train patch size given to the network (adapt to fit in your GPU mem, generally the bigger patch the better)
      patch_shape: [80, 170, 170]
      # train stride between patches
      stride_shape: [80, 170, 170]

    # data augmentation
    transformer:
      raw:
        - name: Standardize
        - name: ToTensor
          expand_dims: true
      label:
        - name: StandardLabelToBoundary
          append_label: true
          blur: true
          sigma: 1.0
        - name: ToTensor
          expand_dims: false