use enums in prototype datasets for demux

Author: pmeier

torchvision/prototype/datasets/_builtin/clevr.py

@@ -1,3 +1,4 @@
+import enum
 import functools
 import io
 import pathlib
@@ -24,6 +25,11 @@
 from torchvision.prototype.features import Label
 
 
+class CLEVRDemux(enum.IntEnum):
+    IMAGES = 0
+    SCENES = 1
+
+
 class CLEVR(Dataset):
     def _make_info(self) -> DatasetInfo:
         return DatasetInfo(
@@ -43,9 +49,9 @@ def resources(self, config: DatasetConfig) -> List[OnlineResource]:
     def _classify_archive(self, data: Tuple[str, Any]) -> Optional[int]:
         path = pathlib.Path(data[0])
         if path.parents[1].name == "images":
-            return 0
+            return CLEVRDemux.IMAGES
         elif path.parent.name == "scenes":
-            return 1
+            return CLEVRDemux.SCENES
         else:
             return None
 

torchvision/prototype/datasets/_builtin/coco.py

@@ -1,3 +1,4 @@
+import enum
 import functools
 import io
 import pathlib
@@ -37,6 +38,11 @@
 from torchvision.prototype.utils._internal import FrozenMapping
 
 
+class CocoDemux(enum.IntEnum):
+    IMAGES_META = 0
+    ANNS_META = 1
+
+
 class Coco(Dataset):
     def _make_info(self) -> DatasetInfo:
         name = "coco"
@@ -144,9 +150,9 @@ def _filter_meta_files(self, data: Tuple[str, Any], *, split: str, year: str, an
     def _classify_meta(self, data: Tuple[str, Any]) -> Optional[int]:
         key, _ = data
         if key == "images":
-            return 0
+            return CocoDemux.IMAGES_META
         elif key == "annotations":
-            return 1
+            return CocoDemux.ANNS_META
         else:
             return None
 

torchvision/prototype/datasets/_builtin/dtd.py

@@ -1,3 +1,4 @@
+import enum
 import io
 import pathlib
 from typing import Any, Callable, Dict, List, Optional, Tuple
@@ -30,6 +31,12 @@
 from torchvision.prototype.features import Label
 
 
+class DTDDemux(enum.IntEnum):
+    SPLIT = 0
+    JOINT_CATEGORIES = 1
+    IMAGES = 2
+
+
 class DTD(Dataset):
     def _make_info(self) -> DatasetInfo:
         return DatasetInfo(
@@ -54,11 +61,11 @@ def _classify_archive(self, data: Tuple[str, Any]) -> Optional[int]:
         path = pathlib.Path(data[0])
         if path.parent.name == "labels":
             if path.name == "labels_joint_anno.txt":
-                return 1
+                return DTDDemux.JOINT_CATEGORIES
 
-            return 0
+            return DTDDemux.SPLIT
         elif path.parents[1].name == "images":
-            return 2
+            return DTDDemux.IMAGES
         else:
             return None
 

torchvision/prototype/datasets/_builtin/mnist.py

@@ -1,4 +1,5 @@
 import abc
+import enum
 import functools
 import io
 import operator
@@ -232,6 +233,11 @@ def _make_info(self) -> DatasetInfo:
     }
 
 
+class EMNISTDemux(enum.IntEnum):
+    IMAGES = 0
+    LABELS = 1
+
+
 class EMNIST(_MNISTBase):
     def _make_info(self) -> DatasetInfo:
         return DatasetInfo(
@@ -273,9 +279,9 @@ def _classify_archive(self, data: Tuple[str, Any], *, config: DatasetConfig) ->
         path = pathlib.Path(data[0])
         (images_file, _), (labels_file, _) = self._files_and_checksums(config)
         if path.name == images_file:
-            return 0
+            return EMNISTDemux.IMAGES
         elif path.name == labels_file:
-            return 1
+            return EMNISTDemux.LABELS
         else:
             return None
 
@@ -320,6 +326,7 @@ def _make_datapipe(
         decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
     ) -> IterDataPipe[Dict[str, Any]]:
         archive_dp = resource_dps[0]
+
         images_dp, labels_dp = Demultiplexer(
             archive_dp,
             2,

torchvision/prototype/datasets/_builtin/oxford_iiit_pet.py

@@ -1,3 +1,4 @@
+import enum
 import functools
 import io
 import pathlib
@@ -24,6 +25,11 @@
 from torchvision.prototype.features import Label
 
 
+class OxfordIITPetDemux(enum.IntEnum):
+    SPLIT_AND_CLASSIFICATION = 0
+    SEGMENTATIONS = 1
+
+
 class OxfordIITPet(Dataset):
     def _make_info(self) -> DatasetInfo:
         return DatasetInfo(
@@ -51,8 +57,8 @@ def resources(self, config: DatasetConfig) -> List[OnlineResource]:
 
     def _classify_anns(self, data: Tuple[str, Any]) -> Optional[int]:
         return {
-            "annotations": 0,
-            "trimaps": 1,
+            "annotations": OxfordIITPetDemux.SPLIT_AND_CLASSIFICATION,
+            "trimaps": OxfordIITPetDemux.SEGMENTATIONS,
         }.get(pathlib.Path(data[0]).parent.name)
 
     def _filter_images(self, data: Tuple[str, Any]) -> bool:

torchvision/prototype/datasets/_builtin/sbd.py

@@ -1,3 +1,4 @@
+import enum
 import functools
 import io
 import pathlib
@@ -33,6 +34,12 @@
 )
 
 
+class SBDDemux(enum.IntEnum):
+    SPLIT = 0
+    IMAGES = 1
+    ANNS = 2
+
+
 class SBD(Dataset):
     def _make_info(self) -> DatasetInfo:
         return DatasetInfo(
@@ -63,12 +70,12 @@ def _classify_archive(self, data: Tuple[str, Any]) -> Optional[int]:
         parent, grandparent, *_ = path.parents
 
         if parent.name == "dataset":
-            return 0
+            return SBDDemux.SPLIT
         elif grandparent.name == "dataset":
             if parent.name == "img":
-                return 1
+                return SBDDemux.IMAGES
             elif parent.name == "cls":
-                return 2
+                return SBDDemux.ANNS
             else:
                 return None
         else:

torchvision/prototype/datasets/_builtin/voc.py

@@ -1,3 +1,4 @@
+import enum
 import functools
 import io
 import pathlib
@@ -31,6 +32,13 @@
     hint_shuffling,
 )
 
+
+class VOCDemux(enum.IntEnum):
+    SPLIT = 0
+    IMAGES = 1
+    ANNS = 2
+
+
 HERE = pathlib.Path(__file__).parent
 
 
@@ -75,11 +83,11 @@ def _is_in_folder(self, data: Tuple[str, Any], *, name: str, depth: int = 1) ->
 
     def _classify_archive(self, data: Tuple[str, Any], *, config: DatasetConfig) -> Optional[int]:
         if self._is_in_folder(data, name="ImageSets", depth=2):
-            return 0
+            return VOCDemux.SPLIT
         elif self._is_in_folder(data, name="JPEGImages"):
-            return 1
+            return VOCDemux.IMAGES
         elif self._is_in_folder(data, name=self._ANNS_FOLDER[config.task]):
-            return 2
+            return VOCDemux.ANNS
         else:
             return None