keras-hub-nightly 0.23.0.dev202509190415__py3-none-any.whl → 0.23.0.dev202509290422__py3-none-any.whl

keras_hub/src/models/depth_anything/depth_anything_loss.py

@@ -0,0 +1,89 @@
+import keras
+from keras import ops
+from keras.src.losses.losses import LossFunctionWrapper
+
+
+class DepthAnythingLoss(LossFunctionWrapper):
+    """Computes the DepthAnything loss between `y_true` & `y_pred`.
+
+    This loss is the Scale-Invariant Logarithmic (SiLog) loss, which is
+    widely used for depth estimation tasks.
+
+    See: [Depth Map Prediction from a Single Image using a Multi-Scale Deep Network](https://arxiv.org/abs/1406.2283)
+
+    Args:
+        lambd: The weighting factor in the scale-invariant log loss formula.
+            Defaults to `0.5`.
+        min_depth: Minimum depth value used to filter `y_pred` and `y_true`.
+            Defaults to `keras.config.epsilon()`.
+        max_depth: Optional maximum depth value used to filter `y_pred` and
+            `y_true`. If not specified, there will be no upper bound.
+        reduction: Type of reduction to apply to the loss. In almost all cases
+            this should be `"sum_over_batch_size"`. Supported options are
+            `"sum"`, `"sum_over_batch_size"`, `"mean"`,
+            `"mean_with_sample_weight"` or `None`. `"sum"` sums the loss,
+            `"sum_over_batch_size"` and `"mean"` sum the loss and divide by the
+            sample size, and `"mean_with_sample_weight"` sums the loss and
+            divides by the sum of the sample weights. `"none"` and `None`
+            perform no aggregation. Defaults to `"sum_over_batch_size"`.
+        name: Optional name for the instance.
+        dtype: The dtype of the loss's computations. Defaults to `None`, which
+            means using `keras.backend.floatx()`. `keras.backend.floatx()` is a
+            `"float32"` unless set to different value
+            (via `keras.backend.set_floatx()`). If a `keras.DTypePolicy` is
+            provided, then the `compute_dtype` will be utilized.
+    """
+
+    def __init__(
+        self,
+        lambd=0.5,
+        min_depth=keras.config.epsilon(),
+        max_depth=None,
+        reduction="sum_over_batch_size",
+        name="depth_anything_loss",
+        dtype=None,
+    ):
+        super().__init__(
+            silog,
+            name=name,
+            reduction=reduction,
+            dtype=dtype,
+            lambd=lambd,
+            min_depth=min_depth,
+            max_depth=max_depth,
+        )
+
+
+def silog(
+    y_true, y_pred, lambd=0.5, min_depth=keras.config.epsilon(), max_depth=None
+):
+    y_pred = ops.convert_to_tensor(y_pred)
+    y_true = ops.convert_to_tensor(y_true, dtype=y_pred.dtype)
+
+    # Apply the valid mask.
+    if max_depth is None:
+        valid_mask = ops.greater_equal(y_true, min_depth)
+    else:
+        valid_mask = ops.logical_and(
+            ops.greater_equal(y_true, min_depth),
+            ops.less_equal(y_true, max_depth),
+        )
+    y_true = ops.multiply(y_true, valid_mask)
+    y_pred = ops.multiply(y_pred, valid_mask)
+
+    diff_log = ops.where(
+        valid_mask,
+        ops.subtract(ops.log(y_true), ops.log(y_pred)),
+        ops.zeros_like(y_true),
+    )
+
+    divisor = ops.sum(ops.cast(valid_mask, y_true.dtype), axis=(1, 2, 3))
+    mean_power2_diff_log = ops.divide_no_nan(
+        ops.sum(ops.power(diff_log, 2), axis=(1, 2, 3)), divisor
+    )
+    power2_mean_diff_log = ops.power(
+        ops.divide_no_nan(ops.sum(diff_log, axis=(1, 2, 3)), divisor), 2
+    )
+    return ops.sqrt(
+        mean_power2_diff_log - ops.multiply(lambd, power2_mean_diff_log)
+    )

keras_hub/src/models/depth_anything/depth_anything_presets.py

@@ -0,0 +1,4 @@
+"""DepthAnything model preset configurations."""
+
+# Metadata for loading pretrained model weights.
+backbone_presets = {}

keras_hub/src/models/depth_anything/interpolate.py

