Creates train/val/test split for Kitti using video ids.

PiperOrigin-RevId: 258103690

Author: TensorFlow Datasets Team

tensorflow_datasets/image/kitti.py

@@ -47,6 +47,7 @@
 _DATA_URL = "https://s3.eu-central-1.amazonaws.com/avg-kitti"
 _IMAGES_FNAME = "data_object_image_2.zip"
 _LABELS_FNAME = "data_object_label_2.zip"
+_DEVKIT_FNAME = "devkit_object.zip"
 _OBJECT_LABELS = [
     "Car",
     "Van",
@@ -57,6 +58,10 @@
     "Tram",
     "Misc",
 ]
+# The percentage of trainset videos to put into validation and test sets.
+# The released test images do not have labels.
+_VALIDATION_SPLIT_PERCENT_VIDEOS = 10
+_TEST_SPLIT_PERCENT_VIDEOS = 10
 
 # Raw Kitti representation of a bounding box. Coordinates are in pixels,
 # measured from the top-left hand corner.
@@ -67,12 +72,13 @@
 class Kitti(tfds.core.GeneratorBasedBuilder):
   """Kitti dataset."""
 
-  VERSION = tfds.core.Version("1.0.0")
+  VERSION = tfds.core.Version("3.0.0")
   SUPPORTED_VERSIONS = [
       tfds.core.Version("2.0.0"),
   ]
   # Version history:
   # 2.0.0: S3 with new hashing function (different shuffle).
+  # 3.0.0: Train/val/test splits based on random video IDs created.
 
   def _info(self):
     # Annotation descriptions are in the object development kit.
@@ -102,26 +108,48 @@ def _split_generators(self, dl_manager):
     filenames = {
         "images": os.path.join(_DATA_URL, _IMAGES_FNAME),
         "annotations": os.path.join(_DATA_URL, _LABELS_FNAME),
+        "devkit": os.path.join(_DATA_URL, _DEVKIT_FNAME),
     }
     files = dl_manager.download(filenames)
+    train_images, validation_images, test_images = _build_splits(
+        dl_manager.iter_archive(files["devkit"]))
+
     return [
         tfds.core.SplitGenerator(
             name=tfds.Split.TRAIN,
             gen_kwargs={
                 "images": dl_manager.iter_archive(files["images"]),
                 "annotations": dl_manager.iter_archive(files["annotations"]),
                 "subdir": "training",
+                "image_ids": train_images,
+            }),
+        tfds.core.SplitGenerator(
+            name=tfds.Split.VALIDATION,
+            gen_kwargs={
+                "images": dl_manager.iter_archive(files["images"]),
+                "annotations": dl_manager.iter_archive(files["annotations"]),
+                "subdir": "training",
+                "image_ids": validation_images,
+            }),
+        tfds.core.SplitGenerator(
+            name=tfds.Split.TEST,
+            gen_kwargs={
+                "images": dl_manager.iter_archive(files["images"]),
+                "annotations": dl_manager.iter_archive(files["annotations"]),
+                "subdir": "training",
+                "image_ids": test_images,
             }),
     ]
 
