mel_processing.py 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153
  1. import math
  2. import os
  3. import random
  4. import torch
  5. from torch import nn
  6. import torch.nn.functional as F
  7. import torch.utils.data
  8. import numpy as np
  9. import librosa
  10. import librosa.util as librosa_util
  11. from librosa.util import normalize, pad_center, tiny
  12. from scipy.signal import get_window
  13. from scipy.io.wavfile import read
  14. from librosa.filters import mel as librosa_mel_fn
  15. MAX_WAV_VALUE = 32768.0
  16. def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
  17. """
  18. PARAMS
  19. ------
  20. C: compression factor
  21. """
  22. return torch.log(torch.clamp(x, min=clip_val) * C)
  23. def dynamic_range_decompression_torch(x, C=1):
  24. """
  25. PARAMS
  26. ------
  27. C: compression factor used to compress
  28. """
  29. return torch.exp(x) / C
  30. def spectral_normalize_torch(magnitudes):
  31. output = dynamic_range_compression_torch(magnitudes)
  32. return output
  33. def spectral_de_normalize_torch(magnitudes):
  34. output = dynamic_range_decompression_torch(magnitudes)
  35. return output
  36. mel_basis = {}
  37. hann_window = {}
  38. def spectrogram_torch(y, n_fft, sampling_rate, hop_size, win_size, center=False):
  39. if torch.min(y) < -1.0:
  40. print("min value is ", torch.min(y))
  41. if torch.max(y) > 1.0:
  42. print("max value is ", torch.max(y))
  43. global hann_window
  44. dtype_device = str(y.dtype) + "_" + str(y.device)
  45. wnsize_dtype_device = str(win_size) + "_" + dtype_device
  46. if wnsize_dtype_device not in hann_window:
  47. hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(
  48. dtype=y.dtype, device=y.device
  49. )
  50. y = torch.nn.functional.pad(
  51. y.unsqueeze(1),
  52. (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)),
  53. mode="reflect",
  54. )
  55. y = y.squeeze(1)
  56. spec = torch.stft(
  57. y,
  58. n_fft,
  59. hop_length=hop_size,
  60. win_length=win_size,
  61. window=hann_window[wnsize_dtype_device],
  62. center=center,
  63. pad_mode="reflect",
  64. normalized=False,
  65. onesided=True,
  66. return_complex=False,
  67. )
  68. spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
  69. return spec
  70. def spec_to_mel_torch(spec, n_fft, num_mels, sampling_rate, fmin, fmax):
  71. global mel_basis
  72. dtype_device = str(spec.dtype) + "_" + str(spec.device)
  73. fmax_dtype_device = str(fmax) + "_" + dtype_device
  74. if fmax_dtype_device not in mel_basis:
  75. mel = librosa_mel_fn(
  76. sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax
  77. )
  78. mel_basis[fmax_dtype_device] = torch.from_numpy(mel).to(
  79. dtype=spec.dtype, device=spec.device
  80. )
  81. spec = torch.matmul(mel_basis[fmax_dtype_device], spec)
  82. spec = spectral_normalize_torch(spec)
  83. return spec
  84. def mel_spectrogram_torch(
  85. y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False
  86. ):
  87. if torch.min(y) < -1.0:
  88. print("min value is ", torch.min(y))
  89. if torch.max(y) > 1.0:
  90. print("max value is ", torch.max(y))
  91. global mel_basis, hann_window
  92. dtype_device = str(y.dtype) + "_" + str(y.device)
  93. fmax_dtype_device = str(fmax) + "_" + dtype_device
  94. wnsize_dtype_device = str(win_size) + "_" + dtype_device
  95. if fmax_dtype_device not in mel_basis:
  96. mel = librosa_mel_fn(
  97. sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax
  98. )
  99. mel_basis[fmax_dtype_device] = torch.from_numpy(mel).to(
  100. dtype=y.dtype, device=y.device
  101. )
  102. if wnsize_dtype_device not in hann_window:
  103. hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(
  104. dtype=y.dtype, device=y.device
  105. )
  106. y = torch.nn.functional.pad(
  107. y.unsqueeze(1),
  108. (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)),
  109. mode="reflect",
  110. )
  111. y = y.squeeze(1)
  112. spec = torch.stft(
  113. y,
  114. n_fft,
  115. hop_length=hop_size,
  116. win_length=win_size,
  117. window=hann_window[wnsize_dtype_device],
  118. center=center,
  119. pad_mode="reflect",
  120. normalized=False,
  121. onesided=True,
  122. return_complex=False,
  123. )
  124. spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
  125. spec = torch.matmul(mel_basis[fmax_dtype_device], spec)
  126. spec = spectral_normalize_torch(spec)
  127. return spec