text: import numpy as np
  
  from blast_ct.nifti.augmention import RandomAugmentation
  
  
  # All patch samplers can't alter state after __init__ because of dataset class when num_workers > 0
  
  def get_patch_and_padding(shape, patch_size, center):
      padding = [[0, 0], [0, 0], [0, 0]]
      patch = []
      for dim, (size, ps, c) in enumerate(zip(shape, patch_size, center)):
          start = c - ps // 2
          end = start + ps
          if start < 0:
              padding[dim][0] = -start
              start = 0
          if end > size:
              padding[dim][1] = end - size
              end = size
          patch.append(slice(start, end, 1))
  
      return patch, padding
  
  
  class PatchSampler(object):
      def __init__(self, image_patch_shape, target_patch_shape, augmentation=None):
          if any((i < l) for i, l in zip(image_patch_shape, target_patch_shape)):
              raise ValueError('Label map patch size must be smaller or equal to image patch size.')
          self.image_patch_size = image_patch_shape
          self.target_patch_size = target_patch_shape
          self.augmentation = [] if augmentation is None else augmentation
          assert isinstance(self.augmentation, list)
          for el in self.augmentation:
              assert isinstance(el, RandomAugmentation)
  
      def get_target_patch(self, target, center):
          if target is None:
              return target
          patch, padding = get_patch_and_padding(target.shape, self.target_patch_size, center)
          target_patch = target[tuple(patch)]
          if any([any(p) for p in padding]):
              target_patch = np.pad(target_patch, padding, mode='constant', constant_values=0)
          return target_patch
  
      def get_image_patch(self, image, center):
          image_shape = image.shape
          patch, padding = get_patch_and_padding(image_shape[1:], self.image_patch_size, center)
          padding = [[0, 0]] + padding
          patch = [slice(0, image_shape[0], 1)] + patch
  
          image_patch = image[tuple(patch)]
  
          if any([any(p) for p in padding]):
              image_patch = np.pad(image_patch, padding, mode='edge')
          return image_patch
  
      def get_patches(self, center, image, target, mask):
          image_patch = self.get_image_patch(image, center)
          target_patch = self.get_target_patch(target, center)
          mask_patch = self.get_target_patch(mask, center)
          return image_patch, target_patch, mask_patch
  
      def sample_patch_center(self, target, mask):
          raise NotImplementedError
  
      def __call__(self, image, target, mask=None):
          # ensure that the mask has at least one point
          mask = np.ones_like(target) if mask is None else mask
          if np.sum(mask) == 0:
              raise ValueError('Empty sampling mask')
  
          center = self.sample_patch_center(target, mask)
          image_patch, target_patch, mask_patch = self.get_patches(center, image, target, mask)
          for augmentation in self.augmentation:
              image_patch, target_patch, mask_patch = augmentation(image_patch, target_patch, mask_patch)
          # these copies are only here because of flipping and should be removed in the future
          return image_patch.copy(), target_patch.copy(), mask_patch.copy()
  
  
  class StochasticPatchSampler(PatchSampler):
      def get_sampling_mask(self, target, mask):
          raise NotImplementedError
  
      def sample_patch_center(self, target, mask):
          sampling_mask = self.get_sampling_mask(target, mask)
          points = np.argwhere(sampling_mask)
          center = points[np.random.choice(len(points))]
          return center
  
  
  class RandomPatchSampler(StochasticPatchSampler):
      def __init__(self, image_patch_shape, target_patch_shape, augmentation=None):
          super().__init__(image_patch_shape, target_patch_shape, augmentation)
  
      def get_sampling_mask(self, target, mask):
          return mask
  
  
  class ConditionalPatchSampler(StochasticPatchSampler):
      def __init__(self, image_patch_shape, target_patch_shape,
                   class_probabilities, augmentation=None, n_tries=3):
          super().__init__(image_patch_shape, target_patch_shape, augmentation)
          self.class_probabilities = class_probabilities
          self.n_tries = n_tries
  
      def get_sampling_mask(self, target, mask):
          # try sampling points from a random label a maximum of n_tries
          for i in range(self.n_tries):
              label = np.random.choice(len(self.class_probabilities), p=self.class_probabilities)
              sampling_mask = np.logical_and(target == label, mask)
              if sampling_mask.any():
                  return sampling_mask
          # if sampling mask is empty default to mask
          return mask
  
  
  class ForegroundBackgroundPatchSampler(StochasticPatchSampler):
      def __init__(self, image_patch_shape, target_patch_shape,
                   foreground_probability=.5, augmentation=None, n_tries=3):
          super().__init__(image_patch_shape, target_patch_shape, augmentation)
          self.foreground_probability = foreground_probability
          self.n_tries = n_tries
  
      def get_sampling_mask(self, target, mask):
          # try sampling points from a random label a maximum of n_tries
          for i in range(self.n_tries):
              is_foreground = np.random.choice((False, True),
                                               p=(1 - self.foreground_probability, self.foreground_probability))
              if is_foreground:
                  sampling_mask = np.logical_and(target != 0, mask)
              else:
                  sampling_mask = np.logical_and(target == 0, mask)
              if sampling_mask.any():
                  return sampling_mask
          # if sampling mask is empty default to mask
          return mask
  
  
  class BoundingBoxCenteredPatchSampler(PatchSampler):
      def __init__(self, image_patch_shape, target_patch_shape, augmentation):
          super().__init__(image_patch_shape, target_patch_shape, augmentation)
  
      def sample_patch_center(self, target, mask):
          return np.array(tuple((np.max(arr) - np.min(arr)) // 2 for arr in np.where(mask)))
  