[WIP] Make Backbone parent classes to allow handling all subclass presets and support extensions

keras_cv/models/backbones/backbone.py

@@ -15,10 +15,11 @@
 
 from keras_cv.api_export import keras_cv_export
 from keras_cv.backend import keras
-from keras_cv.utils.preset_utils import check_preset_class
+from keras_cv.utils.preset_utils import check_config_class
+from keras_cv.utils.preset_utils import list_presets
+from keras_cv.utils.preset_utils import list_subclasses
 from keras_cv.utils.preset_utils import load_from_preset
 from keras_cv.utils.python_utils import classproperty
-from keras_cv.utils.python_utils import format_docstring
 
 
 @keras_cv_export("keras_cv.models.Backbone")
@@ -64,12 +65,18 @@ def from_config(cls, config):
     @classproperty
     def presets(cls):
         """Dictionary of preset names and configs."""
-        return {}
+        presets = list_presets(cls)
+        for subclass in list_subclasses(cls):
+            presets.update(subclass.presets)
+        return presets
 
     @classproperty
     def presets_with_weights(cls):
         """Dictionary of preset names and configs that include weights."""
-        return {}
+        presets = list_presets(cls, with_weights=True)
+        for subclass in list_subclasses(cls):
+            presets.update(subclass.presets)
+        return presets
 
     @classproperty
     def presets_without_weights(cls):
