Merge pull request #1209 from hannesrichter:master

PiperOrigin-RevId: 282059203

Author: copybara-github

tensorflow_datasets/image/deep_weeds.py

@@ -19,13 +19,14 @@
 from __future__ import division
 from __future__ import print_function
 
+import csv
 import os
 
 import tensorflow as tf
 import tensorflow_datasets.public_api as tfds
 
-
 _URL = "https://nextcloud.qriscloud.org.au/index.php/s/a3KxPawpqkiorST/download"
+_URL_LABELS = "https://raw.githubusercontent.com/AlexOlsen/DeepWeeds/master/labels/labels.csv"
 
 _DESCRIPTION = (
     """The DeepWeeds dataset consists of 17,509 images capturing eight different weed species native to Australia """
@@ -34,11 +35,6 @@
     """ "Cluden", "Douglas", "Hervey Range", "Kelso", "McKinlay" and "Paluma"."""
 )
 
-_NAMES = [
-    "Chinee apple", "Snake weed", "Lantana", "Prickly acacia", "Siam weed",
-    "Parthenium", "Rubber vine", "Parkinsonia", "Negative"
-]
-
 _IMAGE_SHAPE = (256, 256, 3)
 
 _CITATION = """\
@@ -73,7 +69,7 @@
 class DeepWeeds(tfds.core.GeneratorBasedBuilder):
   """DeepWeeds Image Dataset Class."""
 
-  VERSION = tfds.core.Version("1.0.0")
+  VERSION = tfds.core.Version("2.0.0")
 
   def _info(self):
     """Define Dataset Info."""
@@ -83,7 +79,7 @@ def _info(self):
         description=(_DESCRIPTION),
         features=tfds.features.FeaturesDict({
             "image": tfds.features.Image(shape=_IMAGE_SHAPE),
-            "label": tfds.features.ClassLabel(names=_NAMES),
+            "label": tfds.features.ClassLabel(num_classes=9),
         }),
         supervised_keys=("image", "label"),
         homepage="https://github.com/AlexOlsen/DeepWeeds",
@@ -93,24 +89,39 @@ def _info(self):
   def _split_generators(self, dl_manager):
     """Define Splits."""
     # The file is in ZIP format, but URL doesn't mention it.
-    path = dl_manager.download_and_extract(
-        tfds.download.Resource(
+    paths = dl_manager.download_and_extract({
+        "image": tfds.download.Resource(
             url=_URL,
-            extract_method=tfds.download.ExtractMethod.ZIP))
+            extract_method=tfds.download.ExtractMethod.ZIP),
+        "label": _URL_LABELS})
 
     return [
         tfds.core.SplitGenerator(
             name="train",
             gen_kwargs={
-                "data_dir_path": path,
+                "data_dir_path": paths["image"],
+                "label_path": paths["label"],
             },
         ),
     ]
 
-  def _generate_examples(self, data_dir_path):
+  def _generate_examples(self, data_dir_path, label_path):
     """Generate images and labels for splits."""
 
+    with tf.io.gfile.GFile(label_path) as f:
+      # Convert to list to reuse the iterator multiple times
+      reader = list(csv.DictReader(f))
+
+    # Extract the mapping int -> str and save the label name string to the
+    # feature
+    label_id_to_name = {int(row["Label"]): row["Species"] for row in reader}
+    self.info.features["label"].names = [
+        v for _, v in sorted(label_id_to_name.items())
+    ]
+
+    filename_to_label = {row["Filename"]: row["Species"] for row in reader}
     for file_name in tf.io.gfile.listdir(data_dir_path):
-      image = os.path.join(data_dir_path, file_name)
-      label = _NAMES[int(file_name.split("-")[2].split(".")[0])]
-      yield file_name, {"image": image, "label": label}
+      yield file_name, {
+          "image": os.path.join(data_dir_path, file_name),
+          "label": filename_to_label[file_name]
+      }

tensorflow_datasets/image/deep_weeds_test.py

@@ -23,7 +23,11 @@
 class DeepWeedsTest(tfds_test.DatasetBuilderTestCase):
   DATASET_CLASS = deep_weeds.DeepWeeds
   SPLITS = {
-      "train": 4,
+      "train": 9,
+  }
+  DL_EXTRACT_RESULT = {
+      "image": "images",
+      "label": "labels.csv",
   }