use enum for category generation

Author: pmeier

torchvision/prototype/datasets/_builtin/dtd.py

@@ -129,7 +129,7 @@ def _make_datapipe(
         return Mapper(dp, self._collate_and_decode_sample, fn_kwargs=dict(decoder=decoder))
 
     def _filter_images(self, data: Tuple[str, Any]) -> bool:
-        return self._classify_archive(data) == 2
+        return self._classify_archive(data) == DTDDemux.IMAGES
 
     def _generate_categories(self, root: pathlib.Path) -> List[str]:
         dp = self.resources(self.default_config)[0].load(pathlib.Path(root) / self.name)

torchvision/prototype/datasets/_builtin/oxford_iiit_pet.py

@@ -141,7 +141,7 @@ def _make_datapipe(
         return Mapper(dp, functools.partial(self._collate_and_decode_sample, decoder=decoder))
 
     def _filter_split_and_classification_anns(self, data: Tuple[str, Any]) -> bool:
-        return self._classify_anns(data) == 0
+        return self._classify_anns(data) == OxfordIITPetDemux.SPLIT_AND_CLASSIFICATION
 
     def _generate_categories(self, root: pathlib.Path) -> List[str]:
         config = self.default_config