Torchcodec decoding

[TyTodd]

Closes #7607

New signatures

Audio

Audio(sampling_rate: Optional[int] = None, mono: bool = True, decode: bool = True, stream_index: Optional[int] = None)

Audio.encode_example(self, value: Union[str, bytes, bytearray, dict, "AudioDecoder"]) -> dict

Audio.decode_example(self, value: dict, token_per_repo_id: Optional[dict[str, Union[str, bool, None]]] = None) -> "AudioDecoder":

Video

Video(decode: bool = True, stream_index: Optional[int] = None, dimension_order: Literal['NCHW', 'NHWC'] = 'NCHW', num_ffmpeg_threads: int = 1, device: Optional[Union[str, "torch.device"]] = 'cpu', seek_mode: Literal['exact', 'approximate'] = 'exact')

Video.encode_example(self, value: Union[str, bytes, bytearray, Example, np.ndarray, "VideoDecoder"]) -> Example:

Video.decode_example(self, value: Union[str, Example], token_per_repo_id: Optional[dict[str, Union[bool, str]]] = None, ) -> "VideoDecoder":

Notes

Audio features constructor takes in 1 new optional param stream_index which is passed to the AudioDecoder constructor to select the stream index of a file.
Audio feature can now take in torchcodec.decoders.AudioDecoder as input to encode_example()
Audio feature decode_example() returns torchcodec.decoders.AudioDecoder

Video feature constructor takes in 5 new optional params stream_index, dimension_order, num_ffmpeg_threads, device, seek_mode all of which are passed to VideoDecoder constructor
Video feature decode_example() returns torchcodec.decoders.VideoDecoder
Video feature can now take in torchcodec.decoders.VideoDecoder as input to encode_example()

All test cases have been updated to reflect these changes
All documentation has also been updated to reflect these changes.

Both VideoDecoder and AudioDecoder when formatted with (np_formatter, tf_formatter, etc) will ignore the type and return themselves. Formatting test cases were updated accordingly to reflect this. (Pretty simple to make this not the case if we want though)

Errors

This test case from tests/packaged_modules/test_audiofolder.py

@require_librosa
@require_sndfile
@pytest.mark.parametrize("streaming", [False, True])
def test_data_files_with_metadata_and_archives(streaming, cache_dir, data_files_with_zip_archives):
    audiofolder = AudioFolder(data_files=data_files_with_zip_archives, cache_dir=cache_dir)
    audiofolder.download_and_prepare()
    datasets = audiofolder.as_streaming_dataset() if streaming else audiofolder.as_dataset()
    for split, data_files in data_files_with_zip_archives.items():
        num_of_archives = len(data_files)  # the metadata file is inside the archive
        expected_num_of_audios = 2 * num_of_archives
        assert split in datasets
        dataset = list(datasets[split])
        assert len(dataset) == expected_num_of_audios
        # make sure each sample has its own audio (all arrays are different) and metadata
        assert (
            sum(np.array_equal(dataset[0]["audio"].get_all_samples().data.numpy(), example["audio"].get_all_samples().data.numpy()) for example in dataset[1:])
            == 0
        )
        assert len({example["text"] for example in dataset}) == expected_num_of_audios
        assert all(example["text"] is not None for example in dataset)

Fails now because AudioDecoder needs to access the files after the lines below are run, but there seems to be some context issues. The file the decoder is trying to read is closed before the decoder gets the chance to decode it.

audiofolder.download_and_prepare()
datasets = audiofolder.as_streaming_dataset() if streaming else audiofolder.as_dataset()

[TyTodd]

@lhoestq any updates on when this will be merged? Let me know if theres anything you need from my end.

[lhoestq]

Awesome ! I added a few comments :)

My main comment is about backward compatibility for audio, the rest looks good to me

[lhoestq]

We also need _hf_encoded in audio now that it's also a reader (useful to speed up re-encoding when saving the a dataset)

[lhoestq]

for backward compatibility we'll need to subclass AudioDecoder to allow this:

audio_numpy_array = decorder["array"]
audio_sampling_rate = decorder["sampling_rate"]

[TyTodd]

Just to make sure I'm understanding correctly, are you saying to instead return a object of a class that extends torchcodec.decoders._audio_decoder.AudioDecoder and overwrites the getitem() method?

[lhoestq]

Yes, hopefully torchcodec allows subclassing.
This is pretty necessary for audio since most training frameworks or scripts use the ["array"] syntax at the moment.

[lhoestq]

Btw I plan to release datasets 4.0 after your PR, this will be a major milestone :)

[TyTodd]

@lhoestq just pushed the new changes.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[lhoestq]

Great ! I took the liberty to move the AudioDecoder to its own file and make small edits in the docs and docstrings

If it looks good to you I think we can merge :)

[lhoestq]

Alright this is ready ! Congrats on implementing this :)

I might do another pass at documentation before doing the new datasets release but this is already quite in good shape