Avoid caching the entire Waymo Open Dataset during loading (#1559)

* Avoid caching the entire Waymo Open Dataset during loading * Formatting
keras-team · Mar 21, 2023 · 1974263 · 1974263
1 parent 4ebbdb9
commit 1974263
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/keras_cv/datasets/waymo/load.py b/keras_cv/datasets/waymo/load.py
@@ -15,7 +15,6 @@
 import os
 
 import tensorflow as tf
-import tensorflow_datasets as tfds
 
 from keras_cv.datasets.waymo import transformer
 from keras_cv.utils import assert_waymo_open_dataset_installed
@@ -28,9 +27,9 @@
 
 def _generate_frames(segments, transformer):
     def _generator():
-        for record in tfds.as_numpy(segments):
+        for record in segments:
             frame = waymo_open_dataset.dataset_pb2.Frame()
-            frame.ParseFromString(record)
+            frame.ParseFromString(record.numpy())
             yield transformer(frame)
 
     return _generator