management:
  wandb_project_name: "segmentation"
  dir_checkpoint: "./checkpoints"
  save_checkpoint: False
  processing_stage: 'Train' # OR 'Test' OR 'Inference'

model:
  name: 'hf'
  hf_path: 'nvidia/mit-b2'
  n_classes: 2
  bilinear: False
  id2label: {0: "background", 1: "foreground"}
  label2id: {"background": 0, "foreground": 1}

data:
  source: 'snowleopard_v2'
  inference_dir: './inference'
  mask_suffix: '_mask.png'
  inference_mask_dir: 'segmentation_mask_result'
  training_percent: null
  num_workers: 2
  n_channels: 3
  img_height: 400
  img_width: 400
  transforms_train: ["random_rotation", "random_crop"]
  transforms_test: 'center_crop'
  transforms_inference: 'resize'
  norm_mean: null
  norm_std: null

train:
  epochs: 25
  batch_size: 2
  optim: "adamw"
  scheduler: "linear"
  scheduler_patience: 2
  lr: 1.0e-5
  amp: False

test:
  use_cuda: True
  path_to_model: ''
  batch_size: 2
