Update model garden to not use tf.python.keras which will be removed soon.

PiperOrigin-RevId: 406002788

official/vision/beta/ops/augment.py

@@ -28,11 +28,10 @@ https://github.com/rwightman/pytorch-image-models
 import math
 from typing import Any, List, Iterable, Optional, Text, Tuple
 
+from keras.layers.preprocessing import image_preprocessing as image_ops
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras.layers.preprocessing import image_preprocessing as image_ops
-
 
 # This signifies the max integer that the controller RNN could predict for the
 # augmentation scheme.

official/vision/image_classification/augment.py

@@ -24,11 +24,11 @@ from __future__ import division
 from __future__ import print_function
 
 import math
+from typing import Any, Dict, List, Optional, Text, Tuple
 
+from keras.layers.preprocessing import image_preprocessing as image_ops
 import tensorflow as tf
-from typing import Any, Dict, List, Optional, Text, Tuple
 
-from tensorflow.python.keras.layers.preprocessing import image_preprocessing as image_ops
 
 # This signifies the max integer that the controller RNN could predict for the
 # augmentation scheme.