a preprocessing layer which randomly crops images during training during training this layer will randomly choose a location to crop images down to a target size the layer will crop all the images in the same batch to the same cropping location at inference time and during training if an input image is smaller than the target size the input will be resized and cropped so as to return the largest possible window in the image that matches the target aspect ratio if you need to apply random cropping at inference time set training to true when calling the layer 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 targetheight targetwidth channels args height integer the height of the output shape width integer the width of the output shape seed integer used to create a random seed kwargs base layer keyword arguments such as name and dtype smartresize will always output float32 so we need to recast a preprocessing layer which randomly crops images during training during training this layer will randomly choose a location to crop images down to a target size the layer will crop all the images in the same batch to the same cropping location at inference time and during training if an input image is smaller than the target size the input will be resized and cropped so as to return the largest possible window in the image that matches the target aspect ratio if you need to apply random cropping at inference time set training to true when calling the layer 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 target_height target_width channels args height integer the height of the output shape width integer the width of the output shape seed integer used to create a random seed kwargs base layer keyword arguments such as name and dtype smart_resize will always output float32 so we need to re cast
from keras import backend
from keras.api_export import keras_export
from keras.layers.preprocessing.tf_data_layer import TFDataLayer
from keras.random.seed_generator import SeedGenerator
from keras.utils import image_utils


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

    def __init__(
        self, height, width, seed=None, data_format=None, name=None, **kwargs
    ):
        super().__init__(name=name, **kwargs)
        self.height = height
        self.width = width
        self.seed = seed or backend.random.make_default_seed()
        self.generator = SeedGenerator(seed)
        self.data_format = backend.standardize_data_format(data_format)

        if self.data_format == "channels_first":
            self.height_axis = -2
            self.width_axis = -1
        elif self.data_format == "channels_last":
            self.height_axis = -3
            self.width_axis = -2

        self.supports_masking = False
        self.supports_jit = False
        self._convert_input_args = False
        self._allow_non_tensor_positional_args = True

    def call(self, inputs, training=True):
        inputs = self.backend.cast(inputs, self.compute_dtype)
        input_shape = self.backend.shape(inputs)
        is_batched = len(input_shape) > 3
        if not is_batched:
            inputs = self.backend.numpy.expand_dims(inputs, axis=0)

        h_diff = input_shape[self.height_axis] - self.height
        w_diff = input_shape[self.width_axis] - self.width

        def random_crop():
            input_height, input_width = (
                input_shape[self.height_axis],
                input_shape[self.width_axis],
            )

            seed_generator = self._get_seed_generator(self.backend._backend)
            h_start = self.backend.cast(
                self.backend.random.uniform(
                    (),
                    0,
                    maxval=float(input_height - self.height + 1),
                    seed=seed_generator,
                ),
                "int32",
            )
            w_start = self.backend.cast(
                self.backend.random.uniform(
                    (),
                    0,
                    maxval=float(input_width - self.width + 1),
                    seed=seed_generator,
                ),
                "int32",
            )
            if self.data_format == "channels_last":
                return self.backend.core.slice(
                    inputs,
                    self.backend.numpy.stack([0, h_start, w_start, 0]),
                    [
                        self.backend.shape(inputs)[0],
                        self.height,
                        self.width,
                        self.backend.shape(inputs)[3],
                    ],
                )
            else:
                return self.backend.core.slice(
                    inputs,
                    self.backend.numpy.stack([0, 0, h_start, w_start]),
                    [
                        self.backend.shape(inputs)[0],
                        self.backend.shape(inputs)[1],
                        self.height,
                        self.width,
                    ],
                )

        def resize():
            outputs = image_utils.smart_resize(
                inputs,
                [self.height, self.width],
                data_format=self.data_format,
                backend_module=self.backend,
            )
            
            return self.backend.cast(outputs, self.compute_dtype)

        if isinstance(h_diff, int) and isinstance(w_diff, int):
            if training and h_diff >= 0 and w_diff >= 0:
                outputs = random_crop()
            else:
                outputs = resize()
        else:
            predicate = self.backend.numpy.logical_and(
                training,
                self.backend.numpy.logical_and(h_diff >= 0, w_diff >= 0),
            )
            outputs = self.backend.cond(
                predicate,
                random_crop,
                resize,
            )

        if not is_batched:
            outputs = self.backend.numpy.squeeze(outputs, axis=0)
        return outputs

    def compute_output_shape(self, input_shape, *args, **kwargs):
        input_shape = list(input_shape)
        input_shape[self.height_axis] = self.height
        input_shape[self.width_axis] = self.width
        return tuple(input_shape)

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