refactor prototype SBD to avoid closing demux streams at construction time

Author: pmeier

test/builtin_dataset_mocks.py

@@ -661,15 +661,15 @@ class SBDMockData:
     _NUM_CATEGORIES = 20
 
     @classmethod
-    def _make_split_files(cls, root_map):
-        ids_map = {
-            split: [f"2008_{idx:06d}" for idx in idcs]
-            for split, idcs in (
-                ("train", [0, 1, 2]),
-                ("train_noval", [0, 2]),
-                ("val", [3]),
-            )
-        }
+    def _make_split_files(cls, root_map, *, split):
+        splits_and_idcs = [
+            ("train", [0, 1, 2]),
+            ("val", [3]),
+        ]
+        if split == "train_noval":
+            splits_and_idcs.append(("train_noval", [0, 2]))
+
+        ids_map = {split: [f"2008_{idx:06d}" for idx in idcs] for split, idcs in splits_and_idcs}
 
         for split, ids in ids_map.items():
             with open(root_map[split] / f"{split}.txt", "w") as fh:
@@ -710,25 +710,27 @@ def _make_segmentation(cls, size):
         return torch.randint(0, cls._NUM_CATEGORIES + 1, size=size, dtype=torch.uint8).numpy()
 
     @classmethod
-    def generate(cls, root):
+    def generate(cls, root, *, split):
         archive_folder = root / "benchmark_RELEASE"
         dataset_folder = archive_folder / "dataset"
         dataset_folder.mkdir(parents=True, exist_ok=True)
 
-        ids, num_samples_map = cls._make_split_files(defaultdict(lambda: dataset_folder, {"train_noval": root}))
+        ids, num_samples_map = cls._make_split_files(
+            defaultdict(lambda: dataset_folder, {"train_noval": root}), split=split
+        )
         sizes = cls._make_anns_folder(dataset_folder, "cls", ids)
         create_image_folder(
             dataset_folder, "img", lambda idx: f"{ids[idx]}.jpg", num_examples=len(ids), size=lambda idx: sizes[idx]
         )
 
         make_tar(root, "benchmark.tgz", archive_folder, compression="gz")
 
-        return num_samples_map
+        return num_samples_map[split]
 
 
 @register_mock(configs=combinations_grid(split=("train", "val", "train_noval")))
 def sbd(root, config):
-    return SBDMockData.generate(root)[config["split"]]
+    return SBDMockData.generate(root, split=config["split"])
 
 
 @register_mock(configs=[dict()])

torchvision/prototype/datasets/_builtin/sbd.py

@@ -49,31 +49,35 @@ def __init__(
         super().__init__(root, dependencies=("scipy",), skip_integrity_check=skip_integrity_check)
 
     def _resources(self) -> List[OnlineResource]:
-        archive = HttpResource(
-            "https://www2.eecs.berkeley.edu/Research/Projects/CS/vision/grouping/semantic_contours/benchmark.tgz",
-            sha256="6a5a2918d5c73ce032fdeba876574d150d9d04113ab87540a1304cbcc715be53",
-        )
-        extra_split = HttpResource(
-            "http://home.bharathh.info/pubs/codes/SBD/train_noval.txt",
-            sha256="0b2068f7a359d2907431803e1cd63bf6162da37d7d503b589d3b08c6fd0c2432",
-        )
-        return [archive, extra_split]
+        resources = [
+            HttpResource(
+                "https://www2.eecs.berkeley.edu/Research/Projects/CS/vision/grouping/semantic_contours/benchmark.tgz",
+                sha256="6a5a2918d5c73ce032fdeba876574d150d9d04113ab87540a1304cbcc715be53",
+            )
+        ]
+        if self._split == "train_noval":
+            resources.append(
+                HttpResource(
+                    "http://home.bharathh.info/pubs/codes/SBD/train_noval.txt",
+                    sha256="0b2068f7a359d2907431803e1cd63bf6162da37d7d503b589d3b08c6fd0c2432",
+                )
+            )
+        return resources
 
     def _classify_archive(self, data: Tuple[str, Any]) -> Optional[int]:
         path = pathlib.Path(data[0])
         parent, grandparent, *_ = path.parents
 
-        if parent.name == "dataset":
-            return 0
-        elif grandparent.name == "dataset":
+        if grandparent.name == "dataset":
             if parent.name == "img":
-                return 1
+                return 0
             elif parent.name == "cls":
-                return 2
-            else:
-                return None
-        else:
-            return None
+                return 1
+
+        if parent.name == "dataset" and self._split != "train_noval":
+            return 2
+
+        return None
 
     def _prepare_sample(self, data: Tuple[Tuple[Any, Tuple[str, BinaryIO]], Tuple[str, BinaryIO]]) -> Dict[str, Any]:
         split_and_image_data, ann_data = data
@@ -93,21 +97,24 @@ def _prepare_sample(self, data: Tuple[Tuple[Any, Tuple[str, BinaryIO]], Tuple[st
         )
 
     def _datapipe(self, resource_dps: List[IterDataPipe]) -> IterDataPipe[Dict[str, Any]]:
-        archive_dp, extra_split_dp = resource_dps
-
-        archive_dp = resource_dps[0]
-        split_dp, images_dp, anns_dp = Demultiplexer(
-            archive_dp,
-            3,
-            self._classify_archive,
-            buffer_size=INFINITE_BUFFER_SIZE,
-            drop_none=True,
-        )
-        split_dp, to_be_closed_dp = (
-            (extra_split_dp, split_dp) if self._split == "train_noval" else (split_dp, extra_split_dp)
-        )
-        for _, file in to_be_closed_dp:
-            file.close()
+        if self._split == "train_noval":
+            archive_dp, split_dp = resource_dps
+            images_dp, anns_dp = Demultiplexer(
+                archive_dp,
+                2,
+                self._classify_archive,
+                buffer_size=INFINITE_BUFFER_SIZE,
+                drop_none=True,
+            )
+        else:
+            archive_dp = resource_dps[0]
+            images_dp, anns_dp, split_dp = Demultiplexer(
+                archive_dp,
+                3,
+                self._classify_archive,
+                buffer_size=INFINITE_BUFFER_SIZE,
+                drop_none=True,
+            )
 
         split_dp = Filter(split_dp, path_comparator("name", f"{self._split}.txt"))
         split_dp = LineReader(split_dp, decode=True)