Add partitioning and improved resuming to the spleeter filtering

This commit is contained in:
James Betker 2021-10-06 17:10:12 -06:00
parent f2977d360c
commit bb891a3a53

View File

@ -78,20 +78,34 @@ def main():
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('--path') parser.add_argument('--path')
parser.add_argument('--out') parser.add_argument('--out')
parser.add_argument('--resume', default=None)
parser.add_argument('--partition_size', default=None)
parser.add_argument('--partition', default=None)
args = parser.parse_args() args = parser.parse_args()
src_dir = args.path src_dir = args.path
out_file = args.out out_file = args.out
output_sample_rate=22050 output_sample_rate=22050
batch_size=16 waiting_for_file = args.resume is not None
resume_file = args.resume
audio_loader = AudioAdapter.default() audio_loader = AudioAdapter.default()
files = find_audio_files(src_dir, include_nonwav=True) files = find_audio_files(src_dir, include_nonwav=True)
# Partition files if needed.
if args.partition_size is not None:
psz = int(args.partition_size)
prt = int(args.partition)
files = files[prt*psz:(prt+1)*psz]
#separator = Separator('pretrained_models/2stems', input_sr=output_sample_rate) #separator = Separator('pretrained_models/2stems', input_sr=output_sample_rate)
separator = Separator('spleeter:2stems') separator = Separator('spleeter:2stems')
unacceptable_files = open(out_file, 'w') unacceptable_files = open(out_file, 'a')
for path in tqdm(files): for e, path in enumerate(tqdm(files)):
print(f"Processing {src_dir}") if waiting_for_file and resume_file not in path:
continue
waiting_for_file = False
print(f"{e}: Processing {path}")
spleeter_ld, sr = audio_loader.load(path, sample_rate=output_sample_rate) spleeter_ld, sr = audio_loader.load(path, sample_rate=output_sample_rate)
sep = separator.separate(spleeter_ld) sep = separator.separate(spleeter_ld)
vocals = sep['vocals'] vocals = sep['vocals']