class_docstr: creates the training/validation batch generator. Samples n_batch_size patients (draws a slice from each patient if 2D)
  from the data set while maintaining foreground-class balance. Returned patches are cropped/padded to pre_crop_size.
  Actual patch_size is obtained after data augmentation.
  :param data: data dictionary as provided by 'load_dataset'.
  :param batch_size: number of patients to sample for the batch
  :return dictionary containing the batch data (b, c, x, y, (z)) / seg (b, 1, x, y, (z)) / pids / class_target
class_docstr_tok: ["creates","the","training\/validation","batch","generator",".","Samples","n_batch_size","patients","(","draws","a","slice","from","each","patient","if","2D",")","from","the","data","set","while","maintaining","foreground-class","balance",".","Returned","patches","are","cropped\/padded","to","pre_crop_size",".","Actual","patch_size","is","obtained","after","data","augmentation",".",":","param","data",":","data","dictionary","as","provided","by","'load_dataset","'",".",":","param","batch_size",":","number","of","patients","to","sample","for","the","batch",":","return","dictionary","containing","the","batch","data","(","b",",","c",",","x",",","y",",","(","z",")",")","\/","seg","(","b",",","1",",","x",",","y",",","(","z",")",")","\/","pids","\/","class_target"]
code: def generate_train_batch(self):
      (
          batch_data,
          batch_segs,
          batch_pids,
          batch_targets,
          batch_patient_labels,
      ) = ([], [], [], [], [])
      class_targets_list = [
          v["class_target"] for (k, v) in self._data.items()
      ]
  
      if self.cf.head_classes > 2:
          # samples patients towards equilibrium of foreground classes on a roi-level (after randomly sampling the ratio "batch_sample_slack).
          batch_ixs = dutils.get_class_balanced_patients(
              class_targets_list,
              self.batch_size,
              self.cf.head_classes - 1,
              slack_factor=self.cf.batch_sample_slack,
          )
      else:
          batch_ixs = np.random.choice(
              len(class_targets_list), self.batch_size
          )
  
      patients = list(self._data.items())
  
      for b in batch_ixs:
          patient = patients[b][1]
  
          data = np.transpose(
              np.load(patient["data"], mmap_mode="r"), axes=(1, 2, 0)
          )[
              np.newaxis
          ]  # (c, y, x, z)
          seg = np.transpose(
              np.load(patient["seg"], mmap_mode="r"), axes=(1, 2, 0)
          )
          batch_pids.append(patient["pid"])
          batch_targets.append(patient["class_target"])
  
          if self.cf.dim == 2:
              # draw random slice from patient while oversampling slices containing foreground objects with p_fg.
              if len(patient["fg_slices"]) > 0:
                  fg_prob = self.p_fg / len(patient["fg_slices"])
                  bg_prob = (1 - self.p_fg) / (
                      data.shape[3] - len(patient["fg_slices"])
                  )
                  slices_prob = [
                      fg_prob if ix in patient["fg_slices"] else bg_prob
                      for ix in range(data.shape[3])
                  ]
                  slice_id = np.random.choice(
                      data.shape[3], p=slices_prob
                  )
              else:
                  slice_id = np.random.choice(data.shape[3])
  
              # if set to not None, add neighbouring slices to each selected slice in channel dimension.
              if self.cf.n_3D_context is not None:
                  padded_data = dutils.pad_nd_image(
                      data[0],
                      [(data.shape[-1] + (self.cf.n_3D_context * 2))],
                      mode="constant",
                  )
                  padded_slice_id = slice_id + self.cf.n_3D_context
                  data = np.concatenate(
                      [
                          padded_data[..., ii][np.newaxis]
                          for ii in range(
                              padded_slice_id - self.cf.n_3D_context,
                              padded_slice_id
                              + self.cf.n_3D_context
                              + 1,
                          )
                      ],
                      axis=0,
                  )
              else:
                  data = data[..., slice_id]
              seg = seg[..., slice_id]
  
          # pad data if smaller than pre_crop_size.
          if np.any(
              [
                  data.shape[dim + 1] < ps
                  for dim, ps in enumerate(self.cf.pre_crop_size)
              ]
          ):
              new_shape = [
                  np.max([data.shape[dim + 1], ps])
                  for dim, ps in enumerate(self.cf.pre_crop_size)
              ]
              data = dutils.pad_nd_image(
                  data, new_shape, mode="constant"
              )
              seg = dutils.pad_nd_image(seg, new_shape, mode="constant")
  
          # crop patches of size pre_crop_size, while sampling patches containing foreground with p_fg.
          crop_dims = [
              dim
              for dim, ps in enumerate(self.cf.pre_crop_size)
              if data.shape[dim + 1] > ps
          ]
          if len(crop_dims) > 0:
              fg_prob_sample = np.random.rand(1)
              # with p_fg: sample random pixel from random ROI and shift center by random value.
              if fg_prob_sample < self.p_fg and np.sum(seg) > 0:
                  seg_ixs = np.argwhere(
                      seg == np.random.choice(np.unique(seg)[1:], 1)
                  )
                  roi_anchor_pixel = seg_ixs[
                      np.random.choice(seg_ixs.shape[0], 1)
                  ][0]
                  assert seg[tuple(roi_anchor_pixel)] > 0
                  # sample the patch center coords. constrained by edges of images - pre_crop_size /2. And by
                  # distance to the desired ROI < patch_size /2.
                  # (here final patch size to account for center_crop after data augmentation).
                  sample_seg_center = {}
                  for ii in crop_dims:
                      low = np.max(
                          (
                              self.cf.pre_crop_size[ii] // 2,
                              roi_anchor_pixel[ii]
                              - (
                                  self.cf.patch_size[ii] // 2
                                  - self.crop_margin[ii]
                              ),
                          )
                      )
                      high = np.min(
                          (
                              data.shape[ii + 1]
                              - self.cf.pre_crop_size[ii] // 2,
                              roi_anchor_pixel[ii]
                              + (
                                  self.cf.patch_size[ii] // 2
                                  - self.crop_margin[ii]
                              ),
                          )
                      )
                      # happens if lesion on the edge of the image. dont care about roi anymore,
                      # just make sure pre-crop is inside image.
                      if low >= high:
                          low = data.shape[ii + 1] // 2 - (
                              data.shape[ii + 1] // 2
                              - self.cf.pre_crop_size[ii] // 2
                          )
                          high = data.shape[ii + 1] // 2 + (
                              data.shape[ii + 1] // 2
                              - self.cf.pre_crop_size[ii] // 2
                          )
                      sample_seg_center[ii] = np.random.randint(
                          low=low, high=high
                      )
  
              else:
                  # not guaranteed to be empty. probability of emptiness depends on the data.
                  sample_seg_center = {
                      ii: np.random.randint(
                          low=self.cf.pre_crop_size[ii] // 2,
                          high=data.shape[ii + 1]
                          - self.cf.pre_crop_size[ii] // 2,
                      )
                      for ii in crop_dims
                  }
  
              for ii in crop_dims:
                  min_crop = int(
                      sample_seg_center[ii]
                      - self.cf.pre_crop_size[ii] // 2
                  )
                  max_crop = int(
                      sample_seg_center[ii]
                      + self.cf.pre_crop_size[ii] // 2
                  )
                  data = np.take(
                      data,
                      indices=range(min_crop, max_crop),
                      axis=ii + 1,
                  )
                  seg = np.take(
                      seg, indices=range(min_crop, max_crop), axis=ii
                  )
  
          batch_data.append(data)
          batch_segs.append(seg[np.newaxis])
  
      data = np.array(batch_data)
      seg = np.array(batch_segs).astype(np.uint8)
      class_target = np.array(batch_targets)
      return {
          "data": data,
          "seg": seg,
          "pid": batch_pids,
          "class_target": class_target,
      }
  
imports: import numpy
  import os
  from collections import OrderedDict
  import pandas
  import pickle
  import time
  import subprocess
  import utils.dataloader_utils
  from batchgenerators.dataloading.data_loader import SlimDataLoaderBase
  from batchgenerators.transforms.spatial_transforms import MirrorTransform
  from batchgenerators.transforms.abstract_transforms import Compose
  from batchgenerators.dataloading.multi_threaded_augmenter import MultiThreadedAugmenter
  from batchgenerators.dataloading import SingleThreadedAugmenter
  from batchgenerators.transforms.spatial_transforms import SpatialTransform
  from batchgenerators.transforms.crop_and_pad_transforms import CenterCropTransform
  from batchgenerators.transforms.utility_transforms import ConvertSegToBoundingBoxCoordinates
label_desc: Use image node_id 2 for calling the BatchGenerator obj's underlying member method code with example usage: obj.generate_train_batch() and returns: dict