keras-hub-nightly 0.16.1.dev202410020340__py3-none-any.whl → 0.19.0.dev202501260345__py3-none-any.whl

keras_hub/src/models/image_object_detector.py

@@ -0,0 +1,87 @@
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.task import Task
+
+
+@keras_hub_export("keras_hub.models.ImageObjectDetector")
+class ImageObjectDetector(Task):
+    """Base class for all image object detection tasks.
+
+    The `ImageObjectDetector` tasks wrap a `keras_hub.models.Backbone` and
+    a `keras_hub.models.Preprocessor` to create a model that can be used for
+    object detection. `ImageObjectDetector` tasks take an additional
+    `num_classes` argument, controlling the number of predicted output classes.
+
+    To fine-tune with `fit()`, pass a dataset containing tuples of `(x, y)`
+    labels where `x` is a string and `y` is dictionary with `boxes` and
+    `classes`.
+
+    All `ImageObjectDetector` tasks include a `from_preset()` constructor which
+    can be used to load a pre-trained config and weights.
+    """
+
+    def compile(
+        self,
+        optimizer="auto",
+        box_loss="auto",
+        classification_loss="auto",
+        metrics=None,
+        **kwargs,
+    ):
+        """Configures the `ImageObjectDetector` task for training.
+
+        The `ImageObjectDetector` task extends the default compilation signature
+        of `keras.Model.compile` with defaults for `optimizer`, `loss`, and
+        `metrics`. To override these defaults, pass any value
+        to these arguments during compilation.
+
+        Args:
+            optimizer: `"auto"`, an optimizer name, or a `keras.Optimizer`
+                instance. Defaults to `"auto"`, which uses the default optimizer
+                for the given model and task. See `keras.Model.compile` and
+                `keras.optimizers` for more info on possible `optimizer` values.
+            box_loss: `"auto"`, a loss name, or a `keras.losses.Loss` instance.
+                Defaults to `"auto"`, where a
+                `keras.losses.Huber` loss will be
+                applied for the object detector task. See
+                `keras.Model.compile` and `keras.losses` for more info on
+                possible `loss` values.
+            classification_loss: `"auto"`, a loss name, or a `keras.losses.Loss`
+                instance. Defaults to `"auto"`, where a
+                `keras.losses.BinaryFocalCrossentropy` loss will be
+                applied for the object detector task. See
+                `keras.Model.compile` and `keras.losses` for more info on
+                possible `loss` values.
+            metrics: `a list of metrics to be evaluated by
+                the model during training and testing. Defaults to `None`.
+                See `keras.Model.compile` and `keras.metrics` for
+                more info on possible `metrics` values.
+            **kwargs: See `keras.Model.compile` for a full list of arguments
+                supported by the compile method.
+        """
+        if optimizer == "auto":
+            optimizer = keras.optimizers.Adam(5e-5)
+        if box_loss == "auto":
+            box_loss = keras.losses.Huber(reduction="sum")
+        if classification_loss == "auto":
+            activation = getattr(self, "activation", None)
+            activation = keras.activations.get(activation)
+            from_logits = activation != keras.activations.sigmoid
+            classification_loss = keras.losses.BinaryFocalCrossentropy(
+                from_logits=from_logits, reduction="sum"
+            )
+        if metrics is not None:
+            raise ValueError("User metrics not yet supported")
+
+        losses = {
+            "bbox_regression": box_loss,
+            "cls_logits": classification_loss,
+        }
+
+        super().compile(
+            optimizer=optimizer,
+            loss=losses,
+            metrics=metrics,
+            **kwargs,
+        )

keras_hub/src/models/image_object_detector_preprocessor.py

