384 lines
15 KiB
Python
384 lines
15 KiB
Python
import math
|
|
import os
|
|
import random
|
|
import torch
|
|
import torch.utils.data
|
|
import numpy as np
|
|
from librosa.util import normalize
|
|
from scipy.io.wavfile import read
|
|
import scipy
|
|
import librosa
|
|
import wave
|
|
from pydub import AudioSegment
|
|
|
|
MAX_WAV_VALUE = 32768.0
|
|
|
|
|
|
def load_wav(full_path):
|
|
try:
|
|
sampling_rate, data = read(full_path)
|
|
if max(data.shape) / sampling_rate < 0.5:
|
|
return None, None
|
|
except FileNotFoundError:
|
|
print(f"File not found: {file_path}")
|
|
return None, None
|
|
except Exception as e:
|
|
print(f"An unexpected error occurred: {e}")
|
|
return None, None
|
|
|
|
if len(data.shape) > 1:
|
|
if data.shape[1] <= 2:
|
|
data = data[...,0]
|
|
else:
|
|
data = data[0,...]
|
|
return data / MAX_WAV_VALUE, sampling_rate
|
|
|
|
def get_wave_duration(file_path):
|
|
"""
|
|
Gets the duration of a WAV file in seconds.
|
|
|
|
:param file_path: Path to the WAV file.
|
|
:return: Duration of the WAV file in seconds.
|
|
"""
|
|
try:
|
|
with wave.open(file_path, 'rb') as wf:
|
|
# Get the number of frames
|
|
num_frames = wf.getnframes()
|
|
# Get the frame rate
|
|
frame_rate = wf.getframerate()
|
|
# Calculate duration
|
|
duration = num_frames / float(frame_rate)
|
|
return duration, frame_rate, num_frames
|
|
except wave.Error as e:
|
|
print(f"Error reading {file_path}: {e}")
|
|
return None, None, None
|
|
except FileNotFoundError:
|
|
print(f"File not found: {file_path}")
|
|
return None, None, None
|
|
except Exception as e:
|
|
print(f"An unexpected error occurred: {e}")
|
|
return None, None, None
|
|
|
|
def read_audio_segment(file_path, start_ms, end_ms):
|
|
"""
|
|
Reads a segment from a WAV file and returns the raw data and its properties.
|
|
|
|
:param file_path: Path to the WAV file.
|
|
:param start_ms: Start time of the segment in milliseconds.
|
|
:param end_ms: End time of the segment in milliseconds.
|
|
:return: A tuple containing the raw audio data, frame rate, sample width, and number of channels.
|
|
"""
|
|
#start_time = time.time()
|
|
try:
|
|
# Load the audio file
|
|
audio = AudioSegment.from_wav(file_path)
|
|
# Extract the segment
|
|
segment = audio[start_ms:end_ms]
|
|
# Get raw audio data
|
|
raw_data = segment.raw_data
|
|
# Get audio properties
|
|
frame_rate = segment.frame_rate
|
|
sample_width = segment.sample_width
|
|
channels = segment.channels
|
|
# Create NumPy array from the raw audio data
|
|
audio_array = np.frombuffer(raw_data, dtype=np.int16)
|
|
|
|
# If stereo, reshape the array to have a second dimension
|
|
if channels > 1:
|
|
audio_array = audio_array.reshape((-1, channels))
|
|
audio_array = audio_array[...,0]
|
|
'''
|
|
if frame_rate !=48000:
|
|
audio_array = audio_array/MAX_WAV_VALUE
|
|
audio_array = librosa.resample(audio_array, frame_rate, 48000)
|
|
audio_array = audio_array * MAX_WAV_VALUE
|
|
frame_rate = 48000
|
|
'''
|
|
#end_time = time.time()
|
|
#time_taken = end_time - start_time
|
|
|
|
#print(f"Successfully read segment from {start_ms}ms to {end_ms}ms in {time_taken:.4f} seconds")
|
|
return audio_array / MAX_WAV_VALUE#, frame_rate #, sample_width, channels
|
|
except Exception as e:
|
|
print(f"An error occurred: {e}")
|
|
return None#, None #, None, None
|
|
|
|
def resample(audio, sr_in, sr_out, target_len=None):
|
|
#audio = audio / MAX_WAV_VALUE
|
|
#audio = normalize(audio) * 0.95
|
|
if target_len is not None:
|
|
audio = scipy.signal.resample(audio, target_len)
|
|
return audio
|
|
resample_factor = sr_out / sr_in
|
|
new_samples = int(len(audio) * resample_factor)
|
|
audio = scipy.signal.resample(audio, new_samples)
|
|
return audio
|
|
|
|
def load_segment(full_path, target_sampling_rate=None, segment_size=None):
|
|
|
|
if segment_size is not None:
|
|
dur,sampling_rate,len_data = get_wave_duration(full_path)
|
|
if sampling_rate is None: return None, None
|
|
if sampling_rate < 44100: return None, None
|
|
|
|
target_dur = segment_size / target_sampling_rate
|
|
if dur < target_dur:
|
|
data, sampling_rate = load_wav(full_path)
|
|
#print(f'data_read: {data.shape}, sampling_rate: {sampling_rate}')
|
|
if data is None: return None, None
|
|
|
|
if target_sampling_rate is not None and sampling_rate != target_sampling_rate:
|
|
data = resample(data, sampling_rate, target_sampling_rate)
|
|
sampling_rate = target_sampling_rate
|
|
data = torch.FloatTensor(data)
|
|
data = data.unsqueeze(0)
|
|
data = torch.nn.functional.pad(data, (0, segment_size - data.size(1)), 'constant')
|
|
data = data.squeeze(0)
|
|
return data.numpy(), sampling_rate
|
|
else:
|
|
dur,sampling_rate,len_data = get_wave_duration(full_path)
|
|
if sampling_rate < 44100: return None, None
|
|
|
|
target_dur = segment_size / target_sampling_rate
|
|
target_len = int(target_dur * sampling_rate)
|
|
start_idx = random.randint(0, (len_data - target_len))
|
|
start_ms = start_idx / sampling_rate * 1000
|
|
end_ms = start_ms + target_dur * 1000
|
|
data = read_audio_segment(full_path, start_ms, end_ms)
|
|
#print(f'data_read: {data.shape}, sampling_rate: {sampling_rate}')
|
|
if data is None: return None, None
|
|
if target_sampling_rate is not None and sampling_rate != target_sampling_rate:
|
|
data = resample(data, sampling_rate, target_sampling_rate)
|
|
sampling_rate = target_sampling_rate
|
|
if len(data) < segment_size:
|
|
data = torch.FloatTensor(data)
|
|
data = data.unsqueeze(0)
|
|
data = torch.nn.functional.pad(data, (0, segment_size - data.size(1)), 'constant')
|
|
data = data.squeeze(0)
|
|
data = data.numpy()
|
|
else:
|
|
start_idx = random.randint(0, (len(data) - segment_size))
|
|
data = data[start_idx:start_idx+segment_size]
|
|
#print(f'data_cut: {data.shape}')
|
|
return data, sampling_rate
|
|
else:
|
|
dur,sampling_rate,len_data = get_wave_duration(full_path)
|
|
if sampling_rate is None: return None, None
|
|
if sampling_rate < 44100: return None, None
|
|
data, sampling_rate = load_wav(full_path)
|
|
if data is None: return None, None
|
|
if target_sampling_rate is not None and sampling_rate != target_sampling_rate:
|
|
data = resample(data, sampling_rate, target_sampling_rate)
|
|
sampling_rate = target_sampling_rate
|
|
return data, sampling_rate
|
|
|
|
def dynamic_range_compression(x, C=1, clip_val=1e-5):
|
|
return np.log(np.clip(x, a_min=clip_val, a_max=None) * C)
|
|
|
|
|
|
def dynamic_range_decompression(x, C=1):
|
|
return np.exp(x) / C
|
|
|
|
|
|
def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
|
|
return torch.log(torch.clamp(x, min=clip_val) * C)
|
|
|
|
|
|
def dynamic_range_decompression_torch(x, C=1):
|
|
return torch.exp(x) / C
|
|
|
|
|
|
def spectral_normalize_torch(magnitudes):
|
|
output = dynamic_range_compression_torch(magnitudes)
|
|
return output
|
|
|
|
|
|
def spectral_de_normalize_torch(magnitudes):
|
|
output = dynamic_range_decompression_torch(magnitudes)
|
|
return output
|
|
|
|
|
|
mel_basis = {}
|
|
hann_window = {}
|
|
|
|
|
|
def mel_spectrogram(y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False):
|
|
'''
|
|
if torch.min(y) < -1.:
|
|
print('min value is ', torch.min(y))
|
|
if torch.max(y) > 1.:
|
|
print('max value is ', torch.max(y))
|
|
'''
|
|
global mel_basis, hann_window
|
|
if fmax not in mel_basis:
|
|
#mel = librosa_mel_fn(sampling_rate, n_fft, num_mels, fmin, fmax)
|
|
# sr, n_fft, n_mels=128, fmin=0.0, fmax
|
|
mel = librosa.filters.mel(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax)
|
|
mel_basis[str(fmax)+'_'+str(y.device)] = torch.from_numpy(mel).float().to(y.device)
|
|
hann_window[str(y.device)] = torch.hann_window(win_size).to(y.device)
|
|
|
|
y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_size)/2), int((n_fft-hop_size)/2)), mode='reflect')
|
|
y = y.squeeze(1)
|
|
|
|
spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[str(y.device)],
|
|
center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=False)
|
|
|
|
spec = torch.sqrt(spec.pow(2).sum(-1)+(1e-9))
|
|
|
|
spec = torch.matmul(mel_basis[str(fmax)+'_'+str(y.device)], spec)
|
|
spec = spectral_normalize_torch(spec)
|
|
|
|
return spec
|
|
|
|
|
|
def get_dataset_filelist_org(a):
|
|
with open(a.input_training_file, 'r', encoding='utf-8') as fi:
|
|
training_files = [os.path.join(a.input_wavs_dir, x.split('|')[0] + '.wav')
|
|
for x in fi.read().split('\n') if len(x) > 0]
|
|
|
|
with open(a.input_validation_file, 'r', encoding='utf-8') as fi:
|
|
validation_files = [os.path.join(a.input_wavs_dir, x.split('|')[0] + '.wav')
|
|
for x in fi.read().split('\n') if len(x) > 0]
|
|
return training_files, validation_files
|
|
|
|
def get_dataset_filelist(a):
|
|
with open(a.input_training_file, 'r', encoding='utf-8') as fi:
|
|
training_files = [x for x in fi.read().split('\n') if len(x) > 0]
|
|
|
|
with open(a.input_validation_file, 'r', encoding='utf-8') as fi:
|
|
validation_files = [x for x in fi.read().split('\n') if len(x) > 0]
|
|
|
|
return training_files, validation_files
|
|
|
|
class MelDataset(torch.utils.data.Dataset):
|
|
def __init__(self, training_files, segment_size, n_fft, num_mels,
|
|
hop_size, win_size, sampling_rate, fmin, fmax, split=True, shuffle=True, n_cache_reuse=1,
|
|
device=None, fmax_loss=None, fine_tuning=False, base_mels_path=None):
|
|
self.audio_files = training_files
|
|
random.seed(1234)
|
|
if shuffle:
|
|
random.shuffle(self.audio_files)
|
|
self.segment_size = segment_size
|
|
self.sampling_rate = sampling_rate
|
|
self.split = split
|
|
self.n_fft = n_fft
|
|
self.num_mels = num_mels
|
|
self.hop_size = hop_size
|
|
self.win_size = win_size
|
|
self.fmin = fmin
|
|
self.fmax = fmax
|
|
self.fmax_loss = fmax_loss
|
|
self.cached_wav = None
|
|
self.n_cache_reuse = n_cache_reuse
|
|
self._cache_ref_count = 0
|
|
self.device = device
|
|
self.fine_tuning = fine_tuning
|
|
self.base_mels_path = base_mels_path
|
|
self.supported_samples = [16000, 22050, 24000] #[4000, 8000, 16000, 22050, 24000, 32000]
|
|
#self.supported_samples = [4000, 8000] #, 16000, 22050, 24000, 32000]
|
|
|
|
def __getitem__(self, index):
|
|
filename = self.audio_files[index]
|
|
while 1:
|
|
#audio, sampling_rate = load_wav(filename)
|
|
audio, sampling_rate = load_segment(filename, self.sampling_rate, self.segment_size)
|
|
if audio is not None: break
|
|
else:
|
|
filename = self.audio_files[random.randint(0,index)]
|
|
#audio, sampling_rate = load_wav(filename)
|
|
#audio, sampling_rate = load_segment(filename, self.sampling_rate, self.segment_size)
|
|
|
|
#audio = audio / MAX_WAV_VALUE
|
|
if not self.fine_tuning:
|
|
audio = normalize(audio) * 0.95
|
|
|
|
sr_out = random.choice(self.supported_samples)
|
|
audio_down = resample(audio, self.sampling_rate, sr_out)
|
|
|
|
target_len = len(audio) #/ downsample_factor
|
|
audio_up = resample(audio_down, None, None, target_len)
|
|
|
|
audio = torch.FloatTensor(audio)
|
|
audio = audio.unsqueeze(0)
|
|
audio_up = torch.FloatTensor(audio_up)
|
|
audio_up = audio_up.unsqueeze(0)
|
|
|
|
mel = mel_spectrogram(audio_up, self.n_fft, self.num_mels,
|
|
self.sampling_rate, self.hop_size, self.win_size, self.fmin, self.fmax,
|
|
center=False)
|
|
|
|
mel_loss = mel_spectrogram(audio, self.n_fft, self.num_mels,
|
|
self.sampling_rate, self.hop_size, self.win_size, self.fmin, self.fmax_loss,
|
|
center=False)
|
|
|
|
return (mel.squeeze(), audio.squeeze(0), filename, mel_loss.squeeze())
|
|
|
|
def __getitem__org(self, index):
|
|
filename = self.audio_files[index]
|
|
if self._cache_ref_count == 0:
|
|
while 1:
|
|
audio, sampling_rate = load_wav(filename)
|
|
if audio is not None: break
|
|
else:
|
|
filename = self.audio_files[random.randint(0,index)]
|
|
audio, sampling_rate = load_wav(filename)
|
|
|
|
audio = audio / MAX_WAV_VALUE
|
|
if not self.fine_tuning:
|
|
audio = normalize(audio) * 0.95
|
|
#self.cached_wav = audio
|
|
if sampling_rate != self.sampling_rate:
|
|
resample_factor = self.sampling_rate / sampling_rate
|
|
new_samples = int(len(audio) * resample_factor)
|
|
audio = scipy.signal.resample(audio, new_samples)#.astype(np.int16)
|
|
#raise ValueError("{} SR doesn't match target {} SR".format(
|
|
# sampling_rate, self.sampling_rate))
|
|
|
|
downsample_factor = 16000 / self.sampling_rate
|
|
new_samples = int(len(audio) * downsample_factor)
|
|
audio_down = scipy.signal.resample(audio, new_samples)
|
|
|
|
new_samples = len(audio) #/ downsample_factor
|
|
audio_up = scipy.signal.resample(audio_down, new_samples)
|
|
#print(f'audio: {audio.shape}, audio_up: {audio_up.shape}')
|
|
#min_idx = min(len(audio), len(audio_up))
|
|
#audio = audio[:min_idx]
|
|
#audio_up = audio_up[:min_idx]
|
|
|
|
self.cached_wav = audio
|
|
self.cached_wav_up = audio_up
|
|
self._cache_ref_count = self.n_cache_reuse
|
|
else:
|
|
audio = self.cached_wav
|
|
audio_up = self.cached_wav_up
|
|
self._cache_ref_count -= 1
|
|
|
|
audio = torch.FloatTensor(audio)
|
|
audio = audio.unsqueeze(0)
|
|
audio_up = torch.FloatTensor(audio_up)
|
|
audio_up = audio_up.unsqueeze(0)
|
|
|
|
if True:
|
|
if self.split:
|
|
if audio.size(1) >= self.segment_size:
|
|
max_audio_start = audio.size(1) - self.segment_size
|
|
audio_start = random.randint(0, max_audio_start)
|
|
audio = audio[:, audio_start:audio_start+self.segment_size]
|
|
audio_up = audio_up[:, audio_start:audio_start+self.segment_size]
|
|
else:
|
|
audio = torch.nn.functional.pad(audio, (0, self.segment_size - audio.size(1)), 'constant')
|
|
audio_up = torch.nn.functional.pad(audio_up, (0, self.segment_size - audio_up.size(1)), 'constant')
|
|
|
|
mel = mel_spectrogram(audio_up, self.n_fft, self.num_mels,
|
|
self.sampling_rate, self.hop_size, self.win_size, self.fmin, self.fmax,
|
|
center=False)
|
|
|
|
mel_loss = mel_spectrogram(audio, self.n_fft, self.num_mels,
|
|
self.sampling_rate, self.hop_size, self.win_size, self.fmin, self.fmax_loss,
|
|
center=False)
|
|
|
|
return (mel.squeeze(), audio.squeeze(0), filename, mel_loss.squeeze())
|
|
|
|
def __len__(self):
|
|
return len(self.audio_files)
|