class_docstr: Tests for youtube_vis dataset.
code: def _download_and_prepare_as_dataset(self, builder):
      super()._download_and_prepare_as_dataset(builder)
  
      if (
          not tf.executing_eagerly()
      ):  # Only test the following in eager mode.
          return
  
      with self.subTest("check_annotations"):
          # Makes certain sanity checks like the number of segmentations
          # should be equal to the number of labeled frame indices, etc.
          splits = builder.as_dataset()
          train_ex = list(splits[tfds.Split.TRAIN])[0]
          val_ex = list(splits[tfds.Split.VALIDATION])[0]
          test_ex = list(splits[tfds.Split.TEST])[0]
          self.assertEqual(
              train_ex["tracks"]["bboxes"].shape[0],
              train_ex["tracks"]["segmentations"].shape[0],
          )
          self.assertEqual(
              train_ex["tracks"]["segmentations"].shape[0],
              train_ex["tracks"]["frames"].shape[0],
          )
          if builder.builder_config.name != "test_config_custom_split":
              # No annotations are provided on the val and test data if not using
              # the custom split.
              self.assertEqual(val_ex["tracks"]["bboxes"].shape[0], 0)
              self.assertEqual(test_ex["tracks"]["bboxes"].shape[0], 0)
  
      with self.subTest("check_video"):
          # Checks that the video data matches the metadata.
          splits = builder.as_dataset()
          train_ex = list(splits[tfds.Split.TRAIN])[0]
          val_ex = list(splits[tfds.Split.VALIDATION])[0]
          test_ex = list(splits[tfds.Split.TEST])[0]
  
          self.assertEqual(
              train_ex["metadata"]["height"].numpy(),
              builder.builder_config.height,
          )
          self.assertEqual(
              train_ex["metadata"]["width"].numpy(),
              builder.builder_config.width,
          )
          self.assertEqual(
              val_ex["metadata"]["height"].numpy(),
              builder.builder_config.height,
          )
          self.assertEqual(
              val_ex["metadata"]["width"].numpy(),
              builder.builder_config.width,
          )
          self.assertEqual(
              test_ex["metadata"]["height"].numpy(),
              builder.builder_config.height,
          )
          self.assertEqual(
              test_ex["metadata"]["width"].numpy(),
              builder.builder_config.width,
          )
          self.assertEqual(
              train_ex["video"].shape,
              (
                  train_ex["metadata"]["num_frames"].numpy(),
                  train_ex["metadata"]["height"].numpy(),
                  train_ex["metadata"]["width"].numpy(),
                  3,
              ),
          )
          self.assertEqual(
              val_ex["video"].shape,
              (
                  val_ex["metadata"]["num_frames"].numpy(),
                  val_ex["metadata"]["height"].numpy(),
                  val_ex["metadata"]["width"].numpy(),
                  3,
              ),
          )
          self.assertEqual(
              test_ex["video"].shape,
              (
                  test_ex["metadata"]["num_frames"].numpy(),
                  test_ex["metadata"]["height"].numpy(),
                  test_ex["metadata"]["width"].numpy(),
                  3,
              ),
          )
  
      with self.subTest("check_segmentations"):
          # Checks that the segmentations match the metadata and have the correct
          # label count.
          splits = builder.as_dataset()
          train_ex = list(splits[tfds.Split.TRAIN])[0]
          for segmentation in train_ex["tracks"]["segmentations"]:
              unique_pixels = set(np.unique(segmentation.numpy()))
              self.assertSetEqual(unique_pixels, set([0, 1]))
              self.assertEqual(
                  (
                      train_ex["metadata"]["height"].numpy(),
                      train_ex["metadata"]["width"].numpy(),
                  ),
                  segmentation.shape[1:3],
              )
  
      if builder.builder_config.name == "test_config_custom_split":
          with self.subTest("check_custom_split"):
              # Checks that all splits contain exactly the same data, since
              # that is what was requested by this custom split.
              splits = builder.as_dataset()
              train_ex = list(splits[tfds.Split.TRAIN])[0]
              val_ex = list(splits[tfds.Split.VALIDATION])[0]
              val_ex = list(splits[tfds.Split.TEST])[0]
              self.assertDictEqual(
                  train_ex["metadata"], val_ex["metadata"]
              )
              self.assertDictEqual(
                  test_ex["metadata"], val_ex["metadata"]
              )
  
file_path: datasets/tensorflow_datasets/video/youtube_vis/youtube_vis_test.py
imports: import numpy
  import tensorflow
  import tensorflow_datasets.public_api
  from tensorflow_datasets.video.youtube_vis import youtube_vis
label_desc: Use image node_id 2 for calling the YoutubeVisTest obj's underlying member method code with example usage: obj._download_and_prepare_as_dataset(builder) without return types