Fixed wrong imports

Author: vfdev-5

test/datasets_utils.py

@@ -584,8 +584,8 @@ def test_transforms(self, config):
 
     @test_all_configs
     def test_transforms_v2_wrapper(self, config):
-        from torchvision.datapoints import wrap_dataset_for_transforms_v2
         from torchvision.datapoints._datapoint import Datapoint
+        from torchvision.datasets import wrap_dataset_for_transforms_v2
 
         try:
             with self.create_dataset(config) as (dataset, _):

test/test_datasets.py

@@ -8,6 +8,7 @@
 import pathlib
 import pickle
 import random
+import re
 import shutil
 import string
 import unittest
@@ -3309,5 +3310,47 @@ def test_bad_input(self):
                 pass
 
 
+class TestDatasetWrapper:
+    def test_unknown_type(self):
+        unknown_object = object()
+        with pytest.raises(
+            TypeError, match=re.escape("is meant for subclasses of `torchvision.datasets.VisionDataset`")
+        ):
+            datasets.wrap_dataset_for_transforms_v2(unknown_object)
+
+    def test_unknown_dataset(self):
+        class MyVisionDataset(datasets.VisionDataset):
+            pass
+
+        dataset = MyVisionDataset("root")
+
+        with pytest.raises(TypeError, match="No wrapper exist"):
+            datasets.wrap_dataset_for_transforms_v2(dataset)
+
+    def test_missing_wrapper(self):
+        dataset = datasets.FakeData()
+
+        with pytest.raises(TypeError, match="please open an issue"):
+            datasets.wrap_dataset_for_transforms_v2(dataset)
+
+    def test_subclass(self, mocker):
+        from torchvision import datapoints
+
+        sentinel = object()
+        mocker.patch.dict(
+            datapoints._dataset_wrapper.WRAPPER_FACTORIES,
+            clear=False,
+            values={datasets.FakeData: lambda dataset: lambda idx, sample: sentinel},
+        )
+
+        class MyFakeData(datasets.FakeData):
+            pass
+
+        dataset = MyFakeData()
+        wrapped_dataset = datasets.wrap_dataset_for_transforms_v2(dataset)
+
+        assert wrapped_dataset[0] is sentinel
+
+
 if __name__ == "__main__":
     unittest.main()

test/test_prototype_datapoints.py

@@ -5,7 +5,6 @@
 
 from PIL import Image
 
-from torchvision import datapoints, datasets
 from torchvision.prototype import datapoints as proto_datapoints
 
 
@@ -163,43 +162,3 @@ def test_bbox_instance(data, format):
     if isinstance(format, str):
         format = datapoints.BoundingBoxFormat.from_str(format.upper())
     assert bboxes.format == format
-
-
-class TestDatasetWrapper:
-    def test_unknown_type(self):
-        unknown_object = object()
-        with pytest.raises(
-            TypeError, match=re.escape("is meant for subclasses of `torchvision.datasets.VisionDataset`")
-        ):
-            datapoints.wrap_dataset_for_transforms_v2(unknown_object)
-
-    def test_unknown_dataset(self):
-        class MyVisionDataset(datasets.VisionDataset):
-            pass
-
-        dataset = MyVisionDataset("root")
-
-        with pytest.raises(TypeError, match="No wrapper exist"):
-            datapoints.wrap_dataset_for_transforms_v2(dataset)
-
-    def test_missing_wrapper(self):
-        dataset = datasets.FakeData()
-
-        with pytest.raises(TypeError, match="please open an issue"):
-            datapoints.wrap_dataset_for_transforms_v2(dataset)
-
-    def test_subclass(self, mocker):
-        sentinel = object()
-        mocker.patch.dict(
-            datapoints._dataset_wrapper.WRAPPER_FACTORIES,
-            clear=False,
-            values={datasets.FakeData: lambda dataset: lambda idx, sample: sentinel},
-        )
-
-        class MyFakeData(datasets.FakeData):
-            pass
-
-        dataset = MyFakeData()
-        wrapped_dataset = datapoints.wrap_dataset_for_transforms_v2(dataset)
-
-        assert wrapped_dataset[0] is sentinel

torchvision/datapoints/__init__.py

@@ -1,11 +1,11 @@
+from torchvision import _BETA_TRANSFORMS_WARNING, _WARN_ABOUT_BETA_TRANSFORMS
+
 from ._bounding_box import BoundingBox, BoundingBoxFormat
 from ._datapoint import _FillType, _FillTypeJIT, _InputType, _InputTypeJIT
 from ._image import _ImageType, _ImageTypeJIT, _TensorImageType, _TensorImageTypeJIT, Image
 from ._mask import Mask
 from ._video import _TensorVideoType, _TensorVideoTypeJIT, _VideoType, _VideoTypeJIT, Video
 
-from torchvision import _BETA_TRANSFORMS_WARNING, _WARN_ABOUT_BETA_TRANSFORMS
-
 if _WARN_ABOUT_BETA_TRANSFORMS:
     import warnings