@@ -0,0 +1,62 @@
+from keras import backend
+from keras import ops
+
+from keras_hub.src.utils.keras_utils import standardize_data_format
+
+
+def interpolate(x, size, data_format=None):
+    """Performs a backend-agnostic version of Torch's `F.interpolate`.
+
+    Args:
+        x: A 4D image tensor.
+        size: A tuple of 2 integers, `(height, width)`.
+        data_format: One of `channels_last` or `channels_first`.
+    """
+    data_format = standardize_data_format(data_format)
+    if backend.backend() == "jax":
+        import jax
+
+        if data_format == "channels_first":
+            x = ops.transpose(x, (0, 2, 3, 1))
+        scale = ops.convert_to_tensor(
+            [
+                (size[0] - 1.0) / (x.shape[1] - 1.0),
+                (size[1] - 1.0) / (x.shape[2] - 1.0),
+            ]
+        )
+        translation = -(scale / 2.0 - 0.5)
+        x = jax.image.scale_and_translate(
+            x,
+            (x.shape[0], *size, x.shape[-1]),
+            method="bilinear",
+            scale=scale,
+            spatial_dims=(1, 2),
+            translation=translation,
+            antialias=False,
+        )
+        if data_format == "channels_first":
+            x = ops.transpose(x, (0, 3, 1, 2))
+    elif backend.backend() == "tensorflow":
+        import tensorflow as tf
+
+        if data_format == "channels_first":
+            x = ops.transpose(x, (0, 2, 3, 1))
+        x = tf.compat.v1.image.resize(
+            x,
+            size=size,
+            method="bilinear",
+            align_corners=True,
+        )
+        if data_format == "channels_first":
+            x = ops.transpose(x, (0, 3, 1, 2))
+    elif backend.backend() == "torch":
+        import torch.nn.functional as F
+
+        if data_format == "channels_last":
+            x = ops.transpose(x, (0, 3, 1, 2))
+        x = F.interpolate(x, size=size, mode="bilinear", align_corners=True)
+        if data_format == "channels_last":
+            x = ops.transpose(x, (0, 2, 3, 1))
+    else:
+        raise NotImplementedError(f"Unsupported backend: {backend.backend()}")
+    return x

keras_hub/src/models/depth_estimator.py

