It... works?

Author: scotts

src/torchcodec/decoders/_video_decoder.py

@@ -19,7 +19,7 @@
     create_decoder,
     ERROR_REPORTING_INSTRUCTIONS,
 )
-from torchcodec.transforms import DecoderTransform, Resize
+from torchcodec.transforms import DecoderTransform, RandomCrop, Resize
 
 
 class VideoDecoder:
@@ -167,7 +167,9 @@ def __init__(
             device = str(device)
 
         device_variant = _get_cuda_backend()
-        transform_specs = _make_transform_specs(transforms)
+        transform_specs = _make_transform_specs(
+            transforms, input_dims=(self.metadata.height, self.metadata.width)
+        )
 
         core.add_video_stream(
             self._decoder,
@@ -450,6 +452,7 @@ def _get_and_validate_stream_metadata(
 
 def _convert_to_decoder_transforms(
     transforms: Sequence[Union[DecoderTransform, nn.Module]],
+    input_dims: Tuple[int, int],
 ) -> List[DecoderTransform]:
     """Convert a sequence of transforms that may contain TorchVision transform
     objects into a list of only TorchCodec transform objects.
@@ -482,21 +485,29 @@ def _convert_to_decoder_transforms(
                     "v2 transforms, but TorchVision is not installed."
                 )
             elif isinstance(transform, v2.Resize):
-                converted_transforms.append(Resize._from_torchvision(transform))
+                transform_tc = Resize._from_torchvision(transform)
+                input_dims = transform_tc._get_output_dims(input_dims)
+                converted_transforms.append(transform_tc)
+            elif isinstance(transform, v2.RandomCrop):
+                transform_tc = RandomCrop._from_torchvision(transform, input_dims)
+                input_dims = transform_tc._get_output_dims(input_dims)
+                converted_transforms.append(transform_tc)
             else:
                 raise ValueError(
                     f"Unsupported transform: {transform}. Transforms must be "
                     "either a TorchCodec DecoderTransform or a TorchVision "
                     "v2 transform."
                 )
         else:
+            intput_dims = transform._get_output_dims(input_dims)
             converted_transforms.append(transform)
 
     return converted_transforms
 
 
 def _make_transform_specs(
     transforms: Optional[Sequence[Union[DecoderTransform, nn.Module]]],
+    input_dims: Tuple[int, int],
 ) -> str:
     """Given a sequence of transforms, turn those into the specification string
        the core API expects.
@@ -516,7 +527,7 @@ def _make_transform_specs(
     if transforms is None:
         return ""
 
-    transforms = _convert_to_decoder_transforms(transforms)
+    transforms = _convert_to_decoder_transforms(transforms, input_dims)
     return ";".join([t._make_transform_spec() for t in transforms])
 
 

src/torchcodec/transforms/__init__.py

@@ -4,4 +4,4 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from ._decoder_transforms import DecoderTransform, Resize  # noqa
+from ._decoder_transforms import DecoderTransform, RandomCrop, Resize  # noqa

src/torchcodec/transforms/_decoder_transforms.py

@@ -7,8 +7,9 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from types import ModuleType
-from typing import Sequence
+from typing import Optional, Sequence, Tuple
 
+import torch
 from torch import nn
 
 
@@ -40,6 +41,9 @@ class DecoderTransform(ABC):
     def _make_transform_spec(self) -> str:
         pass
 
+    def _get_output_dims(self, input_dims: Tuple[int, int]) -> Tuple[int, int]:
+        return input_dims
+
 
 def import_torchvision_transforms_v2() -> ModuleType:
     try:
@@ -69,6 +73,9 @@ def _make_transform_spec(self) -> str:
         assert len(self.size) == 2
         return f"resize, {self.size[0]}, {self.size[1]}"
 
+    def _get_output_dims(self, input_dims: Tuple[int, int]) -> Tuple[int, int]:
+        return self.size
+
     @classmethod
     def _from_torchvision(cls, resize_tv: nn.Module):
         v2 = import_torchvision_transforms_v2()
@@ -92,19 +99,38 @@ def _from_torchvision(cls, resize_tv: nn.Module):
             )
         return cls(size=resize_tv.size)
 
+
 @dataclass
 class RandomCrop(DecoderTransform):
 
     size: Sequence[int]
     _top: Optional[int] = None
     _left: Optional[int] = None
+    _input_dims: Optional[Tuple[int, int]] = None
 
     def _make_transform_spec(self) -> str:
         assert len(self.size) == 2
-        return f"crop, {self.size[0]}, {self.size[1]}, {_left}, {_top}"
+        if self._top is None or self._left is None:
+            assert self._input_dims is not None
+            if self._input_dims[0] < self.size[0] or self._input_dims[1] < self.size[1]:
+                raise ValueError(
+                    f"Input dimensions {input_dims} are smaller than the crop size {self.size}."
+                )
+            self._top = torch.randint(
+                0, self._input_dims[0] - self.size[0] + 1, size=()
+            )
+            self._left = torch.randint(
+                0, self._input_dims[1] - self.size[1] + 1, size=()
+            )
+
+        return f"crop, {self.size[0]}, {self.size[1]}, {self._left}, {self._top}"
+
+    def _get_output_dims(self, input_dims: Tuple[int, int]) -> Tuple[int, int]:
+        self._input_dims = input_dims
+        return self.size
 
     @classmethod
-    def _from_torchvision(cls, random_crop_tv: nn.Module):
+    def _from_torchvision(cls, random_crop_tv: nn.Module, input_dims: Tuple[int, int]):
         v2 = import_torchvision_transforms_v2()
 
         assert isinstance(random_crop_tv, v2.RandomCrop)
@@ -128,5 +154,8 @@ def _from_torchvision(cls, random_crop_tv: nn.Module):
                 "TorchVision RandcomCrop transform must have a (height, width) "
                 f"pair for the size, got {random_crop_tv.size}."
             )
-        params = random_crop_tv.make_params([])
-        return cls(size=random_crop_tv.size)
+        params = random_crop_tv.make_params(
+            torch.empty(size=(3, *input_dims), dtype=torch.uint8)
+        )
+        assert random_crop_tv.size == (params["height"], params["width"])
+        return cls(size=random_crop_tv.size, _top=params["top"], _left=params["left"])

test/test_transform_ops.py

@@ -145,6 +145,49 @@ def test_resize_fails(self):
         ):
             VideoDecoder(NASA_VIDEO.path, transforms=[v2.Resize(size=(100))])
 
+    @pytest.mark.parametrize(
+        "height_scaling_factor, width_scaling_factor",
+        ((0.5, 0.5), (0.25, 0.1)),
+    )
+    @pytest.mark.parametrize("video", [NASA_VIDEO, TEST_SRC_2_720P])
+    def test_random_crop_torchvision(
+        self, video, height_scaling_factor, width_scaling_factor
+    ):
+        height = int(video.get_height() * height_scaling_factor)
+        width = int(video.get_width() * width_scaling_factor)
+
+        torch.manual_seed(0)
+        tc_random_crop = torchcodec.transforms.RandomCrop(size=(height, width))
+        decoder_random_crop = VideoDecoder(video.path, transforms=[tc_random_crop])
+
+        torch.manual_seed(0)
+        decoder_random_crop_tv = VideoDecoder(
+            video.path, transforms=[v2.RandomCrop(size=(height, width))]
+        )
+
+        decoder_full = VideoDecoder(video.path)
+
+        num_frames = len(decoder_random_crop_tv)
+        assert num_frames == len(decoder_full)
+
+        for frame_index in [
+            0,
+            int(num_frames * 0.1),
+            int(num_frames * 0.2),
+            int(num_frames * 0.3),
+            int(num_frames * 0.4),
+            int(num_frames * 0.5),
+            int(num_frames * 0.75),
+            int(num_frames * 0.90),
+            num_frames - 1,
+        ]:
+            frame_random_crop = decoder_random_crop[frame_index]
+            frame_random_crop_tv = decoder_random_crop_tv[frame_index]
+            assert_frames_equal(frame_random_crop, frame_random_crop_tv)
+
+            expected_shape = (video.get_num_color_channels(), height, width)
+            assert frame_random_crop_tv.shape == expected_shape
+
     def test_transform_fails(self):
         with pytest.raises(
             ValueError,