Add VP9 support

Author: NicolasHug

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

@@ -150,10 +150,11 @@ cudaVideoCodec validateCodecSupport(AVCodecID codecId) {
       return cudaVideoCodec_HEVC;
     case AV_CODEC_ID_AV1:
       return cudaVideoCodec_AV1;
+    case AV_CODEC_ID_VP9:
+      return cudaVideoCodec_VP9;
     // TODONVDEC P0: support more codecs
     // case AV_CODEC_ID_MPEG4: return cudaVideoCodec_MPEG4;
     // case AV_CODEC_ID_VP8: return cudaVideoCodec_VP8;
-    // case AV_CODEC_ID_VP9: return cudaVideoCodec_VP9;
     // case AV_CODEC_ID_MJPEG: return cudaVideoCodec_JPEG;
     default: {
       TORCH_CHECK(false, "Unsupported codec type: ", avcodec_get_name(codecId));

test/resources/testsrc2_vp9.webm

@@ -45,6 +45,7 @@
     SINE_MONO_S32_8000,
     TEST_SRC_2_720P,
     TEST_SRC_2_720P_H265,
+    TEST_SRC_2_720P_VP9,
     unsplit_device_str,
 )
 
@@ -588,7 +589,7 @@ def test_get_frame_at_av1(self, device):
             return
 
         if device == "cuda" and in_fbcode():
-            pytest.skip("AV1 decoding on CUDA is not supported internally")
+            pytest.skip("decoding on CUDA is not supported internally")
 
         decoder = VideoDecoder(AV1_VIDEO.path, device=device)
         device, _ = unsplit_device_str(device)
@@ -1451,15 +1452,16 @@ def test_get_frames_at_tensor_indices(self):
             BT709_FULL_RANGE,
             TEST_SRC_2_720P_H265,
             AV1_VIDEO,
+            TEST_SRC_2_720P_VP9,
         ),
     )
     @pytest.mark.parametrize("contiguous_indices", (True, False))
     @pytest.mark.parametrize("seek_mode", ("exact", "approximate"))
     def test_beta_cuda_interface_get_frame_at(
         self, asset, contiguous_indices, seek_mode
     ):
-        if asset == AV1_VIDEO and seek_mode == "approximate":
-            pytest.skip("AV1 asset doesn't work with approximate mode")
+        if asset in (AV1_VIDEO, TEST_SRC_2_720P_VP9) and seek_mode == "approximate":
+            pytest.skip("asset doesn't work with approximate mode")
 
         ref_decoder = VideoDecoder(asset.path, device="cuda", seek_mode=seek_mode)
         beta_decoder = VideoDecoder(
@@ -1493,15 +1495,16 @@ def test_beta_cuda_interface_get_frame_at(
             BT709_FULL_RANGE,
             TEST_SRC_2_720P_H265,
             AV1_VIDEO,
+            TEST_SRC_2_720P_VP9,
         ),
     )
     @pytest.mark.parametrize("contiguous_indices", (True, False))
     @pytest.mark.parametrize("seek_mode", ("exact", "approximate"))
     def test_beta_cuda_interface_get_frames_at(
         self, asset, contiguous_indices, seek_mode
     ):
-        if asset == AV1_VIDEO and seek_mode == "approximate":
-            pytest.skip("AV1 asset doesn't work with approximate mode")
+        if asset in (AV1_VIDEO, TEST_SRC_2_720P_VP9) and seek_mode == "approximate":
+            pytest.skip("asset doesn't work with approximate mode")
 
         ref_decoder = VideoDecoder(asset.path, device="cuda", seek_mode=seek_mode)
         beta_decoder = VideoDecoder(
@@ -1536,12 +1539,13 @@ def test_beta_cuda_interface_get_frames_at(
             BT709_FULL_RANGE,
             TEST_SRC_2_720P_H265,
             AV1_VIDEO,
+            TEST_SRC_2_720P_VP9,
         ),
     )
     @pytest.mark.parametrize("seek_mode", ("exact", "approximate"))
     def test_beta_cuda_interface_get_frame_played_at(self, asset, seek_mode):
-        if asset == AV1_VIDEO and seek_mode == "approximate":
-            pytest.skip("AV1 asset doesn't work with approximate mode")
+        if asset in (AV1_VIDEO, TEST_SRC_2_720P_VP9) and seek_mode == "approximate":
+            pytest.skip("asset doesn't work with approximate mode")
 
         ref_decoder = VideoDecoder(asset.path, device="cuda", seek_mode=seek_mode)
         beta_decoder = VideoDecoder(
@@ -1572,13 +1576,14 @@ def test_beta_cuda_interface_get_frame_played_at(self, asset, seek_mode):
             TEST_SRC_2_720P,
             BT709_FULL_RANGE,
             TEST_SRC_2_720P_H265,
+            TEST_SRC_2_720P_VP9,
             AV1_VIDEO,
         ),
     )
     @pytest.mark.parametrize("seek_mode", ("exact", "approximate"))
     def test_beta_cuda_interface_get_frames_played_at(self, asset, seek_mode):
-        if asset == AV1_VIDEO and seek_mode == "approximate":
-            pytest.skip("AV1 asset doesn't work with approximate mode")
+        if asset in (AV1_VIDEO, TEST_SRC_2_720P_VP9) and seek_mode == "approximate":
+            pytest.skip("asset doesn't work with approximate mode")
 
         ref_decoder = VideoDecoder(asset.path, device="cuda", seek_mode=seek_mode)
         beta_decoder = VideoDecoder(
@@ -1611,12 +1616,13 @@ def test_beta_cuda_interface_get_frames_played_at(self, asset, seek_mode):
             BT709_FULL_RANGE,
             TEST_SRC_2_720P_H265,
             AV1_VIDEO,
+            TEST_SRC_2_720P_VP9,
         ),
     )
     @pytest.mark.parametrize("seek_mode", ("exact", "approximate"))
     def test_beta_cuda_interface_backwards(self, asset, seek_mode):
-        if asset == AV1_VIDEO and seek_mode == "approximate":
-            pytest.skip("AV1 asset doesn't work with approximate mode")
+        if asset in (AV1_VIDEO, TEST_SRC_2_720P_VP9) and seek_mode == "approximate":
+            pytest.skip("asset doesn't work with approximate mode")
 
         ref_decoder = VideoDecoder(asset.path, device="cuda", seek_mode=seek_mode)
         beta_decoder = VideoDecoder(

test/test_decoders.py

@@ -717,3 +717,13 @@ def sample_format(self) -> str:
     },
     frames={0: {}},  # Not needed for now
 )
+
+# ffmpeg -f lavfi -i testsrc2=size=1280x720:rate=30:duration=1 -c:v libvpx-vp9 -b:v 1M output_vp9.webm
+TEST_SRC_2_720P_VP9 = TestVideo(
+    filename="testsrc2_vp9.webm",
+    default_stream_index=0,
+    stream_infos={
+        0: TestVideoStreamInfo(width=1280, height=720, num_color_channels=3),
+    },
+    frames={0: {}},  # Not needed for now
+)