@@ -109,47 +116,19 @@ def from_preset(
             load_weights=False,
         ```
         """
-        # We support short IDs for official presets, e.g. `"bert_base_en"`.
-        # Map these to a Kaggle Models handle.
-        if preset in cls.presets:
-            preset = cls.presets[preset]["kaggle_handle"]
-
-        check_preset_class(preset, cls)
+        preset_cls = check_config_class(preset)
+        if not issubclass(preset_cls, cls):
+            raise ValueError(
+                f"Preset has type `{preset_cls.__name__}` which is not a "
+                f"a subclass of calling class `{cls.__name__}`. Call "
+                f"`from_preset` directly on `{preset_cls.__name__}` instead."
+            )
         return load_from_preset(
             preset,
             load_weights=load_weights,
             config_overrides=kwargs,
         )
 
-    def __init_subclass__(cls, **kwargs):
-        # Use __init_subclass__ to set up a correct docstring for from_preset.
-        super().__init_subclass__(**kwargs)
-
-        # If the subclass does not define from_preset, assign a wrapper so that
-        # each class can have a distinct docstring.
-        if "from_preset" not in cls.__dict__:
-
-            def from_preset(calling_cls, *args, **kwargs):
-                return super(cls, calling_cls).from_preset(*args, **kwargs)
-
-            cls.from_preset = classmethod(from_preset)
-
-        if not cls.presets:
-            cls.from_preset.__func__.__doc__ = """Not implemented.
-
-            No presets available for this class.
-            """
-
-        # Format and assign the docstring unless the subclass has overridden it.
-        if cls.from_preset.__doc__ is None:
-            cls.from_preset.__func__.__doc__ = Backbone.from_preset.__doc__
-            format_docstring(
-                model_name=cls.__name__,
-                example_preset_name=next(iter(cls.presets_with_weights), ""),
-                preset_names='", "'.join(cls.presets),
-                preset_with_weights_names='", "'.join(cls.presets_with_weights),
-            )(cls.from_preset.__func__)
-
     @property
     def pyramid_level_inputs(self):
         """Intermediate model outputs for feature extraction.

keras_cv/models/backbones/backbone_presets.py

@@ -31,8 +31,10 @@
 from keras_cv.models.backbones.video_swin import video_swin_backbone_presets
 from keras_cv.models.backbones.vit_det import vit_det_backbone_presets
 from keras_cv.models.object_detection.yolo_v8 import yolo_v8_backbone_presets
+from keras_cv.models.object_detection_3d import center_pillar_backbone_presets
 
 backbone_presets_no_weights = {
+    **center_pillar_backbone_presets.backbone_presets_no_weights,
     **resnet_v1_backbone_presets.backbone_presets_no_weights,
     **resnet_v2_backbone_presets.backbone_presets_no_weights,
     **mobilenet_v3_backbone_presets.backbone_presets_no_weights,
@@ -47,6 +49,7 @@
 }
 
 backbone_presets_with_weights = {
+    **center_pillar_backbone_presets.backbone_presets_with_weights,
     **resnet_v1_backbone_presets.backbone_presets_with_weights,
     **resnet_v2_backbone_presets.backbone_presets_with_weights,
     **mobilenet_v3_backbone_presets.backbone_presets_with_weights,

keras_cv/models/backbones/csp_darknet/__init__.py

@@ -11,3 +11,43 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+from keras_cv.models.backbones.csp_darknet.csp_darknet_aliases import (
+    CSPDarkNetLBackbone,
+)
+from keras_cv.models.backbones.csp_darknet.csp_darknet_aliases import (
+    CSPDarkNetTinyBackbone,
+)
+from keras_cv.models.backbones.csp_darknet.csp_darknet_backbone import (
+    CSPDarkNetBackbone,
+)
+from keras_cv.models.backbones.csp_darknet.csp_darknet_backbone_presets import (
+    backbone_presets_no_weights,
+)
+from keras_cv.models.backbones.csp_darknet.csp_darknet_backbone_presets import (
+    backbone_presets_with_weights,
+)
+from keras_cv.utils.preset_utils import register_preset
+from keras_cv.utils.preset_utils import register_presets
+
+register_presets(
+    backbone_presets_no_weights, (CSPDarkNetBackbone,), with_weights=False
+)
+register_presets(
+    backbone_presets_with_weights, (CSPDarkNetBackbone,), with_weights=True
+)
+register_presets(
+    backbone_presets_with_weights, (CSPDarkNetBackbone,), with_weights=True
+)
+register_preset(
+    "csp_darknet_tiny_imagenet",
+    backbone_presets_with_weights["csp_darknet_tiny_imagenet"],
+    (CSPDarkNetTinyBackbone,),
+    with_weights=True,
+)
+register_preset(
+    "csp_darknet_l_imagenet",
+    backbone_presets_with_weights["csp_darknet_l_imagenet"],
+    (CSPDarkNetLBackbone,),
+    with_weights=True,
+)

keras_cv/models/backbones/csp_darknet/csp_darknet_aliases.py

@@ -12,16 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import copy
 
 from keras_cv.api_export import keras_cv_export
 from keras_cv.models.backbones.csp_darknet.csp_darknet_backbone import (
     CSPDarkNetBackbone,
 )
-from keras_cv.models.backbones.csp_darknet.csp_darknet_backbone_presets import (
-    backbone_presets,
-)
-from keras_cv.utils.python_utils import classproperty
 
 ALIAS_DOCSTRING = """CSPDarkNetBackbone model with {stackwise_channels} channels
     and {stackwise_depth} depths.
@@ -71,21 +66,6 @@ def __new__(
         )
         return CSPDarkNetBackbone.from_preset("csp_darknet_tiny", **kwargs)
 
-    @classproperty
-    def presets(cls):
-        """Dictionary of preset names and configurations."""
-        return {
-            "csp_darknet_tiny_imagenet": copy.deepcopy(
-                backbone_presets["csp_darknet_tiny_imagenet"]
-            )
-        }
-
-    @classproperty
-    def presets_with_weights(cls):
-        """Dictionary of preset names and configurations that include
-        weights."""
-        return cls.presets
-
 
 @keras_cv_export("keras_cv.models.CSPDarkNetSBackbone")
 class CSPDarkNetSBackbone(CSPDarkNetBackbone):
@@ -106,17 +86,6 @@ def __new__(
         )
         return CSPDarkNetBackbone.from_preset("csp_darknet_s", **kwargs)
 
-    @classproperty
-    def presets(cls):
-        """Dictionary of preset names and configurations."""
-        return {}
-
-    @classproperty
-    def presets_with_weights(cls):
-        """Dictionary of preset names and configurations that include
-        weights."""
-        return {}
-
 
 @keras_cv_export("keras_cv.models.CSPDarkNetMBackbone")
 class CSPDarkNetMBackbone(CSPDarkNetBackbone):
@@ -137,17 +106,6 @@ def __new__(
         )
         return CSPDarkNetBackbone.from_preset("csp_darknet_m", **kwargs)
 
-    @classproperty
-    def presets(cls):
-        """Dictionary of preset names and configurations."""
-        return {}
-
-    @classproperty
-    def presets_with_weights(cls):
-        """Dictionary of preset names and configurations that include
-        weights."""
-        return {}
-
 
 @keras_cv_export("keras_cv.models.CSPDarkNetLBackbone")
 class CSPDarkNetLBackbone(CSPDarkNetBackbone):
@@ -168,21 +126,6 @@ def __new__(
         )
         return CSPDarkNetBackbone.from_preset("csp_darknet_l", **kwargs)
 
-    @classproperty
-    def presets(cls):
-        """Dictionary of preset names and configurations."""
-        return {
-            "csp_darknet_l_imagenet": copy.deepcopy(
-                backbone_presets["csp_darknet_l_imagenet"]
-            )
-        }
-
-    @classproperty
-    def presets_with_weights(cls):
-        """Dictionary of preset names and configurations that include
-        weights."""
-        return cls.presets
-
 
 @keras_cv_export("keras_cv.models.CSPDarkNetXLBackbone")
 class CSPDarkNetXLBackbone(CSPDarkNetBackbone):
@@ -203,17 +146,6 @@ def __new__(
         )
         return CSPDarkNetBackbone.from_preset("csp_darknet_xl", **kwargs)
 
-    @classproperty
-    def presets(cls):
-        """Dictionary of preset names and configurations."""
-        return {}
-
-    @classproperty
-    def presets_with_weights(cls):
-        """Dictionary of preset names and configurations that include
-        weights."""
-        return {}
-
 
 setattr(
     CSPDarkNetTinyBackbone,

keras_cv/models/backbones/csp_darknet/csp_darknet_backbone.py

@@ -13,18 +13,11 @@
 # limitations under the License.
 
 """CSPDarkNet backbone model. """
-import copy
 
 from keras_cv.api_export import keras_cv_export
 from keras_cv.backend import keras
 from keras_cv.models import utils
 from keras_cv.models.backbones.backbone import Backbone
-from keras_cv.models.backbones.csp_darknet.csp_darknet_backbone_presets import (
-    backbone_presets,
-)
-from keras_cv.models.backbones.csp_darknet.csp_darknet_backbone_presets import (
-    backbone_presets_with_weights,
-)
 from keras_cv.models.backbones.csp_darknet.csp_darknet_utils import (
     CrossStagePartial,
 )
@@ -38,7 +31,6 @@
 from keras_cv.models.backbones.csp_darknet.csp_darknet_utils import (
     SpatialPyramidPoolingBottleneck,
 )
-from keras_cv.utils.python_utils import classproperty
 
 
 @keras_cv_export("keras_cv.models.CSPDarkNetBackbone")
@@ -169,14 +161,3 @@ def get_config(self):
             }
         )
         return config
-
-    @classproperty
-    def presets(cls):
-        """Dictionary of preset names and configurations."""
-        return copy.deepcopy(backbone_presets)
-
-    @classproperty
-    def presets_with_weights(cls):
-        """Dictionary of preset names and configurations that include
-        weights."""
-        return copy.deepcopy(backbone_presets_with_weights)

keras_cv/models/backbones/densenet/__init__.py

@@ -11,3 +11,49 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+from keras_cv.models.backbones.densenet.densenet_aliases import (
+    DenseNet121Backbone,
+)
+from keras_cv.models.backbones.densenet.densenet_aliases import (
+    DenseNet169Backbone,
+)
+from keras_cv.models.backbones.densenet.densenet_aliases import (
+    DenseNet201Backbone,
+)
+from keras_cv.models.backbones.densenet.densenet_backbone import (
+    DenseNetBackbone,
+)
+from keras_cv.models.backbones.densenet.densenet_backbone_presets import (
+    backbone_presets_no_weights,
+)
+from keras_cv.models.backbones.densenet.densenet_backbone_presets import (
+    backbone_presets_with_weights,
+)
+from keras_cv.utils.preset_utils import register_preset
+from keras_cv.utils.preset_utils import register_presets
+
+register_presets(
+    backbone_presets_no_weights, (DenseNetBackbone,), with_weights=False
+)
+register_presets(
+    backbone_presets_with_weights, (DenseNetBackbone,), with_weights=True
+)
+register_preset(
+    "densenet121_imagenet",
+    backbone_presets_with_weights["densenet121_imagenet"],
+    (DenseNet121Backbone,),
+    with_weights=True,
+)
+register_preset(
+    "densenet169_imagenet",
+    backbone_presets_with_weights["densenet169_imagenet"],
+    (DenseNet169Backbone,),
+    with_weights=True,
+)
+register_preset(
+    "densenet201_imagenet",
+    backbone_presets_with_weights["densenet201_imagenet"],
+    (DenseNet201Backbone,),
+    with_weights=True,
+)