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

@@ -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

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):

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

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):

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

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):

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