import io
import logging
import os.path
from aim.sdk.num_utils import inst_has_typename
from aim.sdk.objects.io import wavfile
from aim.storage.object import CustomObject
from aim.storage.types import BLOB
logger = logging.getLogger(__name__)
[docs]@CustomObject.alias('aim.audio')
class Audio(CustomObject):
"""Audio object used to store audio objects in Aim repository..
Currently, audio formats are limited to mp3, wav, flac
Args:
data: file path, bytes, io.BaseIO or numpy.array (only for WAV)
format (:obj:`str`): Format of the audio source
rate (:obj:`int`): Rate of the audio file, for WAV defaults to 22500
caption (:obj:`str`, optional): Optional audio caption. '' by default.
"""
AIM_NAME = 'aim.audio'
# supported audio formats
UNKNOWN = ''
MP3 = 'mp3'
WAV = 'wav'
FLAC = 'flac'
audio_formats = (MP3, WAV, FLAC)
def __init__(self, data, format: str = '', caption: str = '', rate: int = None):
super().__init__()
audio_format = format.lower()
if inst_has_typename(data, ['ndarray.numpy']):
# Currently, only WAV audio formats are supported for numpy
audio_format = self.WAV
if not rate:
rate = 22500
logger.info(f'Parameter "rate" is not provided! Using default: {rate}')
bs = wavfile.write(rate, data)
data = bs
# act as a regular file with enforced audio format definition by user side
if not audio_format:
raise ValueError('Audio format must be provided.')
elif audio_format not in self.audio_formats:
raise ValueError(f'Invalid audio format is provided. Must be one of {self.audio_formats}')
if isinstance(data, str):
if not os.path.exists(data) or not os.path.isfile(data):
raise ValueError('Invalid audio file path')
with open(data, 'rb') as FS:
data = FS.read()
elif isinstance(data, io.BytesIO):
data = data.read()
if not isinstance(data, bytes):
raise TypeError('Content is not a byte-stream object')
extra = {
'caption': caption,
'format': audio_format
}
self._prepare(data, **extra)
[docs] def json(self):
"""Dump audio metadata to a dict"""
return {
'caption': self.storage['caption'],
'format': self.storage['format']
}
def _prepare(self, data, **extra) -> None:
assert isinstance(data, bytes)
for k, v in extra.items():
self.storage[k] = v
self.storage['data'] = BLOB(data=data)
[docs] def to_numpy(self):
"""
This method converts WAV to Numpy array.
Other audio formats are not supported at this moment.
Returns: numpy array
"""
assert self.storage['format'] == self.WAV
return wavfile.read(self.get())
[docs] def get(self) -> io.BytesIO:
"""
Reads data from the inner container and writes it to a buffer
Returns: io.BytesIO
"""
bs = self.storage.get('data')
if not bs:
return io.BytesIO()
return io.BytesIO(bytes(bs))