Add ToPureTensor transform

docs/source/transforms.rst

@@ -238,6 +238,7 @@ Conversion
     v2.ConvertImageDtype
     v2.ToDtype
     v2.ConvertBoundingBoxFormat
+    v2.ToPureTensor
 
 Auto-Augmentation
 -----------------

references/classification/presets.py

@@ -68,6 +68,9 @@ def __init__(
         if random_erase_prob > 0:
             transforms.append(T.RandomErasing(p=random_erase_prob))
 
+        if use_v2:
+            transforms.append(T.ToPureTensor())
+
         self.transforms = T.Compose(transforms)
 
     def __call__(self, img):
@@ -107,6 +110,9 @@ def __init__(
             T.Normalize(mean=mean, std=std),
         ]
 
+        if use_v2:
+            transforms.append(T.ToPureTensor())
+
         self.transforms = T.Compose(transforms)
 
     def __call__(self, img):

references/detection/presets.py

@@ -79,6 +79,7 @@ def __init__(
             transforms += [
                 T.ConvertBoundingBoxFormat(datapoints.BoundingBoxFormat.XYXY),
                 T.SanitizeBoundingBoxes(),
+                T.ToPureTensor(),
             ]
 
         self.transforms = T.Compose(transforms)
@@ -103,6 +104,10 @@ def __init__(self, backend="pil", use_v2=False):
             raise ValueError(f"backend can be 'datapoint', 'tensor' or 'pil', but got {backend}")
 
         transforms += [T.ConvertImageDtype(torch.float)]
+
+        if use_v2:
+            transforms += [T.ToPureTensor()]
+
         self.transforms = T.Compose(transforms)
 
     def __call__(self, img, target):

references/segmentation/presets.py

@@ -63,6 +63,8 @@ def __init__(
             transforms += [T.ConvertImageDtype(torch.float)]
 
         transforms += [T.Normalize(mean=mean, std=std)]
+        if use_v2:
+            transforms += [T.ToPureTensor()]
 
         self.transforms = T.Compose(transforms)
 
@@ -98,6 +100,9 @@ def __init__(
             T.ConvertImageDtype(torch.float),
             T.Normalize(mean=mean, std=std),
         ]
+        if use_v2:
+            transforms += [T.ToPureTensor()]
+
         self.transforms = T.Compose(transforms)
 
     def __call__(self, img, target):

test/test_transforms_v2_refactored.py

@@ -2270,3 +2270,24 @@ def test_image_correctness(self, permutation, batch_dims):
         expected = self.reference_image_correctness(image, permutation=permutation)
 
         torch.testing.assert_close(actual, expected)
+
+
+class TestToPureTensor:
+    def test_correctness(self):
+        input = {
+            "img": make_image(),
+            "img_tensor": make_image_tensor(),
+            "img_pil": make_image_pil(),
+            "mask": make_detection_mask(),
+            "video": make_video(),
+            "bbox": make_bounding_box(),
+            "str": "str",
+        }
+
+        out = transforms.ToPureTensor()(input)
+
+        for input_value, out_value in zip(input.values(), out.values()):
+            if isinstance(input_value, datapoints.Datapoint):
+                assert isinstance(out_value, torch.Tensor) and not isinstance(out_value, datapoints.Datapoint)
+            else:
+                assert isinstance(out_value, type(input_value))

torchvision/transforms/v2/__init__.py

@@ -52,7 +52,7 @@
     ToDtype,
 )
 from ._temporal import UniformTemporalSubsample
-from ._type_conversion import PILToTensor, ToImagePIL, ToImageTensor, ToPILImage
+from ._type_conversion import PILToTensor, ToImagePIL, ToImageTensor, ToPILImage, ToPureTensor
 
 from ._deprecated import ToTensor  # usort: skip
 

torchvision/transforms/v2/_type_conversion.py

@@ -80,3 +80,17 @@ def _transform(
 # We changed the name to align them with the new naming scheme. Still, `ToPILImage` is
 # prevalent and well understood. Thus, we just alias it without deprecating the old name.
 ToPILImage = ToImagePIL
+
+
+class ToPureTensor(Transform):
+    """[BETA] Convert all datapoints to pure tensors, removing associated metadata (if any).
+
+    .. v2betastatus:: ToPureTensor transform
+
+    This doesn't scale or change the values, only the type.
+    """
+
+    _transformed_types = (datapoints.Image, datapoints.Video, datapoints.Mask, datapoints.BoundingBoxes)
+
+    def _transform(self, inpt: Any, params: Dict[str, Any]) -> torch.Tensor:
+        return inpt.as_subclass(torch.Tensor)