def __getitem__(self, key):
"""Returns the specified piece of this IOTensor."""
if isinstance(key, slice):
return core_ops.io_audio_readable_read(
self._resource, key.start, key.stop, dtype=self._dtype)
item = core_ops.io_audio_readable_read(
self._resource, key, key + 1, dtype=self._dtype)
if tf.shape(item)[0] == 0:
raise IndexError("index %s is out of range" % key)
return item[0]
def __init__(self, filename, dtype=None):
"""AudioIODataset."""
with tf.name_scope("AudioIODataset"):
if not tf.executing_eagerly():
assert dtype is not None, "dtype must be provided in graph mode"
resource = core_ops.io_audio_readable_init(filename)
if tf.executing_eagerly():
shape, dtype, _ = core_ops.io_audio_readable_spec(resource)
dtype = tf.as_dtype(dtype.numpy())
else:
shape, _, _ = core_ops.io_audio_readable_spec(resource)
capacity = 1024 # kwargs.get("capacity", 4096)
self._resource = resource
dataset = tf.data.Dataset.range(0, shape[0], capacity)
dataset = dataset.map(
lambda index: core_ops.io_audio_readable_read(
resource, index, index + capacity, dtype=dtype))
dataset = dataset.apply(
tf.data.experimental.take_while(
lambda v: tf.greater(tf.shape(v)[0], 0)))
dataset = dataset.unbatch()
self._dataset = dataset
super().__init__(self._dataset._variant_tensor) # pylint: disable=protected-access
def to_tensor(self):
"""Converts this `IOTensor` into a `tf.Tensor`.
Args:
name: A name prefix for the returned tensors (optional).
Returns:
A `Tensor` with value obtained from this `IOTensor`.
"""
return core_ops.io_audio_readable_read(self._resource, 0, -1, dtype=self._dtype)
def __init__(self, resource, shape, dtype, internal=True):
"""AudioGraphIODataset."""
with tf.name_scope("AudioGraphIODataset"):
assert internal
capacity = 1024 #kwargs.get("capacity", 4096)
self._resource = resource
dataset = tf.data.Dataset.range(0, shape[0], capacity)
dataset = dataset.map(
lambda index: core_ops.io_audio_readable_read(
resource, index, index + capacity, dtype=dtype))
dataset = dataset.apply(
tf.data.experimental.take_while(
lambda v: tf.greater(tf.shape(v)[0], 0)))
dataset = dataset.unbatch()
self._dataset = dataset
super(AudioGraphIODataset,
self).__init__(self._dataset._variant_tensor) # pylint: disable=protected-access
def __getitem__(self, key):
"""Returns the specified piece of this IOTensor."""
# always convert to tuple to process
if not isinstance(key, tuple):
key = tuple([key])
# get the start and stop of each element
indices = [(k.start, k.stop) if isinstance(k, slice) else (k, k + 1)
for k in key]
# get the start and stop, and use 0 (start) and -1 (stop) if needed
indices = list(zip(*indices))
start = [0 if e is None else e for e in indices[0]]
stop = [-1 if e is None else e for e in indices[1]]
item = core_ops.io_audio_readable_read(self._resource,
start=start,
stop=stop,
dtype=self._dtype)
# in case certain dimension is not slice, then this dimension will need to
# collapse as `0`, otherwise `:` or `slice(None, None, None)`
indices = [slice(None) if isinstance(k, slice) else 0 for k in key]
return item.__getitem__(indices)
