|
import os |
|
import torch |
|
import torchaudio |
|
import torchaudio.transforms as T |
|
from tqdm import tqdm |
|
import numpy as np |
|
from omegaconf import DictConfig |
|
import hydra |
|
from hydra.utils import to_absolute_path |
|
from transformers import Wav2Vec2FeatureExtractor, AutoModel |
|
|
|
from encoder.mert import FeatureExtractorMERT |
|
from encoder.music2latent import FeatureExtractorM2L |
|
|
|
class AudioProcessor: |
|
def __init__(self, cfg: DictConfig): |
|
self.input_directory = cfg.dataset.input_dir |
|
self.output_directory = cfg.dataset.output_dir |
|
self.segment_duration = cfg.segment_duration |
|
self.resample_rate = cfg.model.sr |
|
self.device_id = cfg.device_id |
|
self.feature_extractor = self._initialize_extractor(cfg.model.name) |
|
self.is_split = cfg.is_split |
|
|
|
def _initialize_extractor(self, model_name: str): |
|
if "MERT" in model_name: |
|
return FeatureExtractorMERT(model_name=model_name, device_id=self.device_id, sr=self.resample_rate) |
|
elif "music2latent" == model_name: |
|
return FeatureExtractorM2L(device_id=self.device_id, sr=self.resample_rate) |
|
else: |
|
raise NotImplementedError(f"Feature extraction for model {model_name} is not implemented.") |
|
|
|
def resample_waveform(self, waveform, original_sample_rate, target_sample_rate): |
|
if original_sample_rate != target_sample_rate: |
|
resampler = T.Resample(original_sample_rate, target_sample_rate) |
|
return resampler(waveform), target_sample_rate |
|
return waveform, original_sample_rate |
|
|
|
def split_audio(self, waveform, sample_rate): |
|
segment_samples = self.segment_duration * sample_rate |
|
total_samples = waveform.size(0) |
|
|
|
segments = [] |
|
for start in range(0, total_samples, segment_samples): |
|
end = start + segment_samples |
|
if end <= total_samples: |
|
segment = waveform[start:end] |
|
segments.append(segment) |
|
|
|
|
|
if len(segments) == 0: |
|
segment = waveform |
|
segments.append(segment) |
|
|
|
return segments |
|
|
|
def process_audio_file(self, file_path, output_dir): |
|
print(f"Processing {file_path}") |
|
waveform, sample_rate = torchaudio.load(file_path) |
|
|
|
if waveform.shape[0] > 1: |
|
waveform = waveform.mean(dim=0).unsqueeze(0) |
|
waveform = waveform.squeeze() |
|
waveform, sample_rate = self.resample_waveform(waveform, sample_rate, self.resample_rate) |
|
|
|
if self.is_split: |
|
segments = self.split_audio(waveform, sample_rate) |
|
for i, segment in enumerate(segments): |
|
segment_save_path = os.path.join(output_dir, f"segment_{i}.npy") |
|
if os.path.exists(segment_save_path): |
|
continue |
|
self.feature_extractor.extract_features_from_segment(segment, sample_rate, segment_save_path) |
|
else: |
|
segment_save_path = os.path.join(output_dir, f"segment_0.npy") |
|
if not os.path.exists(segment_save_path): |
|
self.feature_extractor.extract_features_from_segment(waveform, sample_rate, segment_save_path) |
|
|
|
def process_directory(self): |
|
for root, _, files in os.walk(self.input_directory): |
|
for file in files: |
|
if file.endswith('.mp3'): |
|
file_path = os.path.join(root, file) |
|
relative_path = os.path.relpath(file_path, self.input_directory) |
|
output_file_dir = os.path.join(self.output_directory, os.path.splitext(relative_path)[0]) |
|
os.makedirs(output_file_dir, exist_ok=True) |
|
self.process_audio_file(file_path, output_file_dir) |
|
|
|
@hydra.main(version_base=None, config_path="../config", config_name="prep_config") |
|
def main(cfg: DictConfig): |
|
processor = AudioProcessor(cfg) |
|
processor.process_directory() |
|
|
|
if __name__ == "__main__": |
|
main() |
