| |
| |
| """ |
| @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) |
|
|