@@ -0,0 +1,239 @@
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.task import Task
+
+
+class Multiplier(keras.layers.Layer):
+    def __init__(self, multiplier=None, **kwargs):
+        super().__init__(**kwargs)
+        self.multiplier = float(multiplier) if multiplier is not None else None
+
+    def call(self, inputs):
+        if self.multiplier is not None:
+            inputs = keras.ops.multiply(inputs, self.multiplier)
+        return inputs
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "multiplier": self.multiplier,
+            }
+        )
+        return config
+
+
+@keras_hub_export("keras_hub.models.DepthEstimator")
+class DepthEstimator(Task):
+    """Base class for all depth estimation tasks.
+
+    `DepthEstimator` tasks wrap a `keras_hub.models.Backbone` and
+    a `keras_hub.models.Preprocessor` to create a model that can be used for
+    depth estimation.
+
+    To fine-tune with `fit()`, pass a dataset containing tuples of `(x, y)`
+    labels where `x` is a RGB image and `y` is a depth map. All `DepthEstimator`
+    tasks include a `from_preset()` constructor which can be used to load a
+    pre-trained config and weights.
+
+    Args:
+        backbone: A `keras_hub.models.Backbone` instance or a `keras.Model`.
+        preprocessor: `None`, a `keras_hub.models.Preprocessor` instance,
+            a `keras.Layer` instance, or a callable. If `None` no preprocessing
+            will be applied to the inputs.
+        depth_estimation_type: `"relative"` or `"metric"`. The type of depth map
+            to use. `"relative"` depth maps are up-to-scale, while `"metric"`
+            depth maps have metric meaning (e.g. in meters). Defaults to
+            `"relative"`.
+        min_depth: An float representing the minimum depth value. This value can
+            be used to filter out invalid depth values during training. Defaults
+            to `keras.config.epsilon()`.
+        max_depth: An optional float representing the maximum depth value. This
+            value can be used to filter out invalid depth values during
+            training. When `depth_estimation_type="metric"`, the model's output
+            will be scaled to the range `[0, max_depth]`.
+
+    Examples:
+
+    Call `predict()` to run inference.
+    ```python
+    # Load preset and train
+    images = np.random.randint(0, 256, size=(2, 224, 224, 3))
+    depth_estimator = keras_hub.models.DepthEstimator.from_preset(
+        "depth_anything_v2_small"
+    )
+    depth_estimator.predict(images)
+    ```
+
+    Call `fit()` on a single batch.
+    ```python
+    # Load preset and train
+    images = np.random.randint(0, 256, size=(2, 224, 224, 3))
+    depths = np.random.uniform(0, 10, size=(2, 224, 224))
+    depth_estimator = keras_hub.models.DepthEstimator.from_preset(
+        "depth_anything_v2_small",
+        depth_estimation_type="metric",
+        max_depth=10.0,
+    )
+    depth_estimator.fit(x=images, y=depths, batch_size=2)
+    ```
+
+    Call `fit()` with custom loss, optimizer and backbone.
+    ```python
+    depth_estimator = keras_hub.models.DepthEstimator.from_preset(
+        "depth_anything_v2_small",
+        depth_estimation_type="metric",
+        max_depth=10.0,
+    )
+    depth_estimator.compile(
+        loss=keras.losses.MeanSquaredError(),
+        optimizer=keras.optimizers.Adam(5e-5),
+    )
+    depth_estimator.backbone.trainable = False
+    depth_estimator.fit(x=images, y=depths, batch_size=2)
+    ```
+
+    Custom backbone.
+    ```python
+    images = np.random.randint(0, 256, size=(2, 224, 224, 3))
+    depths = np.random.uniform(0, 10, size=(2, 224, 224))
+    image_encoder = keras_hub.models.DINOV2Backbone.from_preset("dinov2_small")
+    backbone = keras_hub.models.DepthAnythingBackbone(
+        image_encoder=image_encoder,
+        patch_size=image_encoder.patch_size,
+        backbone_hidden_dim=image_encoder.hidden_dim,
+        reassemble_factors=[4, 2, 1, 0.5],
+        neck_hidden_dims=[48, 96, 192, 384],
+        fusion_hidden_dim=64,
+        head_hidden_dim=32,
+        head_in_index=-1,
+    )
+    depth_estimator = keras_hub.models.DepthEstimator(
+        backbone=backbone,
+        depth_estimation_type="metric",
+        max_depth=10.0,
+    )
+    depth_estimator.fit(x=images, y=depths, batch_size=2)
+    ```
+    """
+
+    def __init__(
+        self,
+        backbone,
+        depth_estimation_type,
+        min_depth=keras.config.epsilon(),
+        max_depth=None,
+        preprocessor=None,
+        **kwargs,
+    ):
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        if depth_estimation_type == "relative":
+            self.output_activation = keras.layers.ReLU(
+                dtype=backbone.dtype_policy,
+                name="output_activation",
+            )
+        elif depth_estimation_type == "metric":
+            self.output_activation = keras.layers.Activation(
+                activation="sigmoid",
+                dtype=backbone.dtype_policy,
+                name="output_activation",
+            )
+        else:
+            raise ValueError(
+                "`depth_estimation_type` should be either `'relative'` or "
+                "`'metric'`. "
+                f"Received: depth_estimation_type={depth_estimation_type}."
+            )
+        if max_depth is not None and depth_estimation_type != "metric":
+            raise ValueError(
+                "`max_depth` should only be set when "
+                "`depth_estimation_type='metric'`. "
+                f"Received: depth_estimation_type={depth_estimation_type}, "
+                f"max_depth={max_depth}."
+            )
+        self.multiplier = Multiplier(
+            multiplier=max_depth, dtype=backbone.dtype_policy, name="multiplier"
+        )
+        self.depths = keras.layers.Identity(
+            dtype=backbone.dtype_policy, name="depths"
+        )
+
+        # === Config ===
+        self.depth_estimation_type = depth_estimation_type
+        self.min_depth = float(min_depth) if min_depth is not None else None
+        self.max_depth = float(max_depth) if max_depth is not None else None
+
+        # === Functional Model ===
+        inputs = self.backbone.input
+        depths = self.backbone(inputs)
+        depths = self.output_activation(depths)
+        depths = self.multiplier(depths)
+        depths = self.depths(depths)
+        outputs = {"depths": depths}
+        super().__init__(
+            inputs=inputs,
+            outputs=outputs,
+            **kwargs,
+        )
+
+    def get_config(self):
+        # Backbone serialized in `super`
+        config = super().get_config()
+        config.update(
+            {
+                "depth_estimation_type": self.depth_estimation_type,
+                "min_depth": self.min_depth,
+                "max_depth": self.max_depth,
+            }
+        )
+        return config
+
+    def compile(
+        self,
+        optimizer="auto",
+        loss="auto",
+        *,
+        metrics="auto",
+        **kwargs,
+    ):
+        """Configures the `DepthEstimator` task for training.
+
+        The `DepthEstimator` 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 for the depth estimation task. See
+                `keras.Model.compile` and `keras.losses` for more info on
+                possible `loss` values.
+            metrics: `"auto"`, or a dict of metrics to be evaluated by
+                the model during training and testing. Defaults to `"auto"`,
+                where a `keras.metrics.RootMeanSquaredError` will be applied to
+                track the accuracy 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.
+        """
+        if optimizer == "auto":
+            optimizer = keras.optimizers.AdamW(5e-5)
+        if loss == "auto":
+            loss = {"depths": keras.losses.MeanSquaredError()}
+        if metrics == "auto":
+            metrics = {"depths": keras.metrics.RootMeanSquaredError()}
+        super().compile(
+            optimizer=optimizer,
+            loss=loss,
+            metrics=metrics,
+            **kwargs,
+        )

