Spaces:
Running
Running
| ### | |
| # Author: Kai Li | |
| # Date: 2024-01-22 01:16:22 | |
| # Email: lk21@mails.tsinghua.edu.cn | |
| # LastEditTime: 2024-01-24 00:05:10 | |
| ### | |
| import torch | |
| import torchaudio | |
| import json | |
| from typing import Any, Dict, List, Optional, Tuple | |
| import os | |
| from omegaconf import OmegaConf | |
| import argparse | |
| import pytorch_lightning as pl | |
| import hydra | |
| from pytorch_lightning.strategies.ddp import DDPStrategy | |
| from pytorch_lightning import Callback, LightningDataModule, LightningModule, Trainer | |
| # from pytorch_lightning.loggers import Logger | |
| from omegaconf import DictConfig | |
| import look2hear.system | |
| import look2hear.datas | |
| import look2hear.losses | |
| import look2hear.models | |
| from look2hear.metrics import MetricsTracker | |
| from look2hear.utils import RankedLogger, instantiate, print_only | |
| import warnings | |
| warnings.filterwarnings("ignore") | |
| class Owndata(): | |
| def __init__(self, root): | |
| self.root = root | |
| self.data_lists = [] | |
| for root, dirs, files in os.walk(self.root): | |
| for file in files: | |
| if file == "codec_wav.wav": | |
| self.data_lists.append(os.path.join(root, file)) | |
| def __len__(self): | |
| return len(self.data_lists) | |
| def __getitem__(self, idx): | |
| ori = self.data_lists[idx].replace("codec_wav.wav", "ori_wav.wav") | |
| ori_audio = torchaudio.load(ori)[0] | |
| codec_audio = torchaudio.load(self.data_lists[idx])[0] | |
| return ori_audio, codec_audio, ori | |
| def test(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: | |
| data_val = Owndata("./codec-test") | |
| cfg.model.pop("_target_", None) | |
| model = look2hear.models.GullFullband.from_pretrain(os.path.join(cfg.exp.dir, cfg.exp.name, "best_model.pth"), **cfg.model).cuda() | |
| model.cuda() | |
| os.makedirs(os.path.join(cfg.exp.dir, cfg.exp.name, "results/"), exist_ok=True) | |
| metrics = MetricsTracker(save_file=os.path.join(cfg.exp.dir, cfg.exp.name, "results/")+"metrics.csv") | |
| length = len(data_val) | |
| for idx in range(len(data_val)): | |
| ori_wav, codec_wav, key = data_val[idx] | |
| ori_wav = ori_wav.cuda() | |
| codec_wav = codec_wav.unsqueeze(0).cuda() | |
| with torch.no_grad(): | |
| ests = model(codec_wav) | |
| torchaudio.save(key.replace("ori_wav.wav", "ests_wav.wav"), ests.squeeze(0).cpu(), 44100) | |
| metrics(ori_wav, ests, key) | |
| if idx % 10 == 0: | |
| dicts = metrics.update() | |
| print(f"Processed {idx}/{length} samples: SDR: {dicts['sdr']}, SI-SNR: {dicts['si-snr']}, VISQOL: {dicts['visqol']}") | |
| metrics.final() | |
| print("Finished testing") | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument( | |
| "--conf_dir", | |
| default="Exps/Apollo/config.yaml", | |
| help="Full path to save best validation model", | |
| ) | |
| args = parser.parse_args() | |
| cfg = OmegaConf.load(args.conf_dir) | |
| os.makedirs(os.path.join(cfg.exp.dir, cfg.exp.name), exist_ok=True) | |
| OmegaConf.save(cfg, os.path.join(cfg.exp.dir, cfg.exp.name, "config.yaml")) | |
| test(cfg) | |