text: import paddle
  import numpy as np
  
  
  class ColorizeHint:
      """Get hint and mask images for colorization.
  
      This method is prepared for user guided colorization tasks. Take the original RGB images as imput,
      we will obtain the local hints and correspoding mask to guid colorization process.
  
      Args:
         percent(float): Probability for ignoring hint in an iteration.
         num_points(int): Number of selected hints in an iteration.
         samp(str): Sample method, default is normal.
         use_avg(bool): Whether to use mean in selected hint area.
  
      Return:
          hint(np.ndarray): hint images
          mask(np.ndarray): mask images
      """
      def __init__(self, percent: float, num_points: int = None, samp: str = 'normal', use_avg: bool = True):
          self.percent = percent
          self.num_points = num_points
          self.samp = samp
          self.use_avg = use_avg
  
      def __call__(self, data: np.ndarray, hint: np.ndarray, mask: np.ndarray):
          sample_Ps = [1, 2, 3, 4, 5, 6, 7, 8, 9]
          self.data = data
          self.hint = hint
          self.mask = mask
          N, C, H, W = data.shape
          for nn in range(N):
              pp = 0
              cont_cond = True
              while cont_cond:
                  if self.num_points is None:  # draw from geometric
                      # embed()
                      cont_cond = np.random.rand() < (1 - self.percent)
                  else:  # add certain number of points
                      cont_cond = pp < self.num_points
                  if not cont_cond:  # skip out of loop if condition not met
                      continue
                  P = np.random.choice(sample_Ps)  # patch size
                  # sample location
                  if self.samp == 'normal':  # geometric distribution
                      h = int(np.clip(np.random.normal((H - P + 1) / 2., (H - P + 1) / 4.), 0, H - P))
                      w = int(np.clip(np.random.normal((W - P + 1) / 2., (W - P + 1) / 4.), 0, W - P))
                  else:  # uniform distribution
                      h = np.random.randint(H - P + 1)
                      w = np.random.randint(W - P + 1)
                  # add color point
                  if self.use_avg:
                      # embed()
                      hint[nn, :, h:h + P, w:w + P] = np.mean(np.mean(data[nn, :, h:h + P, w:w + P],
                                                                      axis=2,
                                                                      keepdims=True),
                                                              axis=1,
                                                              keepdims=True).reshape(1, C, 1, 1)
                  else:
                      hint[nn, :, h:h + P, w:w + P] = data[nn, :, h:h + P, w:w + P]
                  mask[nn, :, h:h + P, w:w + P] = 1
                  # increment counter
                  pp += 1
  
          mask -= 0.5
          return hint, mask
  
  
  class ColorizePreprocess:
      """Prepare dataset for image Colorization.
  
      Args:
         ab_thresh(float): Thresh value for setting mask value.
         p(float): Probability for ignoring hint in an iteration.
         num_points(int): Number of selected hints in an iteration.
         samp(str): Sample method, default is normal.
         use_avg(bool): Whether to use mean in selected hint area.
         is_train(bool): Training process or not.
  
      Return:
          data(dict)：The preprocessed data for colorization.
  
      """
      def __init__(self,
                   ab_thresh: float = 0.,
                   p: float = 0.,
                   points: int = None,
                   samp: str = 'normal',
                   use_avg: bool = True):
          self.ab_thresh = ab_thresh
          self.p = p
          self.num_points = points
          self.samp = samp
          self.use_avg = use_avg
          self.gethint = ColorizeHint(percent=self.p, num_points=self.num_points, samp=self.samp, use_avg=self.use_avg)
  
      def __call__(self, data_lab):
          """
          This method seperates the L channel and AB channel, obtain hint, mask and real_B_enc as the input for colorization task.
  
          Args:
             img(np.ndarray|paddle.Tensor): LAB image.
  
          Returns:
              data(dict)：The preprocessed data for colorization.
          """
          if type(data_lab) is not np.ndarray:
              data_lab = data_lab.numpy()
          data = {}
          A = 2 * 110 / 10 + 1
          data['A'] = data_lab[:, [0], :, :]
          data['B'] = data_lab[:, 1:, :, :]
          if self.ab_thresh > 0:  # mask out grayscale images
              thresh = 1. * self.ab_thresh / 110
              mask = np.sum(np.abs(np.max(np.max(data['B'], axis=3), axis=2) - np.min(np.min(data['B'], axis=3), axis=2)),
                            axis=1)
              mask = (mask >= thresh)
              data['A'] = data['A'][mask, :, :, :]
              data['B'] = data['B'][mask, :, :, :]
              if np.sum(mask) == 0:
                  return None
          data_ab_rs = np.round((data['B'][:, :, ::4, ::4] * 110. + 110.) / 10.)  # normalized bin number
          data['real_B_enc'] = data_ab_rs[:, [0], :, :] * A + data_ab_rs[:, [1], :, :]
          data['hint_B'] = np.zeros(shape=data['B'].shape)
          data['mask_B'] = np.zeros(shape=data['A'].shape)
          data['hint_B'], data['mask_B'] = self.gethint(data['B'], data['hint_B'], data['mask_B'])
          data['A'] = paddle.to_tensor(data['A'].astype(np.float32))
          data['B'] = paddle.to_tensor(data['B'].astype(np.float32))
          data['real_B_enc'] = paddle.to_tensor(data['real_B_enc'].astype(np.int64))
          data['hint_B'] = paddle.to_tensor(data['hint_B'].astype(np.float32))
          data['mask_B'] = paddle.to_tensor(data['mask_B'].astype(np.float32))
          return data
  