wehub-resource-sync
58 lines
1.5 KiB
Python
58 lines
1.5 KiB
Python
import io
|
|
from typing import TYPE_CHECKING, Iterator, List, Union
|
|
|
|
from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
|
|
from ray.data._internal.util import _check_import
|
|
from ray.data.block import Block
|
|
from ray.data.datasource.file_based_datasource import FileBasedDatasource
|
|
|
|
if TYPE_CHECKING:
|
|
import pyarrow
|
|
|
|
|
|
class AudioDatasource(FileBasedDatasource):
|
|
_FILE_EXTENSIONS = [
|
|
"mp3",
|
|
"wav",
|
|
"aac",
|
|
"flac",
|
|
"ogg",
|
|
"m4a",
|
|
"wma",
|
|
"alac",
|
|
"aiff",
|
|
"pcm",
|
|
"amr",
|
|
"opus",
|
|
"ra",
|
|
"rm",
|
|
"au",
|
|
"mid",
|
|
"midi",
|
|
"caf",
|
|
]
|
|
|
|
def __init__(
|
|
self,
|
|
paths: Union[str, List[str]],
|
|
**file_based_datasource_kwargs,
|
|
):
|
|
super().__init__(paths, **file_based_datasource_kwargs)
|
|
|
|
_check_import(self, module="soundfile", package="soundfile")
|
|
|
|
def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
|
|
import soundfile
|
|
|
|
# `soundfile` doesn't support reading from a `pyarrow.NativeFile` directly, so
|
|
# we need to read the file into memory first.
|
|
stream = io.BytesIO(f.read())
|
|
amplitude, sample_rate = soundfile.read(stream, always_2d=True, dtype="float32")
|
|
|
|
# (amplitude, channels) -> (channels, amplitude)
|
|
amplitude = amplitude.transpose((1, 0))
|
|
|
|
builder = DelegatingBlockBuilder()
|
|
builder.add({"amplitude": amplitude, "sample_rate": sample_rate})
|
|
yield builder.build()
|