text: # Copyright 2020 - 2021 MONAI Consortium
  # 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.
  
  import logging
  
  import numpy as np
  import torch
  from monai.inferers import SimpleInferer
  from monai.losses import DiceLoss
  from monai.transforms import (
      Activationsd,
      AsDiscreted,
      CropForegroundd,
      EnsureChannelFirstd,
      LoadImaged,
      NormalizeIntensityd,
      Orientationd,
      RandAffined,
      RandFlipd,
      RandShiftIntensityd,
      Resized,
      Spacingd,
      ToTensord,
  )
  
  from monailabel.tasks.train.basic_train import BasicTrainTask
  
  logger = logging.getLogger(__name__)
  
  
  class MyTrain(BasicTrainTask):
      def __init__(
          self,
          model_dir,
          network,
          description="Train Segmentation model for left-atrium",
          **kwargs,
      ):
          self._network = network
          super().__init__(model_dir, description, **kwargs)
  
      def network(self):
          return self._network
  
      def optimizer(self):
          return torch.optim.Adam(self._network.parameters(), lr=0.0001)
  
      def loss_function(self):
          return DiceLoss(to_onehot_y=True, softmax=True)
  
      def train_pre_transforms(self):
          return [
              LoadImaged(keys=("image", "label")),
              EnsureChannelFirstd(keys=("image", "label")),
              Spacingd(
                  keys=("image", "label"),
                  pixdim=(1.0, 1.0, 1.0),
                  mode=("bilinear", "nearest"),
              ),
              Orientationd(keys=("image", "label"), axcodes="RAS"),
              NormalizeIntensityd(keys="image"),
              RandShiftIntensityd(keys="image", offsets=0.1, prob=0.5),
              CropForegroundd(keys=("image", "label"), source_key="image"),
              RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
              RandAffined(
                  keys=["image", "label"],
                  mode=("bilinear", "nearest"),
                  prob=1.0,
                  spatial_size=(256, 256, 128),
                  rotate_range=(0, 0, np.pi / 15),
                  scale_range=(0.1, 0.1, 0.1),
              ),
              ToTensord(keys=("image", "label")),
          ]
  
      def train_post_transforms(self):
          return [
              Activationsd(keys="pred", softmax=True),
              AsDiscreted(
                  keys=("pred", "label"),
                  argmax=(True, False),
                  to_onehot=True,
                  n_classes=2,
              ),
          ]
  
      def val_pre_transforms(self):
          return [
              LoadImaged(keys=("image", "label")),
              EnsureChannelFirstd(keys=("image", "label")),
              Spacingd(
                  keys=("image", "label"),
                  pixdim=(1.0, 1.0, 1.0),
                  mode=("bilinear", "nearest"),
              ),
              Orientationd(keys=("image", "label"), axcodes="RAS"),
              NormalizeIntensityd(keys="image"),
              CropForegroundd(keys=("image", "label"), source_key="image"),
              Resized(keys=("image", "label"), spatial_size=(256, 256, 128), mode=("area", "nearest")),
              ToTensord(keys=("image", "label")),
          ]
  
      def val_inferer(self):
          return SimpleInferer()
  