Closing streams to avoid testing issues

test/test_prototype_builtin_datasets.py

@@ -15,6 +15,17 @@
 from torchvision.prototype import transforms, datasets
 from torchvision.prototype.datasets.utils._internal import INFINITE_BUFFER_SIZE
 from torchvision.prototype.features import Image, Label
+from torchdata.datapipes.utils import StreamWrapper
+
+
+from torchdata.datapipes.iter import (
+    IterDataPipe,
+    Mapper,
+    Demultiplexer,
+    Filter,
+    IterKeyZipper,
+    LineReader,
+)
 
 assert_samples_equal = functools.partial(
     assert_equal, pair_types=(TensorLikePair, ObjectPair), rtol=0, atol=0, equal_nan=True
@@ -64,9 +75,9 @@ def test_smoke(self, dataset_mock, config):
     @parametrize_dataset_mocks(DATASET_MOCKS)
     def test_sample(self, dataset_mock, config):
         dataset, _ = dataset_mock.load(config)
-
         try:
-            sample = next(iter(dataset))
+            iterator = iter(dataset)
+            sample = next(iterator)
         except StopIteration:
             raise AssertionError("Unable to draw any sample.") from None
         except Exception as error:
@@ -78,23 +89,34 @@ def test_sample(self, dataset_mock, config):
         if not sample:
             raise AssertionError("Sample dictionary is empty.")
 
+        list(iterator) # Cleanups and closing streams in buffers
+
     @parametrize_dataset_mocks(DATASET_MOCKS)
     def test_num_samples(self, dataset_mock, config):
         dataset, mock_info = dataset_mock.load(config)
-
         assert len(list(dataset)) == mock_info["num_samples"]
 
     @parametrize_dataset_mocks(DATASET_MOCKS)
     def test_no_vanilla_tensors(self, dataset_mock, config):
+        StreamWrapper.session_streams = {}
         dataset, _ = dataset_mock.load(config)
 
-        vanilla_tensors = {key for key, value in next(iter(dataset)).items() if type(value) is torch.Tensor}
+        iterator = iter(dataset)
+        one_element = next(iterator)
+
+        vanilla_tensors = {key for key, value in one_element.items() if type(value) is torch.Tensor}
         if vanilla_tensors:
             raise AssertionError(
                 f"The values of key(s) "
                 f"{sequence_to_str(sorted(vanilla_tensors), separate_last='and ')} contained vanilla tensors."
             )
 
+        list(iterator) # Cleanups and closing streams in buffers
+
+        if (len(StreamWrapper.session_streams)>0):
+            Demultiplexer.buffers()
+            raise Exception(StreamWrapper.session_streams)
+
     @parametrize_dataset_mocks(DATASET_MOCKS)
     def test_transformable(self, dataset_mock, config):
         dataset, _ = dataset_mock.load(config)

torchvision/prototype/datasets/_builtin/caltech.py

@@ -107,7 +107,9 @@ def _prepare_sample(
         ann_path, ann_buffer = ann_data
 
         image = EncodedImage.from_file(image_buffer)
+        image_buffer.close()
         ann = read_mat(ann_buffer)
+        ann_buffer.close()
 
         return dict(
             label=Label.from_category(category, categories=self._categories),
@@ -186,10 +188,11 @@ def _is_not_rogue_file(self, data: Tuple[str, Any]) -> bool:
 
     def _prepare_sample(self, data: Tuple[str, BinaryIO]) -> Dict[str, Any]:
         path, buffer = data
-
+        image = EncodedImage.from_file(buffer)
+        buffer.close()
         return dict(
             path=path,
-            image=EncodedImage.from_file(buffer),
+            image=image,
             label=Label(int(pathlib.Path(path).parent.name.split(".", 1)[0]) - 1, categories=self._categories),
         )
 

torchvision/prototype/datasets/_builtin/celeba.py

@@ -39,8 +39,8 @@ def __init__(
         self.fieldnames = fieldnames
 
     def __iter__(self) -> Iterator[Tuple[str, Dict[str, str]]]:
-        for _, file in self.datapipe:
-            file = (line.decode() for line in file)
+        for _, fh in self.datapipe:
+            file = (line.decode() for line in fh)
 
             if self.fieldnames:
                 fieldnames = self.fieldnames
@@ -58,6 +58,8 @@ def __iter__(self) -> Iterator[Tuple[str, Dict[str, str]]]:
             for line in csv.DictReader(file, fieldnames=fieldnames, dialect="celeba"):
                 yield line.pop("image_id"), line
 
+            fh.close()
+
 
 NAME = "celeba"
 
@@ -142,6 +144,7 @@ def _prepare_sample(
         path, buffer = image_data
 
         image = EncodedImage.from_file(buffer)
+        buffer.close()
         (_, identity), (_, attributes), (_, bounding_box), (_, landmarks) = ann_data
 
         return dict(

torchvision/prototype/datasets/_builtin/cifar.py

@@ -66,7 +66,9 @@ def _resources(self) -> List[OnlineResource]:
 
     def _unpickle(self, data: Tuple[str, io.BytesIO]) -> Dict[str, Any]:
         _, file = data
-        return cast(Dict[str, Any], pickle.load(file, encoding="latin1"))
+        result = pickle.load(file, encoding="latin1")
+        file.close()
+        return cast(Dict[str, Any], result)
 
     def _prepare_sample(self, data: Tuple[np.ndarray, int]) -> Dict[str, Any]:
         image_array, category_idx = data

torchvision/prototype/datasets/_builtin/clevr.py

@@ -15,6 +15,8 @@
 
 from .._api import register_dataset, register_info
 
+from torchdata import janitor
+
 NAME = "clevr"
 
 
@@ -62,10 +64,12 @@ def _add_empty_anns(self, data: Tuple[str, BinaryIO]) -> Tuple[Tuple[str, Binary
     def _prepare_sample(self, data: Tuple[Tuple[str, BinaryIO], Optional[Dict[str, Any]]]) -> Dict[str, Any]:
         image_data, scenes_data = data
         path, buffer = image_data
+        image = EncodedImage.from_file(buffer)
+        buffer.close()
 
         return dict(
             path=path,
-            image=EncodedImage.from_file(buffer),
+            image=image,
             label=Label(len(scenes_data["objects"])) if scenes_data else None,
         )
 
@@ -97,6 +101,8 @@ def _datapipe(self, resource_dps: List[IterDataPipe]) -> IterDataPipe[Dict[str,
                 buffer_size=INFINITE_BUFFER_SIZE,
             )
         else:
+            for i in scenes_dp:
+                janitor(i)
             dp = Mapper(images_dp, self._add_empty_anns)
 
         return Mapper(dp, self._prepare_sample)

torchvision/prototype/datasets/_builtin/coco.py

@@ -25,6 +25,7 @@
     INFINITE_BUFFER_SIZE,
     getitem,
     read_categories_file,
+    close_buffer,
     path_accessor,
     hint_sharding,
     hint_shuffling,
@@ -174,9 +175,10 @@ def _classify_meta(self, data: Tuple[str, Any]) -> Optional[int]:
 
     def _prepare_image(self, data: Tuple[str, BinaryIO]) -> Dict[str, Any]:
         path, buffer = data
+        image = close_buffer(EncodedImage.from_file, buffer)
         return dict(
             path=path,
-            image=EncodedImage.from_file(buffer),
+            image=image,
         )
 
     def _prepare_sample(
@@ -187,9 +189,11 @@ def _prepare_sample(
         anns, image_meta = ann_data
 
         sample = self._prepare_image(image_data)
+
         # this method is only called if we have annotations
         annotations = cast(str, self._annotations)
         sample.update(self._ANN_DECODERS[annotations](self, anns, image_meta))
+
         return sample
 
     def _datapipe(self, resource_dps: List[IterDataPipe]) -> IterDataPipe[Dict[str, Any]]:

torchvision/prototype/datasets/_builtin/country211.py

@@ -51,11 +51,13 @@ def _resources(self) -> List[OnlineResource]:
 
     def _prepare_sample(self, data: Tuple[str, Any]) -> Dict[str, Any]:
         path, buffer = data
+        image = EncodedImage.from_file(buffer)
+        buffer.close()
         category = pathlib.Path(path).parent.name
         return dict(
             label=Label.from_category(category, categories=self._categories),
             path=path,
-            image=EncodedImage.from_file(buffer),
+            image=image,
         )
 
     def _filter_split(self, data: Tuple[str, Any], *, split: str) -> bool:

torchvision/prototype/datasets/_builtin/cub200.py

@@ -130,12 +130,14 @@ def _2011_prepare_ann(
     ) -> Dict[str, Any]:
         _, (bounding_box_data, segmentation_data) = data
         segmentation_path, segmentation_buffer = segmentation_data
+        segmentation = EncodedImage.from_file(segmentation_buffer)
+        segmentation_buffer.close()
         return dict(
             bounding_box=BoundingBox(
                 [float(part) for part in bounding_box_data[1:]], format="xywh", image_size=image_size
             ),
             segmentation_path=segmentation_path,
-            segmentation=EncodedImage.from_file(segmentation_buffer),
+            segmentation=segmentation,
         )
 
     def _2010_split_key(self, data: str) -> str:
@@ -148,6 +150,7 @@ def _2010_anns_key(self, data: Tuple[str, BinaryIO]) -> Tuple[str, Tuple[str, Bi
     def _2010_prepare_ann(self, data: Tuple[str, Tuple[str, BinaryIO]], image_size: Tuple[int, int]) -> Dict[str, Any]:
         _, (path, buffer) = data
         content = read_mat(buffer)
+        buffer.close()
         return dict(
             ann_path=path,
             bounding_box=BoundingBox(
@@ -169,6 +172,7 @@ def _prepare_sample(
         path, buffer = image_data
 
         image = EncodedImage.from_file(buffer)
+        buffer.close()
 
         return dict(
             prepare_ann_fn(anns_data, image.image_size),

torchvision/prototype/datasets/_builtin/dtd.py

@@ -88,14 +88,16 @@ def _prepare_sample(self, data: Tuple[Tuple[str, List[str]], Tuple[str, BinaryIO
         (_, joint_categories_data), image_data = data
         _, *joint_categories = joint_categories_data
         path, buffer = image_data
+        image = EncodedImage.from_file(buffer)
+        buffer.close()
 
         category = pathlib.Path(path).parent.name
 
         return dict(
             joint_categories={category for category in joint_categories if category},
             label=Label.from_category(category, categories=self._categories),
             path=path,
-            image=EncodedImage.from_file(buffer),
+            image=image,
         )
 
     def _datapipe(self, resource_dps: List[IterDataPipe]) -> IterDataPipe[Dict[str, Any]]:

torchvision/prototype/datasets/_builtin/eurosat.py

@@ -49,10 +49,12 @@ def _resources(self) -> List[OnlineResource]:
     def _prepare_sample(self, data: Tuple[str, Any]) -> Dict[str, Any]:
         path, buffer = data
         category = pathlib.Path(path).parent.name
+        image = EncodedImage.from_file(buffer)
+        buffer.close()
         return dict(
             label=Label.from_category(category, categories=self._categories),
             path=path,
-            image=EncodedImage.from_file(buffer),
+            image=image,
         )
 
     def _datapipe(self, resource_dps: List[IterDataPipe]) -> IterDataPipe[Dict[str, Any]]:

torchvision/prototype/datasets/_builtin/food101.py

@@ -63,10 +63,12 @@ def _classify_archive(self, data: Tuple[str, Any]) -> Optional[int]:
 
     def _prepare_sample(self, data: Tuple[str, Tuple[str, BinaryIO]]) -> Dict[str, Any]:
         id, (path, buffer) = data
+        image = EncodedImage.from_file(buffer)
+        buffer.close()
         return dict(
             label=Label.from_category(id.split("/", 1)[0], categories=self._categories),
             path=path,
-            image=EncodedImage.from_file(buffer),
+            image=image,
         )
 
     def _image_key(self, data: Tuple[str, Any]) -> str:

torchvision/prototype/datasets/_builtin/gtsrb.py

@@ -84,10 +84,12 @@ def _prepare_sample(self, data: Tuple[Tuple[str, Any], Dict[str, Any]]) -> Dict[
             format="xyxy",
             image_size=(int(csv_info["Height"]), int(csv_info["Width"])),
         )
+        image = EncodedImage.from_file(buffer)
+        buffer.close()
 
         return {
             "path": path,
-            "image": EncodedImage.from_file(buffer),
+            "image": image,
             "label": Label(label, categories=self._categories),
             "bounding_box": bounding_box,
         }

torchvision/prototype/datasets/_builtin/imagenet.py

@@ -109,10 +109,11 @@ def _prepare_test_data(self, data: Tuple[str, BinaryIO]) -> Tuple[None, Tuple[st
         return None, data
 
     def _classifiy_devkit(self, data: Tuple[str, BinaryIO]) -> Optional[int]:
+        name, binary_io = data
         return {
             "meta.mat": ImageNetDemux.META,
             "ILSVRC2012_validation_ground_truth.txt": ImageNetDemux.LABEL,
-        }.get(pathlib.Path(data[0]).name)
+        }.get(pathlib.Path(name).name)
 
     # Although the WordNet IDs (wnids) are unique, the corresponding categories are not. For example, both n02012849
     # and n03126707 are labeled 'crane' while the first means the bird and the latter means the construction equipment
@@ -123,12 +124,14 @@ def _classifiy_devkit(self, data: Tuple[str, BinaryIO]) -> Optional[int]:
 
     def _extract_categories_and_wnids(self, data: Tuple[str, BinaryIO]) -> List[Tuple[str, str]]:
         synsets = read_mat(data[1], squeeze_me=True)["synsets"]
-        return [
+        results = [
             (self._WNID_MAP.get(wnid, category.split(",", 1)[0]), wnid)
             for _, wnid, category, _, num_children, *_ in synsets
             # if num_children > 0, we are looking at a superclass that has no direct instance
             if num_children == 0
         ]
+        data[1].close()
+        return results
 
     def _imagenet_label_to_wnid(self, imagenet_label: str, *, wnids: Tuple[str, ...]) -> str:
         return wnids[int(imagenet_label) - 1]
@@ -151,11 +154,13 @@ def _prepare_sample(
         data: Tuple[Optional[Tuple[Label, str]], Tuple[str, BinaryIO]],
     ) -> Dict[str, Any]:
         label_data, (path, buffer) = data
+        image = EncodedImage.from_file(buffer)
+        buffer.close()
 
         return dict(
             dict(zip(("label", "wnid"), label_data if label_data else (None, None))),
             path=path,
-            image=EncodedImage.from_file(buffer),
+            image=image,
         )
 
     def _datapipe(self, resource_dps: List[IterDataPipe]) -> IterDataPipe[Dict[str, Any]]:

torchvision/prototype/datasets/_builtin/mnist.py

@@ -57,6 +57,8 @@ def __iter__(self) -> Iterator[torch.Tensor]:
             for _ in range(stop - start):
                 yield read(dtype=dtype, count=count).reshape(shape)
 
+            file.close()
+
 
 class _MNISTBase(Dataset):
     _URL_BASE: Union[str, Sequence[str]]

torchvision/prototype/datasets/_builtin/oxford_iiit_pet.py

@@ -80,14 +80,18 @@ def _prepare_sample(
         classification_data, segmentation_data = ann_data
         segmentation_path, segmentation_buffer = segmentation_data
         image_path, image_buffer = image_data
+        segmentation = EncodedImage.from_file(segmentation_buffer)
+        segmentation_buffer.close()
+        image = EncodedImage.from_file(image_buffer)
+        image_buffer.close()
 
         return dict(
             label=Label(int(classification_data["label"]) - 1, categories=self._categories),
             species="cat" if classification_data["species"] == "1" else "dog",
             segmentation_path=segmentation_path,
-            segmentation=EncodedImage.from_file(segmentation_buffer),
+            segmentation=segmentation,
             image_path=image_path,
-            image=EncodedImage.from_file(image_buffer),
+            image=image,
         )
 
     def _datapipe(self, resource_dps: List[IterDataPipe]) -> IterDataPipe[Dict[str, Any]]:

torchvision/prototype/datasets/_builtin/pcam.py

@@ -39,6 +39,7 @@ def __iter__(self) -> Iterator[Tuple[str, io.IOBase]]:
                 if self.key is not None:
                     data = data[self.key]
                 yield from data
+            handle.close()
 
 
 _Resource = namedtuple("_Resource", ("file_name", "gdrive_id", "sha256"))