Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add target_sampling_rate parameter to BlissToPcmHDFJob #568

Merged
merged 2 commits into from
Jan 14, 2025
Merged

Add target_sampling_rate parameter to BlissToPcmHDFJob #568

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
79e8080
add resampling to BlissToPcmHDFJob
michelwi Jan 10, 2025
88819b0
add librosa to requirements.txt
michelwi Jan 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
black==22.3.0
h5py
librosa
matplotlib
numpy
soundfile
Expand Down
72 changes: 44 additions & 28 deletions returnn/hdf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from enum import Enum, auto
import glob
import math
import librosa
import numpy as np
import os
import shutil
Expand Down Expand Up @@ -230,6 +231,7 @@ class RoundingScheme(Enum):
"multi_channel_strategy": BaseStrategy(),
"rounding": RoundingScheme.start_and_duration,
"round_factor": 1,
"target_sampling_rate": None,
}

def __init__(
Expand All @@ -241,6 +243,7 @@ def __init__(
returnn_root: Optional[tk.Path] = None,
rounding: RoundingScheme = RoundingScheme.start_and_duration,
round_factor: int = 1,
target_sampling_rate: Optional[int] = None,
):
"""

Expand All @@ -256,6 +259,7 @@ def __init__(
start_and_duration will round down the start time and the duration of the segment
rasr_compatible will round up the start time and round down the end time
:param round_factor: do the rounding based on a sampling rate that is scaled down by this factor
:param target_sampling_rate: desired sampling rate for the HDF, data will be resampled to this rate if needed
"""
self.set_vis_name("Dump audio to HDF")
assert output_dtype in ["float64", "float32", "int32", "int16"]
Expand All @@ -267,6 +271,7 @@ def __init__(
self.returnn_root = returnn_root
self.rounding = rounding
self.round_factor = round_factor
self.target_sampling_rate = target_sampling_rate

self.out_hdf = self.output_path("audio.hdf")

Expand All @@ -284,7 +289,7 @@ def run(self):

if self.segment_file:
with uopen(self.segment_file, "rt") as f:
segments_whitelist = set(l.strip() for l in f.readlines() if len(l.strip()) > 0)
segments_whitelist = {line.strip() for line in f.readlines() if len(line.strip()) > 0}
else:
segments_whitelist = None

Expand All @@ -295,35 +300,46 @@ def run(self):
audio = sf.SoundFile(audio_file)

for segment in recording.segments:
if (not segments_whitelist) or (segment.fullname() in segments_whitelist):
if self.rounding == self.RoundingScheme.start_and_duration:
start = int(segment.start * audio.samplerate / self.round_factor) * self.round_factor
duration = (
int((segment.end - segment.start) * audio.samplerate / self.round_factor)
* self.round_factor
)
elif self.rounding == self.RoundingScheme.rasr_compatible:
start = math.ceil(segment.start * audio.samplerate / self.round_factor) * self.round_factor
duration = (
math.floor(segment.end * audio.samplerate / self.round_factor) * self.round_factor - start
)
else:
raise NotImplementedError(f"RoundingScheme {self.rounding} not implemented.")
audio.seek(start)
data = audio.read(
duration,
always_2d=True,
dtype=self.output_dtype,
if (segments_whitelist is not None) and (segment.fullname() not in segments_whitelist):
continue

# determine correct start and duration values
if self.rounding == self.RoundingScheme.start_and_duration:
start = int(segment.start * audio.samplerate / self.round_factor) * self.round_factor
duration = (
int((segment.end - segment.start) * audio.samplerate / self.round_factor) * self.round_factor
)
if isinstance(self.multi_channel_strategy, self.PickNth):
data = data[:, self.multi_channel_strategy.channel]
else:
assert data.shape[-1] == 1, "Audio has more than one channel, choose a multi_channel_strategy"
out_hdf.insert_batch(
inputs=data.reshape(1, -1, 1),
seq_len=[data.shape[0]],
seq_tag=[segment.fullname()],
elif self.rounding == self.RoundingScheme.rasr_compatible:
start = math.ceil(segment.start * audio.samplerate / self.round_factor) * self.round_factor
duration = (
math.floor(segment.end * audio.samplerate / self.round_factor) * self.round_factor - start
)
else:
raise NotImplementedError(f"RoundingScheme {self.rounding} not implemented.")

# read audio data
audio.seek(start)
data = audio.read(duration, always_2d=True, dtype=self.output_dtype)
if isinstance(self.multi_channel_strategy, self.PickNth):
data = data[:, self.multi_channel_strategy.channel]
else:
assert data.shape[-1] == 1, "Audio has more than one channel, choose a multi_channel_strategy"

# resample if necessary
if (sr := self.target_sampling_rate) is not None and sr != audio.samplerate:
data = librosa.resample(
y=data.astype(float),
orig_sr=audio.samplerate,
target_sr=sr,
axis=0,
).astype(self.output_dtype)

# add audio to hdf
out_hdf.insert_batch(
inputs=data.reshape(1, -1, 1),
seq_len=[data.shape[0]],
seq_tag=[segment.fullname()],
)

audio.close()

Expand Down
Loading