Simplify running KerasCV with Keras 3

README.md

@@ -7,10 +7,10 @@
 [![Contributions Welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg?style=flat)](https://github.com/keras-team/keras-cv/issues)
 
 KerasCV is a library of modular computer vision components that work natively
-with TensorFlow, JAX, or PyTorch. Built on [Keras Core](https://keras.io/keras_core/announcement/),
-these models, layers, metrics, callbacks, etc., can be trained and serialized
-in any framework and re-used in another without costly migrations. See
-"Configuring your backend" below for more details on multi-framework KerasCV.
+with TensorFlow, JAX, or PyTorch. Built on Keras 3, these models, layers, 
+metrics, callbacks, etc., can be trained and serialized in any framework and 
+re-used in another without costly migrations. See "Configuring your backend" 
+below for more details on multi-framework KerasCV.
 
 <img style="width: 440px; max-width: 90%;" src="https://storage.googleapis.com/keras-cv/guides/keras-cv-augmentations.gif">
 
@@ -34,29 +34,44 @@ these common tasks.
 - [API Design Guidelines](.github/API_DESIGN.md)
 
 ## Installation
+KerasCV supports both Keras 2 and Keras 3. We recommend Keras 3 for all new 
+users, as it enables using KerasCV models and layers with JAX, TensorFlow and 
+PyTorch.
 
-To install the latest official release:
+### Keras 2 Installation
+
+To install the latest KerasCV release with Keras 2, simply run:
 
 ```
-pip install keras-cv tensorflow --upgrade
+pip install --upgrade keras-cv
 ```
 
-To install the latest unreleased changes to the library, we recommend using
-pip to install directly from the master branch on github:
+### Keras 3 Installation
+
+There are currently two ways to install Keras 3 with KerasNLP. To install the
+latest changes for KerasCV and Keras, you can use our nightly package.
+
 
 ```
-pip install git+https://github.com/keras-team/keras-cv.git tensorflow --upgrade
+pip install --upgrade keras-cv-nightly
 ```
 
-## Configuring your backend
+To install the stable versions of KerasCV and Keras 3, you should install Keras
+3 **after** installing KerasNLP. This is a temporary step while TensorFlow is
+pinned to Keras 2, and will no longer be necessary after TensorFlow 2.16.
 
-**Keras 3** is an upcoming release of the Keras library which supports
-TensorFlow, Jax or Torch as backends. This is supported today in KerasNLP,
-but will not be enabled by default until the official release of Keras 3. If you
-`pip install keras-cv` and run a script or notebook without changes, you will
-be using TensorFlow and **Keras 2**.
+```
+pip install --upgrade keras-nlp
+pip install keras>=3
+```
+> [!IMPORTANT]
+> Keras 3 will not function with TensorFlow 2.14 or earlier.
+
+## Configuring your backend
 
-If you would like to enable a preview of the Keras 3 behavior, you can do
+If you have Keras 3 installed in your environment (see installation above),
+you can use KerasCV with any of JAX, TensorFlow and PyTorch. To do so, set the
+`KERAS_BACKEND` environment variable. For example:
 so by setting the `KERAS_BACKEND` environment variable. For example:
 
 ```shell
@@ -75,14 +90,6 @@ import keras_cv
 > [!IMPORTANT]
 > Make sure to set the `KERAS_BACKEND` before import any Keras libraries, it
 > will be used to set up Keras when it is first imported.
-Until the Keras 3 release, KerasCV will use a preview of Keras 3 on PyPI named
-[keras-core](https://pypi.org/project/keras-core/).
-
-> [!IMPORTANT]
-> If you set `KERAS_BACKEND` variable, you should `import keras_core as keras`
-> instead of `import keras`. This is a temporary step until Keras 3 is out!
-To restore the default **Keras 2** behavior, `unset KERAS_BACKEND` before
-importing Keras and KerasCV.
 
 Once that configuration step is done, you can just import KerasCV and start
 using it on top of your backend of choice:

keras_cv/backend/__init__.py

@@ -11,6 +11,20 @@
 # 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.
+"""
+Keras backend module.
+
+This module adds a temporary Keras API surface that is fully under KerasCV
+control. The goal is to allow us to write Keras 3-like code everywhere, while
+still supporting Keras 2. We do this by using the `keras_core` package with
+Keras 2 to backport Keras 3 numerics APIs (`keras.ops` and `keras.random`) into
+Keras 2. The sub-modules exposed are as follows:
+
+- `config`: check which version of Keras is being run.
+- `keras`: The full `keras` API with compat shims for older Keras versions.
+- `ops`: `keras.ops` for Keras 3 or `keras_core.ops` for Keras 2.
+- `random`: `keras.random` for Keras 3 or `keras_core.ops` for Keras 2.
+"""
 from keras_cv.backend import config  # noqa: E402
 from keras_cv.backend import keras  # noqa: E402
 from keras_cv.backend import ops  # noqa: E402
@@ -19,12 +33,12 @@
 
 
 def assert_tf_keras(src):
-    if config.multi_backend():
+    if config.keras_3():
         raise NotImplementedError(
-            f"KerasCV component {src} does not yet support Keras Core, and can "
+            f"KerasCV component {src} does not yet support Keras 3, and can "
             "only be used in `tf.keras`."
         )
 
 
 def supports_ragged():
-    return not config.multi_backend()
+    return not config.keras_3()

keras_cv/backend/config.py

@@ -11,91 +11,37 @@
 # 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.
-import json
-import os
-
-_MULTI_BACKEND = False
-
-# Set Keras base dir path given KERAS_HOME env variable, if applicable.
-# Otherwise either ~/.keras or /tmp.
-if "KERAS_HOME" in os.environ:
-    _keras_dir = os.environ.get("KERAS_HOME")
-else:
-    _keras_base_dir = os.path.expanduser("~")
-    if not os.access(_keras_base_dir, os.W_OK):
-        _keras_base_dir = "/tmp"
-    _keras_dir = os.path.join(_keras_base_dir, ".keras")
-
 
-def detect_if_tensorflow_uses_keras_3():
-    # We follow the version of keras that tensorflow is configured to use.
-    from tensorflow import keras
+import os
 
-    # Note that only recent versions of keras have a `version()` function.
-    if hasattr(keras, "version") and keras.version().startswith("3."):
-        return True
+from tensorflow import keras
 
-    # No `keras.version()` means we are on an old version of keras.
-    return False
+# We follow the version of keras that tensorflow is configured to use.
+_USE_KERAS_3 = False
 
+# Note that only recent versions of keras have a `version()` function.
+if hasattr(keras, "version") and keras.version().startswith("3."):
+    _USE_KERAS_3 = True
 
-_USE_KERAS_3 = detect_if_tensorflow_uses_keras_3()
-if _USE_KERAS_3:
-    _MULTI_BACKEND = True
+if not _USE_KERAS_3:
+    backend = os.environ.get("KERAS_BACKEND")
+    if backend and backend != "tensorflow":
+        raise RuntimeError(
+            "When running Keras 2, the `KERAS_BACKEND` environment variable "
+            f"must either be unset or `'tensorflow'`. Received: `{backend}`. "
+            "To set another backend, please install Keras 3. See "
+            "https://github.com/keras-team/keras-cv#installation"
+        )
 
 
 def keras_3():
     """Check if Keras 3 is being used."""
     return _USE_KERAS_3
 
 
-# Attempt to read KerasCV config file.
-_config_path = os.path.expanduser(os.path.join(_keras_dir, "keras_cv.json"))
-if os.path.exists(_config_path):
-    try:
-        with open(_config_path) as f:
-            _config = json.load(f)
-    except ValueError:
-        _config = {}
-    _MULTI_BACKEND = _config.get("multi_backend", _MULTI_BACKEND)
-
-# Save config file, if possible.
-if not os.path.exists(_keras_dir):
-    try:
-        os.makedirs(_keras_dir)
-    except OSError:
-        # Except permission denied and potential race conditions
-        # in multi-threaded environments.
-        pass
-
-if not os.path.exists(_config_path):
-    _config = {
-        "multi_backend": _MULTI_BACKEND,
-    }
-    try:
-        with open(_config_path, "w") as f:
-            f.write(json.dumps(_config, indent=4))
-    except IOError:
-        # Except permission denied.
-        pass
-
-if "KERAS_BACKEND" in os.environ and os.environ["KERAS_BACKEND"]:
-    _MULTI_BACKEND = True
-
-
-def multi_backend():
-    return _MULTI_BACKEND
-
-
 def backend():
     """Check the backend framework."""
-    if not multi_backend():
-        return "tensorflow"
     if not keras_3():
-        import keras_core
-
-        return keras_core.config.backend()
-
-    from tensorflow import keras
+        return "tensorflow"
 
     return keras.config.backend()

keras_cv/backend/keras.py

@@ -30,11 +30,6 @@
     import keras  # noqa: F403, F401
     from keras import *  # noqa: F403, F401
 
-    keras.backend.name_scope = keras.name_scope
-elif config.multi_backend():
-    import keras_core as keras  # noqa: F403, F401
-    from keras_core import *  # noqa: F403, F401
-
     keras.backend.name_scope = keras.name_scope
 else:
     from tensorflow import keras  # noqa: F403, F401

keras_cv/backend/ops.py

@@ -11,10 +11,9 @@
 # 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.backend.config import keras_3
-from keras_cv.backend.config import multi_backend
+from keras_cv.backend import config
 
-if keras_3():
+if config.keras_3():
     from keras.ops import *  # noqa: F403, F401
     from keras.preprocessing.image import smart_resize  # noqa: F403, F401
 
@@ -32,5 +31,5 @@
         from keras_core.src.utils.image_utils import (  # noqa: F403, F401
             smart_resize,
         )
-if not multi_backend():
+if config.backend() == "tensorflow":
     from keras_cv.backend.tf_ops import *  # noqa: F403, F401

keras_cv/backend/scope.py

@@ -18,7 +18,7 @@
 from keras_cv.backend import keras
 from keras_cv.backend import ops
 from keras_cv.backend import tf_ops
-from keras_cv.backend.config import multi_backend
+from keras_cv.backend.config import keras_3
 
 _ORIGINAL_OPS = copy.copy(backend.ops.__dict__)
 _ORIGINAL_SUPPORTS_RAGGED = backend.supports_ragged
@@ -30,7 +30,7 @@
 def tf_data(function):
     @functools.wraps(function)
     def wrapper(*args, **kwargs):
-        if multi_backend() and keras.src.utils.backend_utils.in_tf_graph():
+        if keras_3() and keras.src.utils.backend_utils.in_tf_graph():
             with TFDataScope():
                 return function(*args, **kwargs)
         else:

keras_cv/backend/tf_ops.py

@@ -44,6 +44,3 @@
 from tensorflow import reduce_all as all  # noqa: F403, F401
 from tensorflow import reduce_max as max  # noqa: F403, F401
 from tensorflow import split  # noqa: F403, F401
-from tensorflow.keras.preprocessing.image import (  # noqa: F403, F401
-    smart_resize,
-)

keras_cv/callbacks/pycoco_callback.py

@@ -46,7 +46,6 @@ def __init__(
                 operations and passing `cache=False`.
         """
         assert_pycocotools_installed("PyCOCOCallback")
-        self.model = None
         self.val_data = validation_data
         if cache:
             # We cache the dataset to preserve a consistent iteration order.

keras_cv/callbacks/waymo_evaluation_callback.py

@@ -46,7 +46,6 @@ def __init__(self, validation_data, config=None, **kwargs):
         assert_waymo_open_dataset_installed(
             "keras_cv.callbacks.WaymoEvaluationCallback()"
         )
-        self.model = None
         self.val_data = validation_data
         self.evaluator = WODDetectionEvaluator(
             config=config or self._get_default_config()

keras_cv/conftest.py

@@ -17,7 +17,7 @@
 import tensorflow as tf
 from packaging import version
 
-from keras_cv.backend.config import multi_backend
+from keras_cv.backend.config import keras_3
 
 
 def pytest_addoption(parser):
@@ -69,12 +69,12 @@ def pytest_collection_modifyitems(config, items):
     skip_extra_large = pytest.mark.skipif(
         not run_extra_large_tests, reason="need --run_extra_large option to run"
     )
-    skip_tf_keras_only = pytest.mark.skipif(
-        multi_backend(),
+    skip_keras_2_only = pytest.mark.skipif(
+        keras_3(),
         reason="This test is only supported on tf.keras",
     )
     skip_tf_only = pytest.mark.skipif(
-        multi_backend() and keras_core.backend.backend() != "tensorflow",
+        keras_3() and keras_core.backend.backend() != "tensorflow",
         reason="This test is only supported on TensorFlow",
     )
     for item in items:
@@ -87,6 +87,6 @@ def pytest_collection_modifyitems(config, items):
         if "extra_large" in item.keywords:
             item.add_marker(skip_extra_large)
         if "tf_keras_only" in item.keywords:
-            item.add_marker(skip_tf_keras_only)
+            item.add_marker(skip_keras_2_only)
         if "tf_only" in item.keywords:
             item.add_marker(skip_tf_only)

keras_cv/layers/object_detection/multi_class_non_max_suppression.py

@@ -18,7 +18,7 @@
 from keras_cv.api_export import keras_cv_export
 from keras_cv.backend import keras
 from keras_cv.backend import ops
-from keras_cv.backend.config import multi_backend
+from keras_cv.backend.config import keras_3
 
 
 @keras_cv_export("keras_cv.layers.MultiClassNonMaxSuppression")
@@ -73,7 +73,7 @@ def call(
                 `bounding_box_format` specified in the constructor.
             class_prediction: Dense Tensor of shape [batch, boxes, num_classes].
         """
-        if multi_backend() and keras.backend.backend() != "tensorflow":
+        if keras_3() and keras.backend.backend() != "tensorflow":
             raise NotImplementedError(
                 "MultiClassNonMaxSuppression does not support non-TensorFlow "
                 "backends. Consider using NonMaxSuppression instead."

keras_cv/layers/object_detection/non_max_suppression.py

@@ -20,7 +20,7 @@
 from keras_cv.api_export import keras_cv_export
 from keras_cv.backend import keras
 from keras_cv.backend import ops
-from keras_cv.backend.config import multi_backend
+from keras_cv.backend.config import keras_3
 
 EPSILON = 1e-8
 
@@ -89,7 +89,7 @@ def call(
 
         confidence_prediction = ops.max(class_prediction, axis=-1)
 
-        if not multi_backend() or keras.backend.backend() == "tensorflow":
+        if not keras_3() or keras.backend.backend() == "tensorflow":
             idx, valid_det = tf.image.non_max_suppression_padded(
                 box_prediction,
                 confidence_prediction,
@@ -318,7 +318,7 @@ def suppression_loop_body(boxes, iou_threshold, output_size, idx):
 
     # TODO(ianstenbit): Fix bug in tfnp.take_along_axis that causes this hack.
     # (This will be removed anyway when we use built-in NMS for TF.)
-    if multi_backend() and keras.backend.backend() != "tensorflow":
+    if keras_3() and keras.backend.backend() != "tensorflow":
         idx = ops.take_along_axis(
             ops.reshape(sorted_indices, [-1]), take_along_axis_idx
         )