keras_hub/src/models/depth_estimator_preprocessor.py

@@ -0,0 +1,78 @@
+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.DepthEstimatorPreprocessor")
+class DepthEstimatorPreprocessor(Preprocessor):
+    """Base class for depth estimation preprocessing layers.
+
+    `DepthEstimatorPreprocessor` tasks wraps a
+    `keras_hub.layers.ImageConverter` to create a preprocessing layer for
+    depth estimation tasks. It is intended to be paired with a
+    `keras_hub.models.DepthEstimator` task.
+
+    All `DepthEstimatorPreprocessor` take inputs 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 depths, and `sample_weight` will not be provided.
+
+    The layer will output either `x`, an `(x, y)` tuple if depths were provided,
+    or an `(x, y, sample_weight)` tuple if depths and sample weight were
+    provided. `x` will be the input images after all model preprocessing has
+    been applied.
+
+    All `DepthEstimatorPreprocessor` tasks include a `from_preset()`
+    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.
+
+    Examples.
+    ```python
+    preprocessor = keras_hub.models.DepthEstimatorPreprocessor.from_preset(
+        "depth_anything_v2_small",
+    )
+
+    # Resize a single image for DepthAnythingV2 Small.
+    x = np.random.randint(0, 256, (512, 512, 3))
+    x = preprocessor(x)
+
+    # Resize a labeled image.
+    x = np.random.randint(0, 256, (512, 512, 3))
+    y = np.random.uniform(0, 10, size=(512, 512))
+    x, y = preprocessor(x, y)
+
+    # Resize a batch of labeled images.
+    x = [
+        np.random.randint(0, 256, (512, 512, 3)),
+        np.zeros((512, 512, 3)),
+    ]
+    y = [
+        np.random.uniform(0, 10, size=(512, 512)),
+        np.random.uniform(0, 10, size=(512, 512)),
+    ]
+    x, y = preprocessor(x, y)
+
+    # Use a `tf.data.Dataset`.
+    ds = tf.data.Dataset.from_tensor_slices((x, y)).batch(2)
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
+    ```
+    """
+
+    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/dinov2/dinov2_backbone.py

@@ -1,14 +1,14 @@
 from keras import layers
 
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.models.backbone import Backbone
 from keras_hub.src.models.dinov2.dinov2_layers import DINOV2Embedding
 from keras_hub.src.models.dinov2.dinov2_layers import DINOV2Encoder
+from keras_hub.src.models.feature_pyramid_backbone import FeaturePyramidBackbone
 from keras_hub.src.utils.keras_utils import standardize_data_format
 
 
 @keras_hub_export("keras_hub.models.DINOV2Backbone")
