From d9ac89a31459f455d828a2154349e0cda1dbd0d7 Mon Sep 17 00:00:00 2001
From: Ian Stenbit <3072903+ianstenbit@users.noreply.github.com>
Date: Thu, 7 Sep 2023 11:12:40 -0600
Subject: [PATCH] Remove direct imports of tf.keras (#2048)

* Remove direct imports of keras

* Undo change in predict_utils

* Undo augmentation layer changes
---
 keras_cv/callbacks/pycoco_callback.py           | 2 +-
 keras_cv/callbacks/waymo_evaluation_callback.py | 7 +++----
 keras_cv/layers/preprocessing/random_zoom.py    | 2 +-
 keras_cv/models/legacy/utils.py                 | 2 +-
 keras_cv/models/legacy/utils_test.py            | 2 +-
 keras_cv/models/legacy/weights.py               | 2 +-
 keras_cv/models/utils_test.py                   | 2 +-
 7 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/keras_cv/callbacks/pycoco_callback.py b/keras_cv/callbacks/pycoco_callback.py
index 889fa361e1..0de03fe799 100644
--- a/keras_cv/callbacks/pycoco_callback.py
+++ b/keras_cv/callbacks/pycoco_callback.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import numpy as np
-from keras.callbacks import Callback
+from tensorflow.keras.callbacks import Callback
 
 from keras_cv import bounding_box
 from keras_cv.api_export import keras_cv_export
diff --git a/keras_cv/callbacks/waymo_evaluation_callback.py b/keras_cv/callbacks/waymo_evaluation_callback.py
index 458b3796f3..cc73eb9ae8 100644
--- a/keras_cv/callbacks/waymo_evaluation_callback.py
+++ b/keras_cv/callbacks/waymo_evaluation_callback.py
@@ -12,8 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import tensorflow as tf
-from keras.callbacks import Callback
+from tensorflow.keras.callbacks import Callback
 
+from keras_cv.api_export import keras_cv_export
+from keras_cv.bounding_box_3d import CENTER_XYZ_DXDYDZ_PHI
 from keras_cv.utils import assert_waymo_open_dataset_installed
 
 try:
@@ -26,9 +28,6 @@
 except ImportError:
     WODDetectionEvaluator = None
 
-from keras_cv.api_export import keras_cv_export
-from keras_cv.bounding_box_3d import CENTER_XYZ_DXDYDZ_PHI
-
 
 @keras_cv_export("keras_cv.callbacks.WaymoEvaluationCallback")
 class WaymoEvaluationCallback(Callback):
diff --git a/keras_cv/layers/preprocessing/random_zoom.py b/keras_cv/layers/preprocessing/random_zoom.py
index 2e860da649..7af69a2bbf 100644
--- a/keras_cv/layers/preprocessing/random_zoom.py
+++ b/keras_cv/layers/preprocessing/random_zoom.py
@@ -14,7 +14,7 @@
 
 
 import tensorflow as tf
-from keras import backend
+from tensorflow.keras import backend
 
 from keras_cv.api_export import keras_cv_export
 from keras_cv.layers.preprocessing.vectorized_base_image_augmentation_layer import (  # noqa: E501
diff --git a/keras_cv/models/legacy/utils.py b/keras_cv/models/legacy/utils.py
index ba3e852b52..df11081a95 100644
--- a/keras_cv/models/legacy/utils.py
+++ b/keras_cv/models/legacy/utils.py
@@ -14,8 +14,8 @@
 # ==============================================================================
 """Utility functions for models"""
 
-from keras import layers
 from tensorflow import keras
+from tensorflow.keras import layers
 
 
 def parse_model_inputs(input_shape, input_tensor):
diff --git a/keras_cv/models/legacy/utils_test.py b/keras_cv/models/legacy/utils_test.py
index df9e15fc2e..497d573fce 100644
--- a/keras_cv/models/legacy/utils_test.py
+++ b/keras_cv/models/legacy/utils_test.py
@@ -13,8 +13,8 @@
 # limitations under the License.
 """Tests for KerasCV model utils."""
 
-from keras import layers
 from tensorflow import keras
+from tensorflow.keras import layers
 
 from keras_cv.models.legacy import utils
 from keras_cv.tests.test_case import TestCase
diff --git a/keras_cv/models/legacy/weights.py b/keras_cv/models/legacy/weights.py
index fc99606000..5f93d1ba97 100644
--- a/keras_cv/models/legacy/weights.py
+++ b/keras_cv/models/legacy/weights.py
@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 import tensorflow as tf
-from keras import utils
+from tensorflow.keras import utils
 
 
 def parse_weights(weights, include_top, model_type):
diff --git a/keras_cv/models/utils_test.py b/keras_cv/models/utils_test.py
index 712bd702b1..cc6c1ac596 100644
--- a/keras_cv/models/utils_test.py
+++ b/keras_cv/models/utils_test.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 """Tests for KerasCV model utils."""
 
-from keras import layers
+from tensorflow.keras import layers
 
 from keras_cv.models import utils
 from keras_cv.tests.test_case import TestCase