@@ -0,0 +1,57 @@
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.preprocessor import Preprocessor
+from keras_hub.src.utils.tensor_utils import preprocessing_function
+
+
+@keras_hub_export("keras_hub.models.ImageObjectDetectorPreprocessor")
+class ImageObjectDetectorPreprocessor(Preprocessor):
+    """Base class for object detector preprocessing layers.
+
+    `ImageObjectDetectorPreprocessor` tasks wraps a
+    `keras_hub.layers.Preprocessor` to create a preprocessing layer for
+    object detection tasks. It is intended to be paired with a
+    `keras_hub.models.ImageObjectDetector` task.
+
+    All `ImageObjectDetectorPreprocessor` take three inputs, `x`, `y`, and
+    `sample_weight`. `x`, the first input, should always be included. It can
+    be a image or batch of images. See examples below. `y` and `sample_weight`
+    are optional inputs that will be passed through unaltered. Usually, `y` will
+    be the a dict of `{"boxes": Tensor(batch_size, num_boxes, 4),
+    "classes": (batch_size, num_boxes)}.
+
+    The layer will returns either `x`, an `(x, y)` tuple if labels were
+    provided, or an `(x, y, sample_weight)` tuple if labels and sample weight
+    were provided. `x` will be the input images after all model preprocessing
+    has been applied.
+
+    All `ImageObjectDetectorPreprocessor` tasks include a `from_preset()`
+    constructor which can be used to load a pre-trained config and vocabularies.
+    You can call the `from_preset()` constructor directly on this base class, in
+    which case the correct class for your model will be automatically
+    instantiated.
+
+    Args:
+        image_converter: Preprocessing pipeline for images.
+
+    Examples.
+    ```python
+    preprocessor = keras_hub.models.ImageObjectDetectorPreprocessor.from_preset(
+        "retinanet_resnet50",
+    )
+    """
+
+    def __init__(
+        self,
+        image_converter=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.image_converter = image_converter
+
+    @preprocessing_function
+    def call(self, x, y=None, sample_weight=None):
+        if self.image_converter:
+            x = self.image_converter(x)
+        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)

keras_hub/src/models/image_segmenter.py

