135 lines
5.1 KiB
Python
135 lines
5.1 KiB
Python
import librosa
|
|
import numpy as np
|
|
import torch
|
|
import os
|
|
from scipy.io.wavfile import write, read
|
|
from model import BigramLanguageModel
|
|
from divideStereo import split_stereo, add_stereo
|
|
|
|
class ProceesAudio():
|
|
audio_data = []
|
|
|
|
final_audio = []
|
|
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
|
|
target_audio = []
|
|
|
|
def _split_audio_s(self, file_path, segment_length=0.5, overlap=0):
|
|
print(file_path)
|
|
# Load the stereo audio file
|
|
audio, sr = librosa.load(file_path , sr=None, mono=False, dtype=np.float32)
|
|
|
|
# Calculate segment and overlap samples
|
|
segment_samples = int(segment_length * sr)
|
|
overlap_samples = int(segment_samples * overlap)
|
|
|
|
# Split the stereo audio into segments while preserving stereo channels
|
|
segments = []
|
|
for start in range(0, audio.shape[1], segment_samples - overlap_samples):
|
|
segment = audio[:, start:start + segment_samples]
|
|
if segment.shape[1] == segment_samples:
|
|
segments.append(segment)
|
|
|
|
return segments, sr
|
|
|
|
def _split_audio(self, file_path, segment_length=0.1, overlap=0):
|
|
audio, sr = librosa.load(file_path, sr=None)
|
|
segment_samples = int(segment_length * sr)
|
|
overlap_samples = int(segment_samples * overlap)
|
|
segments = []
|
|
for start in range(0, len(audio), segment_samples - overlap_samples):
|
|
segment = audio[start:start + segment_samples]
|
|
if len(segment) == segment_samples:
|
|
segments.append(segment)
|
|
return segments, sr
|
|
|
|
def _calculate_average_amplitude(self, segments, sr, n_fft=2048, hop_length=256, num_frequency_bands=100):
|
|
# for segment in segments:
|
|
# stft = librosa.stft(segment, n_fft=n_fft, hop_length=hop_length)
|
|
# magnitude = np.abs(stft)
|
|
# max_amplitude = max(max_amplitude, np.max(magnitude))
|
|
ret=[]
|
|
audios = []
|
|
max_amp = 0
|
|
for segment in segments:
|
|
max_amp = max(max_amp, np.max(np.abs(segment)))
|
|
for segment in segments:
|
|
amps, _ = split_stereo(segment=segment, max_amp=max_amp, sr=sr, num_parts=10)
|
|
ret.append(amps)
|
|
return ret, audios
|
|
|
|
def _process_main(self, file_path):
|
|
segments, sr = self._split_audio_s(file_path)
|
|
amps, _ = self._calculate_average_amplitude(segments, sr)
|
|
for amp in amps:
|
|
self.audio_data.append(amp)
|
|
|
|
def _get_output_amps(self, input_amps, index):
|
|
model = BigramLanguageModel()
|
|
model.to("cpu", dtype=float)
|
|
model.load_state_dict(torch.load('amp_net.pth', map_location=torch.device('cpu')))
|
|
return model.generate(torch.tensor(input_amps[:index+1]).view(1, index+1, len(input_amps[0])).to(self.device), len(input_amps)-index + 1)
|
|
|
|
|
|
def make_smooth(self, audio, gain, prev_gain):
|
|
smooth_index = 1000
|
|
mult_arr_in = np.linspace(prev_gain, gain, num=smooth_index)
|
|
for i in range(smooth_index):
|
|
audio[:smooth_index][i] *= mult_arr_in[i]
|
|
audio[smooth_index:] *= gain
|
|
return audio
|
|
|
|
|
|
def perform_modulation(self, file, index, output_file_name, num_frequency_bands=100):
|
|
segments, sr = self._split_audio_s(file)
|
|
max_amp = 0
|
|
for segment in segments:
|
|
max_amp = max(max_amp, np.max(np.abs(segment)))
|
|
|
|
x, _ = self._calculate_average_amplitude(segments=segments, sr=sr)
|
|
#print(x.shape)
|
|
y = self._get_output_amps(x, index)
|
|
modified_segs = []
|
|
prev_gains = np.ones(10)
|
|
for segment, mod in zip(segments, y[0]):
|
|
_, audios = split_stereo(segment=segment, max_amp=max_amp, sr=sr, num_parts=10)
|
|
final_audios = []
|
|
curr_gains = []
|
|
|
|
for audio, target_amp, i in zip(audios, mod, range(10)):
|
|
|
|
gain = (target_amp.item()/np.mean(np.abs(audio)))*max_amp
|
|
|
|
if np.mean(np.abs(audio)) == 0:
|
|
gain=0
|
|
elif gain <= 50:
|
|
gain = gain/50
|
|
else:
|
|
gain=1
|
|
|
|
audio = self.make_smooth(audio, gain, prev_gains[i])
|
|
curr_gains.append(gain)
|
|
final_audios.append(audio)
|
|
|
|
prev_gains = curr_gains
|
|
modified_seg = add_stereo(final_audios, len(final_audios[0]), sample_rate=sr)
|
|
modified_segs.append(modified_seg)
|
|
|
|
modified_segs = np.concatenate(modified_segs)
|
|
write("sitare_modified.wav", rate=sr, data=modified_segs.astype(np.float32))
|
|
|
|
def get_training_data(self, file_path, data_dir):
|
|
for song in os.listdir(data_dir):
|
|
self.audio_data = []
|
|
self._process_main(os.path.join(data_dir, song))
|
|
audio_data = np.array(self.audio_data)
|
|
tensor_data = torch.tensor(audio_data, dtype=torch.float32)
|
|
torch.save(tensor_data, file_path + '/' + song + '.data.pt')
|
|
|
|
|
|
|
|
pa = ProceesAudio()
|
|
#pa.get_training_data('extracted_data', 'songs')
|
|
pa.perform_modulation('sitare.wav', 0, 'output.wav')
|