Export task models as both models.task.Model and models.Model (#2093)

* Export task models as both models.task.Model and models.Model * Same for ImageClassifier
keras-team · Oct 2, 2023 · 69de849 · 69de849
1 parent f2721af
commit 69de849
Show file tree

Hide file tree

Showing 6 changed files with 29 additions and 6 deletions.
diff --git a/keras_cv/models/classification/image_classifier.py b/keras_cv/models/classification/image_classifier.py
@@ -28,7 +28,12 @@
 from keras_cv.utils.python_utils import classproperty
 
 
-@keras_cv_export("keras_cv.models.ImageClassifier")
+@keras_cv_export(
+    [
+        "keras_cv.models.ImageClassifier",
+        "keras_cv.models.classification.ImageClassifier",
+    ]
+)
 class ImageClassifier(Task):
     """Image classifier with pooling and dense layer prediction head.
 

diff --git a/keras_cv/models/object_detection/retinanet/retinanet.py b/keras_cv/models/object_detection/retinanet/retinanet.py
@@ -41,7 +41,9 @@
 BOX_VARIANCE = [0.1, 0.1, 0.2, 0.2]
 
 
-@keras_cv_export("keras_cv.models.RetinaNet")
+@keras_cv_export(
+    ["keras_cv.models.RetinaNet", "keras_cv.models.object_detection.RetinaNet"]
+)
 class RetinaNet(Task):
     """A Keras model implementing the RetinaNet meta-architecture.
 

diff --git a/keras_cv/models/object_detection/yolo_v8/yolo_v8_detector.py b/keras_cv/models/object_detection/yolo_v8/yolo_v8_detector.py
@@ -320,7 +320,12 @@ def dist2bbox(distance, anchor_points):
     return ops.concatenate((x1y1, x2y2), axis=-1)  # xyxy bbox
 
 
-@keras_cv_export("keras_cv.models.YOLOV8Detector")
+@keras_cv_export(
+    [
+        "keras_cv.models.YOLOV8Detector",
+        "keras_cv.models.object_detection.YOLOV8Detector",
+    ]
+)
 class YOLOV8Detector(Task):
     """Implements the YOLOV8 architecture for object detection.
 

diff --git a/keras_cv/models/segmentation/deeplab_v3_plus/deeplab_v3_plus.py b/keras_cv/models/segmentation/deeplab_v3_plus/deeplab_v3_plus.py
@@ -29,7 +29,12 @@
 from keras_cv.utils.train import get_feature_extractor
 
 
-@keras_cv_export("keras_cv.models.DeepLabV3Plus")
+@keras_cv_export(
+    [
+        "keras_cv.models.DeepLabV3Plus",
+        "keras_cv.models.segmentation.DeepLabV3Plus",
+    ]
+)
 class DeepLabV3Plus(Task):
     """A Keras model implementing the DeepLabV3+ architecture for semantic
     segmentation.

diff --git a/keras_cv/models/segmentation/segformer/segformer.py b/keras_cv/models/segmentation/segformer/segformer.py
@@ -27,7 +27,9 @@
 from keras_cv.utils.train import get_feature_extractor
 
 
-@keras_cv_export("keras_cv.models.segmentation.SegFormer")
+@keras_cv_export(
+    ["keras_cv.models.SegFormer", "keras_cv.models.segmentation.SegFormer"]
+)
 class SegFormer(Task):
     """A Keras model implementing the SegFormer architecture for semantic
     segmentation.

diff --git a/keras_cv/models/segmentation/segment_anything/sam.py b/keras_cv/models/segmentation/segment_anything/sam.py
@@ -31,7 +31,11 @@
 
 
 @keras_cv_export(
-    "keras_cv.models.SegmentAnythingModel", package="keras_cv.models"
+    [
+        "keras_cv.models.SegmentAnythingModel",
+        "keras_cv.models.segmentation.SegmentAnythingModel",
+    ],
+    package="keras_cv.models",
 )
 class SegmentAnythingModel(Task):
     """