a preprocessing layer which randomly flips images during training this layer will flip the images horizontally and or vertically based on the mode attribute during inference time the output will be identical to input call the layer with trainingtrue to flip the input input pixel values can be of any range e g 0 1 or 0 255 and of integer or floating point dtype by default the layer will output floats note this layer is safe to use inside a tf data pipeline independently of which backend you re using input shape 3d unbatched or 4d batched tensor with shape height width channels in channelslast format output shape 3d unbatched or 4d batched tensor with shape height width channels in channelslast format args mode string indicating which flip mode to use can be horizontal vertical or horizontalandvertical horizontal is a leftright flip and vertical is a topbottom flip defaults to horizontalandvertical seed integer used to create a random seed kwargs base layer keyword arguments such as name and dtype a preprocessing layer which randomly flips images during training this layer will flip the images horizontally and or vertically based on the mode attribute during inference time the output will be identical to input call the layer with training true to flip the input input pixel values can be of any range e g 0 1 or 0 255 and of integer or floating point dtype by default the layer will output floats note this layer is safe to use inside a tf data pipeline independently of which backend you re using input shape 3d unbatched or 4d batched tensor with shape height width channels in channels_last format output shape 3d unbatched or 4d batched tensor with shape height width channels in channels_last format args mode string indicating which flip mode to use can be horizontal vertical or horizontal_and_vertical horizontal is a left right flip and vertical is a top bottom flip defaults to horizontal_and_vertical seed integer used to create a random seed kwargs base layer keyword arguments such as name and dtype
from keras.api_export import keras_export
from keras.layers.preprocessing.tf_data_layer import TFDataLayer
from keras.random.seed_generator import SeedGenerator

HORIZONTAL = "horizontal"
VERTICAL = "vertical"
HORIZONTAL_AND_VERTICAL = "horizontal_and_vertical"


@keras_export("keras.layers.RandomFlip")
class RandomFlip(TFDataLayer):
    

    def __init__(self, mode=HORIZONTAL_AND_VERTICAL, seed=None, **kwargs):
        super().__init__(**kwargs)
        self.seed = seed
        self.generator = SeedGenerator(seed)
        self.mode = mode
        self._convert_input_args = False
        self._allow_non_tensor_positional_args = True

    def _randomly_flip_inputs(self, inputs):
        inputs_shape = self.backend.shape(inputs)
        unbatched = len(inputs_shape) == 3
        if unbatched:
            inputs = self.backend.numpy.expand_dims(inputs, axis=0)
            inputs_shape = self.backend.shape(inputs)

        batch_size = inputs_shape[0]
        flipped_outputs = inputs
        seed_generator = self._get_seed_generator(self.backend._backend)
        if self.mode == HORIZONTAL or self.mode == HORIZONTAL_AND_VERTICAL:
            flipped_outputs = self.backend.numpy.where(
                self.backend.random.uniform(
                    shape=(batch_size, 1, 1, 1), seed=seed_generator
                )
                <= 0.5,
                self.backend.numpy.flip(flipped_outputs, axis=-2),
                flipped_outputs,
            )
        if self.mode == VERTICAL or self.mode == HORIZONTAL_AND_VERTICAL:
            flipped_outputs = self.backend.numpy.where(
                self.backend.random.uniform(
                    shape=(batch_size, 1, 1, 1), seed=seed_generator
                )
                <= 0.5,
                self.backend.numpy.flip(flipped_outputs, axis=-3),
                flipped_outputs,
            )
        if unbatched:
            flipped_outputs = self.backend.numpy.squeeze(
                flipped_outputs, axis=0
            )
        return flipped_outputs

    def call(self, inputs, training=True):
        inputs = self.backend.cast(inputs, self.compute_dtype)
        if training:
            return self._randomly_flip_inputs(inputs)
        else:
            return inputs

    def compute_output_shape(self, input_shape):
        return input_shape

    def get_config(self):
        config = super().get_config()
        config.update({"seed": self.seed, "mode": self.mode})
        return config