Rename the key of the pyramid outputs in DINOV2.

Author: james77777778

keras_hub/src/models/depth_anything/depth_anything_backbone_test.py

@@ -29,7 +29,7 @@ def setUp(self):
             "fusion_hidden_dim": 128,
             "head_hidden_dim": 16,
             "head_in_index": -1,
-            "feature_keys": ["Stage1", "Stage2", "Stage3", "Stage4"],
+            "feature_keys": ["stage1", "stage2", "stage3", "stage4"],
         }
         self.input_data = np.ones((2, 70, 70, 3), dtype="float32")
 

keras_hub/src/models/depth_anything/depth_anything_depth_estimator_test.py

@@ -43,7 +43,7 @@ def setUp(self):
             fusion_hidden_dim=128,
             head_hidden_dim=16,
             head_in_index=-1,
-            feature_keys=["Stage1", "Stage2", "Stage3", "Stage4"],
+            feature_keys=["stage1", "stage2", "stage3", "stage4"],
         )
         self.init_kwargs = {
             "backbone": self.backbone,

keras_hub/src/models/dinov2/dinov2_backbone.py

@@ -196,7 +196,7 @@ def __init__(
         pyramid_outputs = {}
         image_input = layers.Input(shape=image_shape, name="images")
         x = self.embeddings(image_input)
-        pyramid_outputs["Stem"] = x
+        pyramid_outputs["stem"] = x
         x, encoder_pyramid_outputs = self.encoder(x)
         pyramid_outputs.update(encoder_pyramid_outputs)
         x = self.layernorm(x)

keras_hub/src/models/dinov2/dinov2_backbone_test.py

@@ -36,7 +36,7 @@ def test_backbone_basics(self):
             init_kwargs=self.init_kwargs,
             input_data=self.input_data,
             expected_output_shape=(2, sequence_length, hidden_dim),
-            expected_pyramid_output_keys=["Stem", "Stage1", "Stage2"],
+            expected_pyramid_output_keys=["stem", "stage1", "stage2"],
             expected_pyramid_image_sizes=[(sequence_length, hidden_dim)] * 3,
             run_data_format_check=False,
         )
@@ -131,7 +131,7 @@ def test_backbone_basics(self):
             init_kwargs=self.init_kwargs,
             input_data=self.input_data,
             expected_output_shape=(2, sequence_length, hidden_dim),
-            expected_pyramid_output_keys=["Stem", "Stage1", "Stage2"],
+            expected_pyramid_output_keys=["stem", "stage1", "stage2"],
             expected_pyramid_image_sizes=[(sequence_length, hidden_dim)] * 3,
             run_data_format_check=False,
         )

keras_hub/src/models/dinov2/dinov2_layers.py

@@ -870,7 +870,7 @@ def call(self, inputs, training=None):
         x = inputs
         for layer_index, layer in enumerate(self.layers, start=1):
             x = layer(x, training=training)
-            pyramid_outputs[f"Stage{str(layer_index)}"] = x
+            pyramid_outputs[f"stage{str(layer_index)}"] = x
         return x, pyramid_outputs
 
     def get_config(self):
@@ -892,5 +892,5 @@ def get_config(self):
     def compute_output_shape(self, input_shape):
         pyramid_outputs = {}
         for layer_index in range(1, len(self.layers) + 1):
-            pyramid_outputs[f"Stage{str(layer_index)}"] = input_shape
+            pyramid_outputs[f"stage{str(layer_index)}"] = input_shape
         return input_shape, pyramid_outputs

tools/checkpoint_conversion/convert_depth_anything_checkpoints.py

@@ -64,7 +64,6 @@ def convert_model(hf_model, dtype=None):
     model_config = hf_model.config.to_dict()
     # In KerasHub, the stage names are capitalized.
     feature_keys = model_config["backbone_config"]["out_features"]
-    feature_keys = [key.replace("stage", "Stage") for key in feature_keys]
     assert model_config["depth_estimation_type"] == "relative"
     assert model_config["max_depth"] in (None, 1.0)
     return DepthAnythingBackbone(