-  def _generate_examples(self, images, annotations, subdir):
+  def _generate_examples(self, images, annotations, subdir, image_ids):
     """Yields images and annotations.
 
     Args:
       images: object that iterates over the archive of images.
       annotations: object that iterates over the archive of annotations.
       subdir: subdirectory from which to extract images and annotations, e.g.
         training or testing.
+      image_ids: file ids for images in this split.
 
     Yields:
       A tuple containing the example's key, and the example.
@@ -145,8 +173,10 @@ def _generate_examples(self, images, annotations, subdir):
         continue
       if prefix.split("/")[0] != subdir:
         continue
-
-      annotations = all_annotations[int(prefix[-6:])]
+      image_id = int(prefix[-6:])
+      if image_id not in image_ids:
+        continue
+      annotations = all_annotations[image_id]
       img = cv2.imdecode(np.fromstring(fobj.read(), dtype=np.uint8),
                          cv2.IMREAD_COLOR)
       img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
@@ -211,3 +241,71 @@ def _parse_kitti_annotations(annotations_csv):
         "rotation_y": float(rotation_y),
     })
   return annotations
+
+
+def _build_splits(devkit):
+  """Splits the train data into train/val/test by video.
+
+  Ensures that images from the same video do not traverse the splits.
+
+  Args:
+    devkit: object that iterates over the devkit archive.
+
+  Returns:
+    train_images: File ids for the training set images.
+    validation_images: File ids for the validation set images.
+    test_images: File ids for the test set images.
+  """
+  mapping_line_ids = None
+  mapping_lines = None
+  for fpath, fobj in devkit:
+    if fpath == "mapping/train_rand.txt":
+      # Converts 1-based line index to 0-based line index.
+      mapping_line_ids = [
+          int(x.strip()) - 1 for x in fobj.read().decode("utf-8").split(",")
+      ]
+    if fpath == "mapping/train_mapping.txt":
+      mapping_lines = fobj.readlines()
+      mapping_lines = [x.decode("utf-8") for x in mapping_lines]
+
+  assert mapping_line_ids
+  assert mapping_lines
+
+  video_to_image = collections.defaultdict(list)
+  for image_id, mapping_lineid in enumerate(mapping_line_ids):
+    line = mapping_lines[mapping_lineid]
+    video_id = line.split(" ")[1]
+    video_to_image[video_id].append(image_id)
+
+  # Sets numpy random state.
+  numpy_original_state = np.random.get_state()
+  np.random.seed(seed=123)
+
+  # Max 1 for testing.
+  num_test_videos = max(1,
+                        _TEST_SPLIT_PERCENT_VIDEOS * len(video_to_image) // 100)
+  num_validation_videos = max(
+      1,
+      _VALIDATION_SPLIT_PERCENT_VIDEOS * len(video_to_image) // 100)
+  test_videos = set(
+      np.random.choice(
+          list(video_to_image.keys()), num_test_videos, replace=False))
+  validation_videos = set(
+      np.random.choice(
+          list(set(video_to_image.keys()) - set(test_videos)),
+          num_validation_videos,
+          replace=False))
+  test_images = []
+  validation_images = []
+  train_images = []
+  for k, v in video_to_image.items():
+    if k in test_videos:
+      test_images.extend(v)
+    elif k in validation_videos:
+      validation_images.extend(v)
+    else:
+      train_images.extend(v)
+
+  # Resets numpy random state.
+  np.random.set_state(numpy_original_state)
+  return train_images, validation_images, test_images

tensorflow_datasets/image/kitti_test.py

@@ -26,11 +26,14 @@
 class KittiTest(testing.DatasetBuilderTestCase):
   DATASET_CLASS = kitti.Kitti
   SPLITS = {
-      "train": 5,
+      "train": 6,
+      "validation": 2,
+      "test": 2,
   }
   DL_EXTRACT_RESULT = {
       "images": "data_object_image_2.zip",
       "annotations": "data_object_label_2.zip",
+      "devkit": "devkit_object.zip",
   }
 
 

tensorflow_datasets/testing/kitti.py

@@ -40,7 +40,8 @@
                     "Path to tensorflow_datasets directory")
 
 FLAGS = flags.FLAGS
-NUM_IMAGES = 5
+NUM_IMAGES = 10
+NUM_VIDEOS = 5
 HEIGHT = 375
 WIDTH = 1242
 OBJECTS = [
@@ -139,6 +140,28 @@ def _get_label_file(annotation):
   return fobj.name
 
 
+def _get_mapping_files():
+  """Returns dummy image to video mapping files."""
+  # Random indices file.
+  train_rand = np.random.permutation(range(1, NUM_IMAGES + 1))  # 1-based index
+  fobj_rand = tempfile.NamedTemporaryFile(
+      delete=False, mode="wb", suffix=".txt")
+  fobj_rand.write(",".join([str(x) for x in train_rand]))
+  fobj_rand.close()
+
+  # Mapping file.
+  fobj_map = tempfile.NamedTemporaryFile(delete=False, mode="wb", suffix=".txt")
+  assert NUM_IMAGES > NUM_VIDEOS
+  assert NUM_IMAGES % NUM_VIDEOS == 0
+  vid_ids = list(range(NUM_VIDEOS)) * (NUM_IMAGES // NUM_VIDEOS)
+  for vid in vid_ids:
+    row = "2011_09_26 2011_09_26_drive_00{:02d}_sync 0000000123".format(vid)
+    fobj_map.write(row + "\n")
+  fobj_map.close()
+
+  return fobj_rand.name, fobj_map.name
+
+
 def _create_zip_files():
   """Saves png and label using name index."""
   if not os.path.exists(_output_dir()):
@@ -161,6 +184,13 @@ def _create_zip_files():
           label,
           os.path.join("training", "label_2", "label_{:06d}.txt".format(i)))
 
+  devkit_out_path = os.path.join(_output_dir(), "devkit_object.zip")
+  with zipfile.ZipFile(devkit_out_path, "w") as devkit_zip:
+    train_rand, train_mapping = _get_mapping_files()
+    devkit_zip.write(train_rand, os.path.join("mapping", "train_rand.txt"))
+    devkit_zip.write(train_mapping, os.path.join("mapping",
+                                                 "train_mapping.txt"))
+
 
 def main(argv):
   if len(argv) > 1:

tensorflow_datasets/testing/test_data/fake_examples/kitti/data_object_image_2.zip

@@ -1,2 +1,3 @@
 https://s3.eu-central-1.amazonaws.com/avg-kitti/data_object_image_2.zip 12569945557 351c5a2aa0cd9238b50174a3a62b846bc5855da256b82a196431d60ff8d43617
 https://s3.eu-central-1.amazonaws.com/avg-kitti/data_object_label_2.zip 5601213 4efc76220d867e1c31bb980bbf8cbc02599f02a9cb4350effa98dbb04aaed880
+https://s3.eu-central-1.amazonaws.com/avg-kitti/devkit_object.zip 63778 cfde67e531832618ea0fd4844f91e34a45068025e2bef79e278cb812dc2537d0