110 lines
4.0 KiB
Python
110 lines
4.0 KiB
Python
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
@File : layers.py
|
|
@Time : 2024/4/22 下午2:40
|
|
@Author : waytan
|
|
@Contact : waytan@tencent.com
|
|
@License : (C)Copyright 2024, Tencent
|
|
"""
|
|
import os
|
|
import json
|
|
import time
|
|
import logging
|
|
import argparse
|
|
from datetime import datetime
|
|
|
|
|
|
import torch
|
|
|
|
from models.apply import BagOfModels
|
|
from models.pretrained import get_model_from_yaml
|
|
|
|
|
|
class Separator:
|
|
def __init__(self, dm_model_path, dm_config_path, gpu_id=0) -> None:
|
|
if torch.cuda.is_available() and gpu_id < torch.cuda.device_count():
|
|
self.device = torch.device(f"cuda:{gpu_id}")
|
|
else:
|
|
self.device = torch.device("cpu")
|
|
self.demucs_model = self.init_demucs_model(dm_model_path, dm_config_path)
|
|
|
|
def init_demucs_model(self, model_path, config_path) -> BagOfModels:
|
|
model = get_model_from_yaml(config_path, model_path)
|
|
model.to(self.device)
|
|
model.eval()
|
|
return model
|
|
|
|
def run(self, audio_path, output_dir, ext=".flac"):
|
|
name, _ = os.path.splitext(os.path.split(audio_path)[-1])
|
|
output_paths = []
|
|
for stem in self.demucs_model.sources:
|
|
output_path = os.path.join(output_dir, f"{name}_{stem}{ext}")
|
|
if os.path.exists(output_path):
|
|
output_paths.append(output_path)
|
|
if len(output_paths) == 4:
|
|
drums_path, bass_path, other_path, vocal_path = output_paths
|
|
else:
|
|
drums_path, bass_path, other_path, vocal_path = self.demucs_model.separate(audio_path, output_dir, device=self.device)
|
|
data_dict = {
|
|
"vocal_path": vocal_path,
|
|
"bgm_path": [drums_path, bass_path, other_path]
|
|
}
|
|
return data_dict
|
|
|
|
|
|
def json_io(input_json, output_json, model_dir, dst_dir, gpu_id=0):
|
|
current_datetime = datetime.now()
|
|
current_datetime_str = current_datetime.strftime('%Y-%m-%d-%H:%M')
|
|
logging.basicConfig(filename=os.path.join(dst_dir, f'logger-separate-{os.path.split(input_json)[1]}-{current_datetime_str}.log'), level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
|
|
|
|
sp = Separator(os.path.join(model_dir, "htdemucs.pth"), os.path.join(model_dir, "htdemucs.yaml"), gpu_id=gpu_id)
|
|
with open(input_json, "r") as fp:
|
|
lines = fp.readlines()
|
|
t1 = time.time()
|
|
success_num = 0
|
|
fail_num = 0
|
|
total_num = len(lines)
|
|
sep_items = []
|
|
for line in lines:
|
|
item = json.loads(line)
|
|
flac_file = item["path"]
|
|
try:
|
|
fix_data = sp.run(flac_file, dst_dir)
|
|
except Exception as e:
|
|
fail_num += 1
|
|
logging.error(f"process-{success_num + fail_num}/{total_num}|success-{success_num}|fail-{fail_num}|{item['idx']} process fail for {str(e)}")
|
|
continue
|
|
|
|
item["vocal_path"] = fix_data["vocal_path"]
|
|
item["bgm_path"] = fix_data["bgm_path"]
|
|
sep_items.append(item)
|
|
success_num += 1
|
|
logging.debug(f"process-{success_num + fail_num}/{total_num}|success-{success_num}|fail-{fail_num}|{item['idx']} process success")
|
|
|
|
with open(output_json, "w", encoding='utf-8') as fw:
|
|
for item in sep_items:
|
|
fw.write(json.dumps(item, ensure_ascii=False) + "\n")
|
|
|
|
t2 = time.time()
|
|
logging.debug(f"total cost {round(t2-t1, 3)}s")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description='')
|
|
parser.add_argument("-m", dest="model_dir")
|
|
parser.add_argument("-d", dest="dst_dir")
|
|
parser.add_argument("-j", dest="input_json")
|
|
parser.add_argument("-o", dest="output_json")
|
|
parser.add_argument("-gid", dest="gpu_id", default=0, type=int)
|
|
args = parser.parse_args()
|
|
|
|
if not args.dst_dir:
|
|
dst_dir = os.path.join(os.getcwd(), "separate_result")
|
|
os.makedirs(dst_dir, exist_ok=True)
|
|
else:
|
|
dst_dir = os.path.join(args.dst_dir, "separate_result")
|
|
os.makedirs(dst_dir, exist_ok=True)
|
|
|
|
json_io(args.input_json, args.output_json, args.model_dir, dst_dir, gpu_id=args.gpu_id)
|