Added width and height options for ResolutionSubsampler (#274)

* Added width and height parameters for ResolutionSubsampler

* Added tests

* Fixed linting

* Linting

* Black formatting

---------

Co-authored-by: iejMac <kilianmaciej6@gmail.com>

tests/test_subsamplers.py

@@ -97,14 +97,14 @@ def test_clipping_subsampler(clips):
 
 
 @pytest.mark.parametrize("size,resize_mode", [(144, ["scale"]), (1620, ["scale", "crop", "pad"])])
-def test_resolution_subsampler(size, resize_mode):
+def test_resolution_subsampler_video_size(size, resize_mode):
     current_folder = os.path.dirname(__file__)
     # video lenght - 2:02, 1080x1920
     video = os.path.join(current_folder, "test_files/test_video.mp4")
     with open(video, "rb") as vid_f:
         video_bytes = vid_f.read()
 
-    subsampler = ResolutionSubsampler(size, resize_mode)
+    subsampler = ResolutionSubsampler(video_size=size, resize_mode=resize_mode)
 
     streams = {"video": [video_bytes]}
     subsampled_streams, _, error_message = subsampler(streams)
@@ -125,6 +125,36 @@ def test_resolution_subsampler(size, resize_mode):
             assert w_vid == size
 
 
+@pytest.mark.parametrize("height,width,resize_mode", [(-1,128, ["scale"]), (1620,1620, ["scale", "crop", "pad"])])
+def test_resolution_subsampler_height_and_width(height, width, resize_mode):
+    current_folder = os.path.dirname(__file__)
+    # video lenght - 2:02, 1080x1920
+    video = os.path.join(current_folder, "test_files/test_video.mp4")
+    with open(video, "rb") as vid_f:
+        video_bytes = vid_f.read()
+
+    subsampler = ResolutionSubsampler(height=height, width=width, resize_mode=resize_mode)
+
+    streams = {"video": [video_bytes]}
+    subsampled_streams, _, error_message = subsampler(streams)
+    assert error_message is None
+    subsampled_videos = subsampled_streams["video"]
+
+    with tempfile.NamedTemporaryFile() as tmp:
+        tmp.write(subsampled_videos[0])
+
+        probe = ffmpeg.probe(tmp.name)
+        video_stream = [stream for stream in probe["streams"] if stream["codec_type"] == "video"][0]
+        h_vid, w_vid = video_stream["height"], video_stream["width"]
+
+        if resize_mode == ["scale"]:
+            assert h_vid == 72
+            assert w_vid == 128
+        else:
+            assert h_vid == height
+            assert w_vid == width
+
+
 @pytest.mark.parametrize("target_frame_rate", [6, 15, 30])
 def test_frame_rate_subsampler(target_frame_rate):
     current_folder = os.path.dirname(__file__)

video2dataset/subsamplers/resolution_subsampler.py

@@ -4,26 +4,42 @@
 import os
 import ffmpeg
 import tempfile
+from typing import Literal
 
 from .subsampler import Subsampler
 
 
 class ResolutionSubsampler(Subsampler):
     """
     Adjusts the resolution of the videos to the specified height and width.
+    Please do not set both video_size and height/width. This will result in an error.
+    If both height and width are set, scale mode output will have the specified height (ignoring width).
 
     Args:
-        video_size (int): Target resolution of the videos.
         resize_mode (list[str]): List of resize modes to apply. Possible options are:
-            scale: scale video keeping aspect ratios (currently always picks video height)
+            scale: scale video keeping aspect ratios
             crop: center crop to video_size x video_size
             pad: center pad to video_size x video_size
+        height (int): Height of video.
+        width (int): Width of video.
+        video_size (int): Both height and width.
         encode_format (str): Format to encode in (i.e. mp4)
     """
 
-    def __init__(self, video_size, resize_mode, encode_format="mp4"):
-        self.video_size = video_size
+    def __init__(
+        self,
+        resize_mode: Literal["scale", "crop", "pad"],
+        height: int = -1,
+        width: int = -1,
+        video_size: int = -1,
+        encode_format: str = "mp4",
+    ):
+        if video_size > 0 and (height > 0 or width > 0):
+            raise Exception("Either set video_size, or set height and/or width")
         self.resize_mode = resize_mode
+        self.height = height if video_size < 0 else video_size
+        self.width = width if video_size < 0 else video_size
+        self.video_size = video_size
         self.encode_format = encode_format
 
     def __call__(self, streams, metadata=None):
@@ -36,11 +52,14 @@ def __call__(self, streams, metadata=None):
                 try:
                     _ = ffmpeg.input(f"{tmpdir}/input.mp4")
                     if "scale" in self.resize_mode:
-                        _ = _.filter("scale", -2, self.video_size)
+                        if self.height > 0:
+                            _ = _.filter("scale", -2, self.height)
+                        else:
+                            _ = _.filter("scale", self.width, -2)
                     if "crop" in self.resize_mode:
-                        _ = _.filter("crop", w=self.video_size, h=self.video_size)
+                        _ = _.filter("crop", w=self.width, h=self.height)
                     if "pad" in self.resize_mode:
-                        _ = _.filter("pad", w=self.video_size, h=self.video_size)
+                        _ = _.filter("pad", w=self.width, h=self.height)
                     _ = _.output(f"{tmpdir}/output.mp4", reset_timestamps=1).run(capture_stdout=True, quiet=True)
                 except Exception as err:  # pylint: disable=broad-except
                     return [], None, str(err)