@@ -16,11 +16,6 @@ class ImageSegmenter(Task):
     be used to load a pre-trained config and weights.
     """
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        # Default compilation.
-        self.compile()
-
     def compile(
         self,
         optimizer="auto",

keras_hub/src/models/image_segmenter_preprocessor.py

@@ -19,9 +19,12 @@ class ImageSegmenterPreprocessor(Preprocessor):
 
     - `x`: The first input, should always be included. It can be an image or
       a batch of images.
-    - `y`: (Optional) Usually the segmentation mask(s), will be passed through
-      unaltered.
+    - `y`: (Optional) Usually the segmentation mask(s), if `resize_output_mask`
+        is set to `True` this will be resized to input image shape else will be
+        passed through unaltered.
     - `sample_weight`: (Optional) Will be passed through unaltered.
+    - `resize_output_mask` bool: If set to `True` the output mask will be
+      resized to the same size as the input image. Defaults to `False`.
 
     The layer will output either `x`, an `(x, y)` tuple if labels were provided,
     or an `(x, y, sample_weight)` tuple if labels and sample weight were
@@ -29,7 +32,7 @@ class ImageSegmenterPreprocessor(Preprocessor):
     been applied.
 
     All `ImageSegmenterPreprocessor` tasks include a `from_preset()`
-    constructor which can be used to load a pre-trained config and vocabularies.
+    constructor which can be used to load a pre-trained config.
     You can call the `from_preset()` constructor directly on this base class, in
     which case the correct class for your model will be automatically
     instantiated.
@@ -49,7 +52,8 @@ class ImageSegmenterPreprocessor(Preprocessor):
     x, y = preprocessor(x, y)
 
     # Resize a batch of images and masks.
-    x, y = [np.ones((512, 512, 3)), np.zeros((512, 512, 3))], [np.ones((512, 512, 1)), np.zeros((512, 512, 1))]
+    x, y = [np.ones((512, 512, 3)), np.zeros((512, 512, 3))],
+           [np.ones((512, 512, 1)), np.zeros((512, 512, 1))]
     x, y = preprocessor(x, y)
 
     # Use a `tf.data.Dataset`.
@@ -61,13 +65,34 @@ class ImageSegmenterPreprocessor(Preprocessor):
     def __init__(
         self,
         image_converter=None,
+        resize_output_mask=False,
         **kwargs,
     ):
         super().__init__(**kwargs)
         self.image_converter = image_converter
+        self.resize_output_mask = resize_output_mask
 
     @preprocessing_function
     def call(self, x, y=None, sample_weight=None):
         if self.image_converter:
             x = self.image_converter(x)
+
+        if y is not None and self.image_converter and self.resize_output_mask:
+            y = keras.layers.Resizing(
+                height=(
+                    self.image_converter.image_size[0]
+                    if self.image_converter.image_size
+                    else None
+                ),
+                width=(
+                    self.image_converter.image_size[1]
+                    if self.image_converter.image_size
+                    else None
+                ),
+                crop_to_aspect_ratio=self.image_converter.crop_to_aspect_ratio,
+                interpolation="nearest",
+                data_format=self.image_converter.data_format,
+                dtype=self.dtype_policy,
+                name="mask_resizing",
+            )(y)
         return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)

keras_hub/src/models/image_to_image.py

@@ -0,0 +1,417 @@
+import itertools
+from functools import partial
+
+import keras
+from keras import ops
+from keras import random
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.task import Task
+from keras_hub.src.utils.keras_utils import standardize_data_format
+
+try:
+    import tensorflow as tf
+except ImportError:
+    tf = None
+
+
+@keras_hub_export("keras_hub.models.ImageToImage")
+class ImageToImage(Task):
+    """Base class for image-to-image tasks.
+
+    `ImageToImage` tasks wrap a `keras_hub.models.Backbone` and
+    a `keras_hub.models.Preprocessor` to create a model that can be used for
+    generation and generative fine-tuning.
+
+    `ImageToImage` tasks provide an additional, high-level `generate()` function
+    which can be used to generate image by token with a (image, string) in,
+    image out signature.
+
+    All `ImageToImage` tasks include a `from_preset()` constructor which can be
+    used to load a pre-trained config and weights.
+
+    Example:
+
+    ```python
+    # Load a Stable Diffusion 3 backbone with pre-trained weights.
+    reference_image = np.ones((1024, 1024, 3), dtype="float32")
+    image_to_image = keras_hub.models.ImageToImage.from_preset(
+        "stable_diffusion_3_medium",
+    )
+    image_to_image.generate(
+        reference_image,
+        "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
+    )
+
+    # Load a Stable Diffusion 3 backbone at bfloat16 precision.
+    image_to_image = keras_hub.models.ImageToImage.from_preset(
+        "stable_diffusion_3_medium",
+        dtype="bfloat16",
+    )
+    image_to_image.generate(
+        reference_image,
+        "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
+    )
+    ```
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # Default compilation.
+        self.compile()
+
+    @property
+    def support_negative_prompts(self):
+        """Whether the model supports `negative_prompts` key in `generate()`."""
+        return bool(True)
+
+    @property
+    def image_shape(self):
+        return tuple(self.backbone.image_shape)
+
+    @property
+    def latent_shape(self):
+        return tuple(self.backbone.latent_shape)
+
+    def compile(
+        self,
+        optimizer="auto",
+        loss="auto",
+        *,
+        metrics="auto",
+        **kwargs,
+    ):
+        """Configures the `ImageToImage` task for training.
+
+        The `ImageToImage` task extends the default compilation signature of
+        `keras.Model.compile` with defaults for `optimizer`, `loss`, and
+        `metrics`. To override these defaults, pass any value
+        to these arguments during compilation.
+
+        Args:
+            optimizer: `"auto"`, an optimizer name, or a `keras.Optimizer`
+                instance. Defaults to `"auto"`, which uses the default optimizer
+                for the given model and task. See `keras.Model.compile` and
+                `keras.optimizers` for more info on possible `optimizer` values.
+            loss: `"auto"`, a loss name, or a `keras.losses.Loss` instance.
+                Defaults to `"auto"`, where a
+                `keras.losses.MeanSquaredError` loss will be applied. See
+                `keras.Model.compile` and `keras.losses` for more info on
+                possible `loss` values.
+            metrics: `"auto"`, or a list of metrics to be evaluated by
+                the model during training and testing. Defaults to `"auto"`,
+                where a `keras.metrics.MeanSquaredError` will be applied to
+                track the loss of the model during training. See
+                `keras.Model.compile` and `keras.metrics` for more info on
+                possible `metrics` values.
+            **kwargs: See `keras.Model.compile` for a full list of arguments
+                supported by the compile method.
+        """
+        # Ref: https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py#L410-L414
+        if optimizer == "auto":
+            optimizer = keras.optimizers.AdamW(
+                1e-4, weight_decay=1e-2, epsilon=1e-8, clipnorm=1.0
+            )
+        if loss == "auto":
+            loss = keras.losses.MeanSquaredError()
+        if metrics == "auto":
+            metrics = [keras.metrics.MeanSquaredError()]
+        super().compile(
+            optimizer=optimizer,
+            loss=loss,
+            metrics=metrics,
+            **kwargs,
+        )
+        self.generate_function = None
+
+    def generate_step(self, *args, **kwargs):
+        """Run generation on batches of input."""
+        raise NotImplementedError
+
+    def make_generate_function(self):
+        """Create or return the compiled generation function."""
+        if self.generate_function is not None:
+            return self.generate_function
+
+        self.generate_function = self.generate_step
+        if keras.config.backend() == "torch":
+            import torch
+
+            def wrapped_function(*args, **kwargs):
+                with torch.no_grad():
+                    return self.generate_step(*args, **kwargs)
+
+            self.generate_function = wrapped_function
+        elif keras.config.backend() == "tensorflow" and not self.run_eagerly:
+            self.generate_function = tf.function(
+                self.generate_step, jit_compile=self.jit_compile
+            )
+        elif keras.config.backend() == "jax" and not self.run_eagerly:
+            import jax
+
+            @partial(jax.jit)
+            def compiled_function(state, *args, **kwargs):
+                (
+                    trainable_variables,
+                    non_trainable_variables,
+                ) = state
+                mapping = itertools.chain(
+                    zip(self.trainable_variables, trainable_variables),
+                    zip(self.non_trainable_variables, non_trainable_variables),
+                )
+
+                with keras.StatelessScope(state_mapping=mapping):
+                    outputs = self.generate_step(*args, **kwargs)
+                return outputs
+
+            def wrapped_function(*args, **kwargs):
+                # Create an explicit tuple of all variable state.
+                state = (
+                    # Use the explicit variable.value to preserve the
+                    # sharding spec of distribution.
+                    [v.value for v in self.trainable_variables],
+                    [v.value for v in self.non_trainable_variables],
+                )
+                outputs = compiled_function(state, *args, **kwargs)
+                return outputs
+
+            self.generate_function = wrapped_function
+        return self.generate_function
+
+    def _normalize_generate_inputs(self, inputs):
+        """Normalize user input to the generate function.
+
+        This function converts all inputs to tensors, adds a batch dimension if
+        necessary, and returns a iterable "dataset like" object (either an
+        actual `tf.data.Dataset` or a list with a single batch element).
+
+        The input format must be one of the following:
+        - A dict with "images", "prompts" and/or "negative_prompts" keys
+        - A tf.data.Dataset with "images", "prompts" and/or "negative_prompts"
+            keys
+
+        The output will be a dict with "images", "prompts" and/or
+        "negative_prompts" keys.
+        """
+        if tf and isinstance(inputs, tf.data.Dataset):
+            _inputs = {
+                "images": inputs.map(lambda x: x["images"]).as_numpy_iterator(),
+                "prompts": inputs.map(
+                    lambda x: x["prompts"]
+                ).as_numpy_iterator(),
+            }
+            if self.support_negative_prompts:
+                _inputs["negative_prompts"] = inputs.map(
+                    lambda x: x["negative_prompts"]
+                ).as_numpy_iterator()
+            return _inputs, False
+
+        if (
+            not isinstance(inputs, dict)
+            or "images" not in inputs
+            or "prompts" not in inputs
+        ):
+            raise ValueError(
+                '`inputs` must be a dict with "images" and "prompts" keys or a'
+                f"tf.data.Dataset. Received: inputs={inputs}"
+            )
+
+        def normalize(x):
+            if isinstance(x, str):
+                return [x], True
+            if tf and isinstance(x, tf.Tensor) and x.shape.rank == 0:
+                return x[tf.newaxis], True
+            return x, False
+
+        def normalize_images(x):
+            data_format = getattr(
+                self.backbone, "data_format", standardize_data_format(None)
+            )
+            input_is_scalar = False
+            x = ops.convert_to_tensor(x)
+            if len(ops.shape(x)) < 4:
+                x = ops.expand_dims(x, axis=0)
+                input_is_scalar = True
+            x = ops.image.resize(
+                x,
+                (self.backbone.image_shape[0], self.backbone.image_shape[1]),
+                interpolation="nearest",
+                data_format=data_format,
+            )
+            return x, input_is_scalar
+
+        def get_dummy_prompts(x):
+            dummy_prompts = [""] * len(x)
+            if tf and isinstance(x, tf.Tensor):
+                return tf.convert_to_tensor(dummy_prompts)
+            else:
+                return dummy_prompts
+
+        for key in inputs:
+            if key == "images":
+                inputs[key], input_is_scalar = normalize_images(inputs[key])
+            else:
+                inputs[key], input_is_scalar = normalize(inputs[key])
+
+        if self.support_negative_prompts and "negative_prompts" not in inputs:
+            inputs["negative_prompts"] = get_dummy_prompts(inputs["prompts"])
+
+        return [inputs], input_is_scalar
+
+    def _normalize_generate_outputs(self, outputs, input_is_scalar):
+        """Normalize user output from the generate function.
+
+        This function converts all output to numpy with a value range of
+        `[0, 255]`. If a batch dimension was added to the input, it is removed
+        from the output.
+        """
+
+        def normalize(x):
+            outputs = ops.concatenate(x, axis=0)
+            outputs = ops.clip(ops.divide(ops.add(outputs, 1.0), 2.0), 0.0, 1.0)
+            outputs = ops.cast(ops.round(ops.multiply(outputs, 255.0)), "uint8")
+            outputs = ops.squeeze(outputs, 0) if input_is_scalar else outputs
+            return ops.convert_to_numpy(outputs)
+
+        if isinstance(outputs[0], dict):
+            normalized = {}
+            for key in outputs[0]:
+                normalized[key] = normalize([x[key] for x in outputs])
+            return normalized
+        return normalize([x for x in outputs])
+
+    def generate(
+        self,
+        inputs,
+        num_steps,
+        strength,
+        guidance_scale=None,
+        seed=None,
+    ):
+        """Generate image based on the provided `inputs`.
+
+        Typically, `inputs` is a dict with `"images"` and `"prompts"` keys.
+        `"images"` are reference images within a value range of
+        `[-1.0, 1.0]`, which will be resized to `self.backbone.height` and
+        `self.backbone.width`, then encoded into latent space by the VAE
+        encoder. `"prompts"` are strings that will be tokenized and encoded by
+        the text encoder.
+
+        Some models support a `"negative_prompts"` key, which helps steer the
+        model away from generating certain styles and elements. To enable this,
+        add `"negative_prompts"` to the input dict.
+
+        If `inputs` are a `tf.data.Dataset`, outputs will be generated
+        "batch-by-batch" and concatenated. Otherwise, all inputs will be
+        processed as batches.
+
+        Args:
+            inputs: python data, tensor data, or a `tf.data.Dataset`. The format
+                must be one of the following:
+                - A dict with `"images"`, `"prompts"` and/or
+                    `"negative_prompts"` keys.
+                - A `tf.data.Dataset` with `"images"`, `"prompts"` and/or
+                    `"negative_prompts"` keys.
+            num_steps: int. The number of diffusion steps to take.
+            strength: float. Indicates the extent to which the reference
+                `images` are transformed. Must be between `0.0` and `1.0`. When
+                `strength=1.0`, `images` is essentially ignore and added noise
+                is maximum and the denoising process runs for the full number of
+                iterations specified in `num_steps`.
+            guidance_scale: Optional float. The classifier free guidance scale
+                defined in [Classifier-Free Diffusion Guidance](
+                https://arxiv.org/abs/2207.12598). A higher scale encourages
+                generating images more closely related to the prompts, typically
+                at the cost of lower image quality. Note that some models don't
+                utilize classifier-free guidance.
+            seed: optional int. Used as a random seed.
+        """
+        num_steps = int(num_steps)
+        strength = float(strength)
+        guidance_scale = (
+            float(guidance_scale) if guidance_scale is not None else None
+        )
+        if strength < 0.0 or strength > 1.0:
+            raise ValueError(
+                "`strength` must be between `0.0` and `1.0`. "
+                f"Received strength={strength}."
+            )
+        if guidance_scale is not None and guidance_scale > 1.0:
+            guidance_scale = ops.convert_to_tensor(float(guidance_scale))
+        else:
+            guidance_scale = None
+        starting_step = int(num_steps * (1.0 - strength))
+        starting_step = ops.convert_to_tensor(starting_step, "int32")
+        num_steps = ops.convert_to_tensor(int(num_steps), "int32")
+
+        # Check `inputs` format.
+        required_keys = ["images", "prompts"]
+        if tf and isinstance(inputs, tf.data.Dataset):
+            spec = inputs.element_spec
+            if not all(key in spec for key in required_keys):
+                raise ValueError(
+                    "Expected a `tf.data.Dataset` with the following keys:"
+                    f"{required_keys}. Received: inputs.element_spec={spec}"
+                )
+        else:
+            if not isinstance(inputs, dict):
+                raise ValueError(
+                    "Expected a `dict` or `tf.data.Dataset`. "
+                    f"Received: inputs={inputs} of type {type(inputs)}."
+                )
+            if not all(key in inputs for key in required_keys):
+                raise ValueError(
+                    "Expected a `dict` with the following keys:"
+                    f"{required_keys}. "
+                    f"Received: inputs.keys={list(inputs.keys())}"
+                )
+
+        # Setup our three main passes.
+        # 1. Preprocessing strings to dense integer tensors.
+        # 2. Generate outputs via a compiled function on dense tensors.
+        # 3. Postprocess dense tensors to a value range of `[0, 255]`.
+        generate_function = self.make_generate_function()
+
+        def preprocess(x):
+            if self.preprocessor is not None:
+                return self.preprocessor.generate_preprocess(x)
+            else:
+                return x
+
+        def generate(images, x):
+            token_ids = x[0] if self.support_negative_prompts else x
+
+            # Initialize noises.
+            if isinstance(token_ids, dict):
+                arbitrary_key = list(token_ids.keys())[0]
+                batch_size = ops.shape(token_ids[arbitrary_key])[0]
+            else:
+                batch_size = ops.shape(token_ids)[0]
+            noise_shape = (batch_size,) + self.latent_shape[1:]
+            noises = random.normal(noise_shape, dtype="float32", seed=seed)
+
+            return generate_function(
+                images, noises, x, starting_step, num_steps, guidance_scale
+            )
+
+        # Normalize and preprocess inputs.
+        inputs, input_is_scalar = self._normalize_generate_inputs(inputs)
+        if self.support_negative_prompts:
+            images = [x["images"] for x in inputs]
+            token_ids = [preprocess(x["prompts"]) for x in inputs]
+            negative_token_ids = [
+                preprocess(x["negative_prompts"]) for x in inputs
+            ]
+            # Tuple format: (images, (token_ids, negative_token_ids)).
+            inputs = [
+                x for x in zip(images, zip(token_ids, negative_token_ids))
+            ]
+        else:
+            images = [x["images"] for x in inputs]
+            token_ids = [preprocess(x["prompts"]) for x in inputs]
+            # Tuple format: (images, token_ids).
+            inputs = [x for x in zip(images, token_ids)]
+
+        # Image-to-image.
+        outputs = [generate(*x) for x in inputs]
+        return self._normalize_generate_outputs(outputs, input_is_scalar)