Added segmentation mask to Oxford-IIIT pets dataset.

Extended oxford_iiit_pet.py and other related files

PiperOrigin-RevId: 257082403

Author: TensorFlow Datasets Team

tensorflow_datasets/image/oxford_iiit_pet.py

@@ -48,10 +48,13 @@ class OxfordIIITPet(tfds.core.GeneratorBasedBuilder):
 
   VERSION = tfds.core.Version("1.1.0")
   SUPPORTED_VERSIONS = [
+      tfds.core.Version("2.1.0", experiments={tfds.core.Experiment.S3: True}),
       tfds.core.Version("2.0.0", experiments={tfds.core.Experiment.S3: True}),
+      tfds.core.Version("1.2.0"),
       tfds.core.Version("1.1.0"),
   ]
   # Version history:
+  # 2.1.0, 1.2.0: addition of the segmentation_mask feature.
   # 2.0.0: S3 (new shuffling, sharding and slicing mechanism).
 
   def _info(self):
@@ -62,6 +65,7 @@ def _info(self):
             "image": tfds.features.Image(),
             "label": tfds.features.ClassLabel(num_classes=37),
             "file_name": tfds.features.Text(),
+            "segmentation_mask": tfds.features.Image(shape=(None, None, 1))
         }),
         supervised_keys=("image", "label"),
         urls=["http://www.robots.ox.ac.uk/~vgg/data/pets/"],
@@ -91,6 +95,7 @@ def _split_generators(self, dl_manager):
         num_shards=_NUM_SHARDS,
         gen_kwargs={
             "images_dir_path": images_path_dir,
+            "annotations_dir_path": annotations_path_dir,
             "images_list_file": os.path.join(annotations_path_dir,
                                              "trainval.txt"),
             },
@@ -100,25 +105,31 @@ def _split_generators(self, dl_manager):
         num_shards=_NUM_SHARDS,
         gen_kwargs={
             "images_dir_path": images_path_dir,
+            "annotations_dir_path": annotations_path_dir,
             "images_list_file": os.path.join(annotations_path_dir,
                                              "test.txt")
             },
         )
 
     return [train_split, test_split]
 
-  def _generate_examples(self, images_dir_path, images_list_file):
+  def _generate_examples(self, images_dir_path, annotations_dir_path,
+                         images_list_file):
     with tf.io.gfile.GFile(images_list_file, "r") as images_list:
       for line in images_list:
         image_name, label, _, _ = line.strip().split(" ")
 
+        trimaps_dir_path = os.path.join(annotations_dir_path, "trimaps")
+
+        trimap_name = image_name + ".png"
         image_name += ".jpg"
         label = int(label) - 1
 
         record = {
             "image": os.path.join(images_dir_path, image_name),
             "label": int(label),
-            "file_name": image_name
+            "file_name": image_name,
+            "segmentation_mask": os.path.join(trimaps_dir_path, trimap_name)
         }
         if self.version.implements(tfds.core.Experiment.S3):
           yield image_name, record

tensorflow_datasets/testing/fake_data_utils.py

@@ -32,17 +32,17 @@
 CHANNELS_NB = 3
 
 
-def get_random_picture(height=None, width=None):
+def get_random_picture(height=None, width=None, channels=CHANNELS_NB):
   """Returns random picture as np.ndarray (int)."""
   height = height or random.randrange(MIN_HEIGHT_WIDTH, MAX_HEIGHT_WIDTH)
   width = width or random.randrange(MIN_HEIGHT_WIDTH, MAX_HEIGHT_WIDTH)
   return np.random.randint(
-      256, size=(height, width, CHANNELS_NB), dtype=np.uint8)
+      256, size=(height, width, channels), dtype=np.uint8)
 
 
-def get_random_jpeg(height=None, width=None):
+def get_random_jpeg(height=None, width=None, channels=CHANNELS_NB):
   """Returns path to JPEG picture."""
-  image = get_random_picture(height, width)
+  image = get_random_picture(height, width, channels)
   jpeg = tf.image.encode_jpeg(image)
   with utils.nogpu_session() as sess:
     res = sess.run(jpeg)
@@ -52,9 +52,9 @@ def get_random_jpeg(height=None, width=None):
   return fobj.name
 
 
-def get_random_png(height=None, width=None):
+def get_random_png(height=None, width=None, channels=CHANNELS_NB):
   """Returns path to PNG picture."""
-  image = get_random_picture(height, width)
+  image = get_random_picture(height, width, channels)
   png = tf.image.encode_png(image)
   with utils.nogpu_session() as sess:
     res = sess.run(png)

tensorflow_datasets/testing/oxford_iiit_pet.py

@@ -60,8 +60,16 @@ def _generate_data():
 
   # Generate annotations
   annotations_dir = os.path.join(_output_dir(), 'annotations')
+
   if not tf.io.gfile.exists(annotations_dir):
     tf.io.gfile.makedirs(annotations_dir)
+
+  # Generate trimaps
+  trimaps_dir = os.path.join(annotations_dir, 'trimaps')
+
+  if not tf.io.gfile.exists(trimaps_dir):
+    tf.io.gfile.makedirs(trimaps_dir)
+
   global_count = 0
   for filename, num_examples in [('trainval.txt', _TRAIN_IMAGES_NUMBER),
                                  ('test.txt', _TEST_IMAGES_NUMBER)]:
@@ -73,6 +81,13 @@ def _generate_data():
     tf.io.gfile.copy(fobj.name, os.path.join(annotations_dir, filename),
                      overwrite=True)
 
+  # Create trimaps
+  for i in range(_TRAIN_IMAGES_NUMBER + _TEST_IMAGES_NUMBER):
+    trimap_name = 'image{:03d}.png'.format(i)
+    tf.io.gfile.copy(fake_data_utils.get_random_png(channels=1),
+                     os.path.join(trimaps_dir, trimap_name),
+                     overwrite=True)
+
 
 def main(argv):
   if len(argv) > 1: