title: Deal with random parameters in keras preprocessing layers
body: <p>I have a tensorflow preprocessing model which create 3D Synthetic images with 3D label maps.
  The inputs have 5 dimension: (batchsize, label_shape, channel) and the outputs should have 5 dimensions aswell</p>
  <p>Here is the code of the model I got: <a href="https://github.com/BBillot/SynthSeg/blob/master/SynthSeg/labels_to_image_model.py" rel="nofollow noreferrer">https://github.com/BBillot/SynthSeg/blob/master/SynthSeg/labels_to_image_model.py</a> (1)</p>
  <p>At the end of this model, I would like to add a layer which should only save a slab of the image and the label map with a probability. To make my slab, I randomly select an axis, then I randomly select a center on this axis and I make a segment with a fixed size on this axis. All the voxels which have not one index on this segment are fixed to 0.</p>
  <p>Here is the code I have made:</p>
  <pre><code>class RandomSlab(Layer):
      def __init__(self, size_slab, prob, **kwargs):
          super(RandomSlab, self).__init__(**kwargs)
          self.prob = prob
          self.dim = [1, 2, 3]
          self.size_slab = size_slab
          self.axis_selected = None
          self.shape_vol = None
  
      def get_config(self):
          config = super(self).get_config()
          config['prob'] = self.prob
          config['dim'] = self.dim
          config['size_slab'] = self.size_slab
          config['axis_selected'] = self.axis_selected
          config['shape_vol'] = self.shape_vol
  
          return config
  
      def call(self, inputs, **kwargs):
          image = inputs[0]
          labels = inputs[1]
          action = random.uniform(0, 1)
          value_to_take = np.zeros(self.shape_vol, dtype=bool)
  
          self.axis_selected = random.choice(self.dim)
  
          if inputs[0].shape[self.axis_selected] &lt;= self.size_slab:
              self.size_slab = inputs[0].shape[self.axis_selected]
          print(&quot;axis, &quot;, self.axis_selected)
  
          if action &gt; self.prob:
              return inputs
          print(len(value_to_take[int(self.size_slab / 2):int(1.5*self.size_slab), :, :, :]))
  
          if self.axis_selected == 1:
              slab_center = random.randint(int(self.size_slab/2), int(self.shape_vol[0]-self.size_slab/2))
              value_to_take[int(slab_center - self.size_slab / 2):int(slab_center + self.size_slab), :, :, :] = True
  
          elif self.axis_selected == 2:
              slab_center = random.randint(int(self.size_slab/2), int(self.shape_vol[1]-self.size_slab/2))
              value_to_take[:, int(slab_center - self.size_slab / 2):int(slab_center + self.size_slab), :, :] = True
  
          else:
              slab_center = random.randint(int(self.size_slab/2), int(self.shape_vol[2]-self.size_slab/2))
              value_to_take[:, :, int(slab_center - self.size_slab / 2):int(slab_center + self.size_slab), :] = True
          image = tf.where(value_to_take, image, 0)
          labels = tf.where(value_to_take, labels, 0)
          print(&quot;slab&quot;)
          return [image, labels]
  
      def build(self, input_shape):
          self.shape_vol = (input_shape[0][1], input_shape[0][2], input_shape[0][3], input_shape[0][4])
          self.built = True
          super(RandomSlab, self).build(input_shape)
  </code></pre>
  <p>and then I had the line <code>[image, labels] = layers.RandomSlab(size_slab=31, prob=1)([image, labels])</code> at line 240 of (1). To check.</p>
  <p>This solution doesn't have the good behavior because the probability is computed once at the creation of the model. I wanted the probability and the axis selected to be computed each time an input is given. How can I do this?</p>
  <p>Thank you for your help</p>