-class DINOV2Backbone(Backbone):
+class DINOV2Backbone(FeaturePyramidBackbone):
     """DINOV2 core network with hyperparameters.
 
     DINOV2 offers a powerful, generalist visual backbone learned entirely from
@@ -19,6 +19,10 @@ class DINOV2Backbone(Backbone):
     DINOV2 model with any number of layers, heads, and embedding dimensions. To
     load preset architectures and weights, use the `from_preset` constructor.
 
+    Note that this backbone is a Feature Pyramid Backbone that can output
+    intermediate feature maps from different stages of the model. See the
+    example below for how to access these feature pyramid outputs.
+
     Note that this backbone supports interpolation of the position embeddings
     to the input image shape. This is useful when the input image shape is
     different from the shape used to train the position embeddings. The
@@ -50,6 +54,8 @@ class DINOV2Backbone(Backbone):
             embeddings to the actual input shape. Defaults to `(518, 518)`.
         antialias_in_interpolation: bool. Whether to use antialiasing in the
             interpolation of the position embeddings. Defaults to `False`.
+        apply_layernorm: bool. Whether to apply layer normalization to the
+            outputs of each stage in the feature pyramid. Defaults to `False`.
         data_format: `None` or str. If specified, either `"channels_last"` or
             `"channels_first"`. The ordering of the dimensions in the
             inputs. `"channels_last"` corresponds to inputs with shape
@@ -95,6 +101,16 @@ class DINOV2Backbone(Backbone):
         position_embedding_shape=(518, 518),
     )
     model(input_data)
+
+    # Accessing feature pyramid outputs.
+    backbone = keras_hub.models.DINOV2Backbone.from_preset(
+        "dinov2_base", image_shape=(224, 224, 3)
+    )
+    model = keras.Model(
+        inputs=backbone.inputs,
+        outputs=backbone.pyramid_outputs,
+    )
+    features = model(input_data)
     ```
     """
 
@@ -114,6 +130,7 @@ class DINOV2Backbone(Backbone):
         image_shape=(224, 224, 3),
         position_embedding_shape=(518, 518, 3),
         antialias_in_interpolation=False,
+        apply_layernorm=False,
         data_format=None,
         dtype=None,
         name=None,
@@ -176,10 +193,16 @@ class DINOV2Backbone(Backbone):
         )
 
         # === Functional Model ===
+        pyramid_outputs = {}
         image_input = layers.Input(shape=image_shape, name="images")
         x = self.embeddings(image_input)
-        x = self.encoder(x)
+        pyramid_outputs["stem"] = x
+        x, encoder_pyramid_outputs = self.encoder(x)
+        pyramid_outputs.update(encoder_pyramid_outputs)
         x = self.layernorm(x)
+        if apply_layernorm:
+            for key in pyramid_outputs:
+                pyramid_outputs[key] = self.layernorm(pyramid_outputs[key])
         outputs = x
         super().__init__(
             inputs={"images": image_input},
@@ -204,6 +227,8 @@ class DINOV2Backbone(Backbone):
         self.image_shape = image_shape
         self.position_embedding_shape = position_embedding_shape
         self.antialias_in_interpolation = bool(antialias_in_interpolation)
+        self.apply_layernorm = apply_layernorm
+        self.pyramid_outputs = pyramid_outputs
 
     def get_config(self):
         config = super().get_config()
@@ -223,6 +248,7 @@ class DINOV2Backbone(Backbone):
                 "image_shape": self.image_shape,
                 "position_embedding_shape": self.position_embedding_shape,
                 "antialias_in_interpolation": self.antialias_in_interpolation,
+                "apply_layernorm": self.apply_layernorm,
             }
         )
         return config

keras_hub/src/models/dinov2/dinov2_layers.py

@@ -1,3 +1,4 @@
+import keras
 from keras import backend
 from keras import config
 from keras import initializers
@@ -290,6 +291,10 @@ class DINOV2Embedding(layers.Layer):
                 output_shape[1] = 1 + self.num_register_tokens + patch_num**2
         return output_shape
 
+    def compute_output_spec(self, inputs):
+        output_shape = self.compute_output_shape(inputs.shape)
+        return keras.KerasTensor(output_shape, dtype=self.compute_dtype)
+
     @staticmethod
     def _interpolate_position_embeddings(
         position_embeddings,
@@ -861,10 +866,12 @@ class DINOV2Encoder(layers.Layer):
             input_shape = layer.compute_output_shape(input_shape)
 
     def call(self, inputs, training=None):
+        pyramid_outputs = {}
         x = inputs
-        for layer in self.layers:
+        for layer_index, layer in enumerate(self.layers, start=1):
             x = layer(x, training=training)
-        return x
+            pyramid_outputs[f"stage{str(layer_index)}"] = x
+        return x, pyramid_outputs
 
     def get_config(self):
         config = super().get_config()
@@ -883,4 +890,7 @@ class DINOV2Encoder(layers.Layer):
         return config
 
     def compute_output_shape(self, input_shape):
-        return input_shape
+        pyramid_outputs = {}
+        for layer_index in range(1, len(self.layers) + 1):
+            pyramid_outputs[f"stage{str(layer_index)}"] = input_shape
+        return input_shape, pyramid_outputs