98 lines
3.9 KiB
Python
98 lines
3.9 KiB
Python
|
|
||
|
|
|
||
|
|
"""
|
||
|
|
Master script that processes all MP3 files found in an input directory. Splits those files up into sub-files of a
|
||
|
|
predetermined duration.
|
||
|
|
"""
|
||
|
|
import argparse
|
||
|
|
import functools
|
||
|
|
import os
|
||
|
|
from multiprocessing.pool import ThreadPool
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
import torch
|
||
|
|
import torchaudio
|
||
|
|
import numpy as np
|
||
|
|
from tqdm import tqdm
|
||
|
|
|
||
|
|
from data.util import find_audio_files, find_files_of_type
|
||
|
|
from trainer.injectors.audio_injectors import TorchMelSpectrogramInjector
|
||
|
|
from utils.util import load_audio
|
||
|
|
|
||
|
|
|
||
|
|
def report_progress(progress_file, file):
|
||
|
|
with open(progress_file, 'a', encoding='utf-8') as f:
|
||
|
|
f.write(f'{file}\n')
|
||
|
|
|
||
|
|
|
||
|
|
spec_fn = TorchMelSpectrogramInjector({'n_mel_channels': 256, 'mel_fmax': 11000, 'filter_length': 16000,
|
||
|
|
'true_normalization': True, 'normalize': True, 'in': 'in', 'out': 'out'}, {}).cuda()
|
||
|
|
|
||
|
|
def produce_mel(audio):
|
||
|
|
return spec_fn({'in': audio.unsqueeze(0)})['out'].squeeze(0)
|
||
|
|
|
||
|
|
|
||
|
|
def process_folder(folder, base_path, output_path, progress_file, max_duration, sampling_rate=22050):
|
||
|
|
outdir = os.path.join(output_path, f'{os.path.relpath(folder, base_path)}')
|
||
|
|
os.makedirs(outdir, exist_ok=True)
|
||
|
|
|
||
|
|
files = list(os.listdir(folder))
|
||
|
|
i = 0
|
||
|
|
output_i = 0
|
||
|
|
while i < len(files):
|
||
|
|
last_ordinal = -1
|
||
|
|
total_progress = 0
|
||
|
|
to_combine = []
|
||
|
|
while i < len(files) and total_progress < max_duration:
|
||
|
|
audio_file = os.path.join(folder, files[i], "no_vocals.wav")
|
||
|
|
if not os.path.exists(audio_file):
|
||
|
|
break
|
||
|
|
to_combine.append(load_audio(audio_file, 22050))
|
||
|
|
file_ordinal = int(files[i])
|
||
|
|
if last_ordinal != -1 and file_ordinal != last_ordinal+1:
|
||
|
|
break
|
||
|
|
else:
|
||
|
|
i += 1
|
||
|
|
total_progress += 30
|
||
|
|
if total_progress > 30:
|
||
|
|
combined = torch.cat(to_combine, dim=-1).cuda()
|
||
|
|
mel = produce_mel(combined).cpu().numpy()
|
||
|
|
np.savez(os.path.join(outdir, f'{output_i}'), mel)
|
||
|
|
output_i += 1
|
||
|
|
report_progress(progress_file, folder)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument('--path', type=str, help='Path to search for files', default='Y:\\separated')
|
||
|
|
parser.add_argument('--progress_file', type=str, help='Place to store all files that have already been processed', default='Y:\\separated\\large_mels\\already_processed.txt')
|
||
|
|
parser.add_argument('--output_path', type=str, help='Path for output files', default='Y:\\separated\\large_mels')
|
||
|
|
parser.add_argument('--num_threads', type=int, help='Number of concurrent workers processing files.', default=3)
|
||
|
|
parser.add_argument('--max_duration', type=int, help='Duration per clip in seconds', default=120)
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
os.makedirs(args.output_path, exist_ok=True)
|
||
|
|
processed_files = set()
|
||
|
|
if os.path.exists(args.progress_file):
|
||
|
|
with open(args.progress_file, 'r', encoding='utf-8') as f:
|
||
|
|
for line in f.readlines():
|
||
|
|
processed_files.add(line.strip())
|
||
|
|
|
||
|
|
cache_path = os.path.join(args.output_path, 'cache.pth')
|
||
|
|
if os.path.exists(cache_path):
|
||
|
|
root_music_files = torch.load(cache_path)
|
||
|
|
else:
|
||
|
|
path = Path(args.path)
|
||
|
|
def collect(p):
|
||
|
|
return str(os.path.dirname(os.path.dirname(p)))
|
||
|
|
root_music_files = {collect(p) for p in path.rglob("*no_vocals.wav")}
|
||
|
|
torch.save(root_music_files, cache_path)
|
||
|
|
|
||
|
|
orig_len = len(root_music_files)
|
||
|
|
folders = root_music_files - processed_files
|
||
|
|
print(f"Found {len(folders)} files to process. Total processing is {100 * (orig_len - len(folders)) / orig_len}% complete.")
|
||
|
|
|
||
|
|
with ThreadPool(args.num_threads) as pool:
|
||
|
|
list(tqdm(pool.imap(functools.partial(process_folder, output_path=args.output_path, base_path=args.path,
|
||
|
|
progress_file=args.progress_file, max_duration=args.max_duration), folders), total=len(folders)))
|